[bug fix] ZeRO2: preserve fp32 grad accumulation precision in partition storage

[stas00]

When training bf16 models with data_types.grad_accum_dtype: fp32, ZeRO-2 could still store the reduced owned gradient partition in bf16 storage. That defeats the fp32 accumulation setting after communication: values that are representable in fp32 but lie between bf16 values are rounded before APIs such as safe_get_full_grad() read them back.

This matters for request accumulation and similar uneven-gradient cases. If two ranks contribute gradients 1.0 and 1.0078125, fp32 communication produces the exact average 1.00390625. Storing that result in a bf16 partition rounds it back to 1.0, so the training step loses the precision that grad_accum_dtype: fp32 requested.

This PR makes ZeRO-1/2 use fp32 owned gradient partition storage when fp32 gradient accumulation is active. The ZeRO-2 path now writes reduced fp32 communication results into a compact fp32 partition buffer and skips allocating the legacy bf16 grads_in_partition buffer for that path. This keeps the memory behavior aligned with the feature: the old bf16 owned partition is replaced rather than kept alongside the fp32 one, so bf16 + fp32 grad accumulation should use less GPU memory than the previous mixed-storage path.

ZeRO-3 was checked separately. It already allocates grad_partitions_flat_buffer with gradient_accumulation_dtype and casts reduced partitions to that dtype before storing them, so it does not have the same persistent bf16 partition-storage issue. A ZeRO-3 regression case is included to guard that behavior.

Memory saved

Run with a 1B-parameter bf16 toy model on 2 H200 ranks.

  Per rank:

   metric                       before a77aeb676    after 87e384ec3        change
  ━━━━━━━━━━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━
   peak allocated                     14,336 MiB         14,336 MiB     unchanged
  ───────────────────────────  ──────────────────  ─────────────────  ────────────
   peak reserved                      16,386 MiB         14,338 MiB    -2,048 MiB
  ───────────────────────────  ──────────────────  ─────────────────  ────────────
   reserved delta after init          12,290 MiB         10,242 MiB    -2,048 MiB 

Config was ZeRO-2, bf16 model, communication_data_type=fp32, data_types.grad_accum_dtype=fp32, GAS=4, numel=1,073,741,824, 2 ranks.

PyTorch peak allocated is unchanged, but peak reserved GPU memory drops by 2 GiB per rank.

Implementation

• Adds compact fp32 owned gradient partition buffers for ZeRO-1/2 when the model dtype differs from fp32 gradient accumulation dtype.
• Writes fp32 all-reduce results into that compact partition storage instead of relying on the bf16-owned partition copy.
• Avoids allocating the legacy bf16 grads_in_partition buffer on the compact fp32 path.
• Keeps the regression observable through safe_get_full_grad() rather than depending on private storage layout, with the ZeRO-2 internal dtype assertion kept only where it directly checks the fixed path.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 87e384ec3c

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[stas00]

@sfc-gh-truwase, I think the modal ci is the problem, it skips in the PR and then runs in the merge queue and fails there, trying to figure out if I can find the actual failures.

Details

[stas00]

So modal CI workload is broken, it completes only 24% of tests and then exits:

DS_CI_FAILURE_CLASS=timeout: Sandbox lifetime exhausted after 5400s (run pytest failed with exit code 137)
Error: Process completed with exit code 124.

[Copilot]

🟡 Changes recommended

CPU-offload, non-contiguous, and dynamic state-offload paths do not correctly preserve or manage the new fp32 partition storage.

3 open findings

• Preserve flat partition ownership across gradient offload and reload · New
• Preserve fp32 collective results with CPU offload and noncontiguous gradients · New
• Add ZeRO-1 coverage for fp32 reduction staging · New

What changed in this PR

Preserves fp32 gradient accumulation precision in ZeRO-1/2 partition storage.

Changes:

• Adds compact fp32 owned-gradient buffers.
• Routes reduced gradients into fp32 storage.
• Adds precision and compatibility tests.

File	Description
deepspeed/​runtime/​zero/​stage_1_and_2.py	Implements fp32 partition storage and reduction routing.
tests/​unit/​runtime/​zero/​test_zero_tensor_fragment.py	Adds gradient precision regression coverage.
tests/​unit/​v1/​zero/​test_zero2_gradient_safety.py	Tests legacy bucket compatibility.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.