Fix merged LoRA cache invalidation across FSDP2 updates

[kiddyboots216]

Summary

• invalidate canonical merged-LoRA folds when the published adapter generation changes instead of relying on tensor version and pointer identity
• advance generations for optimizer updates, logical/checkpoint restores, and rollback; route rank-0 broadcast restores through the generation-aware API
• invalidate direct-optimizer trainer paths after successful mutation

Why

FSDP2/DTensor scratch publication can change parameter bytes without advancing Tensor._version, while allocator reuse can preserve the same pointer. The old cache key could therefore reuse stale folded MoE weights after many successful steps.

Validation

• 143 focused tests passed
• two-rank H100 real-FSDP2 lifecycle test passed, including an observed reused version/pointer key and byte-identical sync-fold versus next-forward outputs after the update
• Ruff, formatting, diff, and public-hygiene checks passed

[broly-code-security-scanner]

Broly Security Scan

Note

✅ Clean scan
No vulnerabilities detected in this PR.

Note

Re-scan this PR anytime with /broly scan — useful after /broly undismiss, or to refresh findings without a new push.

Broly — SAST (zai-org/GLM-5.2) · Secrets · SCA · IaC · GH Actions · Base Images · Supply Chain Threats · Exploit Chains · Adversarial Verification

We're continuously improving Broly's accuracy and finding quality — your feedback is valuable. False positives, missed findings, bugs, and feature requests all welcome.

 

[kiddyboots216]

Superseded by #58. The identical merged-LoRA cache fix is now commit b5ef860 on that branch.