Unmanaged gradient accumulation: ZeRO offload support (#8225)
## Summary
* Extends unmanaged gradient accumulation
(`managed_gradient_accumulation=false`) to **ZeRO optimizer-state and
parameter offload** (CPU/NVMe). Follow-up to #8217 (ZeRO stage 3, now
merged).
* Stage 2/3: grads still reduce/partition every `backward()`; `step()`
finalizes deferred offload boundary work (grad norms + FP32/NVMe copy)
via `finalize_gradient_accumulation_boundary()`.
* Stage 1: continues to reduce at `step()` via `allreduce_gradients()`,
which already performs offload boundary finalization when the boundary
flag is true.
* Pipeline parallelism, DeepCompile, Apex AMP, and stage-0/1
`overlap_comm` remain unsupported.
## Test plan
Validated on a 2-GPU node:
* [x] Full `-k Unmanaged` suite (**31 passed**), including:
* `test_unmanaged_matches_managed_optimizer_offload[1|2|3]`
* `test_unmanaged_matches_managed_param_offload` (stage 3)
* existing non-offload unmanaged equivalence / varying-GAS / rejection
tests
* [x] Docs updated (`config-json.md`, `training.rst`); 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: Ma, Guokai <guokai.ma@gmail.com>
