Unmanaged gradient accumulation: ZeRO stage 2 support (#8203)
## Summary
Extends unmanaged gradient accumulation
(`managed_gradient_accumulation=false`) to **ZeRO stage 2**. Stacked on
top of the stage 0/1 foundation in #8184 (base branch
`sfc-gh-truwase/gas_mgmt_zero01`); review that PR first.
Unlike stage 0/1 (where `backward()` accumulates locally and `step()`
performs the reduction), ZeRO stage 2 must reduce/partition gradients on
**every** `backward()` to preserve its memory characteristics. This is
compatible with unmanaged mode because reduce-scatter is linear:
accumulating the reduced partitions across N caller-controlled backwards
is equivalent to reducing once at the boundary. Micro-step tracking
stays disabled and the caller still owns the boundary; only the
`averaged_gradients` finalization is deferred to `step()`.
- `ZeROOptimizer.finalize_gradient_accumulation_boundary()` (stage 1/2)
builds `averaged_gradients` from the accumulated `all_grad_tensors` at
`step()`.
- Validation relaxed to allow stage 2 (`not partition_weights`), still
rejecting stage 3 and ZeRO offload (follow-up PRs).
- **`overlap_comm` is now supported for stage 2**: its async reduction
is confined to the per-backward path and the epilogue synchronizes
before finalizing, so it behaves exactly as in managed mode. It remains
rejected for stage 0/1 (where reduction is deferred to `step()`).
## Test plan
Validated on a 2-GPU node (full `-k Unmanaged` suite, 22 passed):
- [x] `test_unmanaged_matches_managed[2]` — unmanaged stage-2 matches
managed stage-2
- [x] `test_unmanaged_varying_backward_count[2]` — variable backward
count per step, stage 2
- [x] `test_unmanaged_matches_managed_overlap_comm` — unmanaged stage-2
with `overlap_comm=True` matches managed reference
- [x] `test_unmanaged_rejects_stage3`,
`test_unmanaged_rejects_zero_offload` — stage 3 / offload rejected
- [x] `test_unmanaged_rejects_overlap_comm[0,1]` — overlap_comm still
rejected for stage 0/1
- [x] `pre-commit` (yapf/flake8/codespell) clean
Docs (`config-json.md`, `training.rst`) updated to describe stage-2
behavior and the overlap_comm support; previewable on `rtd-staging`.
Made with [Cursor](https://cursor.com)
---------
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>
