fix(pp): keep accumulated grads across stage re-initialization (#3530)
`_initialize_pp_stages` ends by calling `_post_metadata_inference_cleanup` on every stage, which sets `param.grad = None` on FSDP modules. Upstream that is safe: initialization happens once, before any real backward, and the only gradients present are the throwaway ones left by dynamic metadata inference. `reset_pp_stage_shapes` breaks that assumption. It re-initializes the stages on every sequence-length change, which for VLM batches is most steps, and under gradient accumulation that lands between micro-batches. Everything accumulated so far is discarded, so only the final micro-batch reaches the optimizer. Restore the gradients when the stage is in static-metadata mode, where no inference forward/backward ran and there is nothing stale to clear. Dynamic mode is left alone. Measured on a Gemma4 proxy at pp_size=2, gradient norm versus a single-rank reference on identical data: accumulation=2 before 1.94, 1.99, 1.20 after 1.07, 0.96, 0.98 accumulation=3 before 1.67, 1.34, 1.03 after 1.02, 1.00, 0.99 Signed-off-by: Abhishree <abhishreetm@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
A
Abhishree Thittenamane committed
c4c0e4b8f069de16f095ad35498de7c4502e7af8
Parent: 875637f
Committed by GitHub <noreply@github.com>
on 8/20/2026, 5:37:16 AM