SIGN IN SIGN UP

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