Repository navigation
ZeRO-2 CPU offload: accumulate in the compact partition between requests - #9044
Conversation
With stage-2 CPU offload, the running gradient sum now lives in one partition-sized buffer in grad_accum_dtype during each request, and is published to the host FP32 partition at the request boundary. The GPU partition is freed between requests when the host copy can restore the sum exactly (same dtype, or FP32 host holding bf16/fp16 values), and is restored by the next backward in the same optimizer window. The compact path is limited to stage-2 CPU offload. Stage 1 (cumulative unreduced grad_accum across reentrant phases), Muon, and ZenFlow keep their existing gating. Config now rejects data_types.grad_accum_dtype="fp32" combined with fp16/bf16 master weights and grads, since the optimizer step runs in the master dtype and would discard the extra accumulation precision. Tests: update the stubs in the gradient-safety tests to the legacy path, read the running partition sum for grad-accum assertions, skip the fp32-accum overflow case under low-precision masters, and add config tests for the new rejection rule. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 00bcc411a5
ℹ️ 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".
| if (self.grad_accum_dtype is not None and DtypeEnum(self.grad_accum_dtype).value == torch.float32 | ||
| and (self.float16_config.fp16_master_weights_and_grads | ||
| or self.bfloat16_config.bf16_master_weights_and_grads)): | ||
| raise ValueError("data_types.grad_accum_dtype='fp32' is incompatible with " |
There was a problem hiding this comment.
This unconditional configuration check rejects every ZeRO stage using grad_accum_dtype="fp32" with FP16/BF16 master weights, including Stage 1, Stage 3, and non-offload paths unrelated to the new compact Stage-2 buffer. The extra precision is not simply discarded: summing microbatch gradients in FP32 before one final cast can avoid overflow and rounding loss that occurs when each addition is performed in the master dtype. Previously valid production configurations now fail during initialization, and the changed overflow test merely skips this combination; scope the rejection to a genuinely unsupported implementation path or preserve the conversion at publication.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Agreed that the check was too broad, and it is already gone on this branch (be4ef875).
The compact Stage-2 path does not need the rejection: when the host copy cannot restore the running sum exactly (bf16/fp16 masters with fp32 accumulation), the GPU partition stays resident until step(). The original check ran at config time for every ZeRO stage and offload mode, including paths that never use that buffer.
One correction to the production-config claim: Stage 1 / Stage 3 without CPU-Adam offload already fail on master (stage_1_and_2.py and stage3.py require offload + DeepSpeedCPUAdam for master-weight features). The check mainly blocked Stage 2 without offload, plus CPU-offload configs that still train.
On 4× B200 with the rejection removed: test_finite_step_after_overflow passed 8/8, including the fp32-accumulation cases with low-precision masters. A scratch smoke test also trained with finite losses for Stage 2 with and without CPU offload, and for Stage 1 with CPU offload. Stage 1 without offload still hits the existing master-weight assertion.
Not claimed: that fp32 accumulation is more accurate than bf16 accumulation on this path. That comparison is not in the PR.
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Remove the config-time rejection of data_types.grad_accum_dtype='fp32' combined with fp16/bf16 master weights and grads. The rejection applied to every ZeRO stage and offload mode, although the compact-buffer path it was meant to protect already keeps the device sum resident when the bf16/fp16 host copy cannot restore it exactly. Validated on 4x B200: the overflow-recovery tests pass for all fp32-accumulation cases with low-precision masters, and a bf16-master fp32-accumulation smoke test trains with finite losses for stage 2 with and without CPU offload and for stages 1/2 with CPU offload. Stage 1 without offload is still rejected by the existing master-weight assertion in stage_1_and_2.py. Tests: drop the rejection test and the overflow-test skip that covered the removed rule. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
…s norm test Check the compact partition's overflow after the cast to the host gradient dtype, so an fp32 sum that exceeds the fp16 range is caught when fp16 master weights are used. Cover the host-accumulator path across request boundaries, the fp16 host overflow, and the freed-buffer state in the accumulation norm test. Drop a config test that never exercised a rejection. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
…lation on 16-bit models Restore the gate that applied before this series. Default bf16/fp16 accumulation keeps the per-parameter path, so it does not hold a partition-sized buffer on the device. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Characterize the window sum for stage x master x accumulation with bf16-exact micro-step gradients: stage 2 with fp32 masters and fp32 accumulation is exact; other cells round once (bf16 accumulation, or the host copy for bf16 masters). Also pin safe_set_full_grad: a value set after a freed boundary is kept and accumulated; one set while the partition is resident is overwritten at the boundary. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
|
Raw verification data for the numbers in the PR description, with base and PR side by side. Source: B200 job logs job-20261109T232543Z and job-20261109T235123Z. 'base' is a4c989b and 'PR' is ebf0112. Between-request memory is torch.cuda.memory_allocated after one boundary backward, on rank 0; window peak is the last window's max_memory_allocated. Backward and step times are medians. Synthetic probe (eight 4096x4096 linears, bf16), base / PR side by side, last window:
Qwen3-4B (random-init, seq 512, 4 boundary backwards/window), base / PR side by side, last window:
The released fp32 path is faster to hold memory but slower per backward on Qwen3-4B (for example 474→733 ms, stage 2, 1 GPU). See the PR description. |
With ZeRO-2 CPU offload and fp32 accumulation on a 16-bit model, the compact GPU partition is freed after each request boundary by default. Setting release_compact_grad_partition=false keeps it resident until step(), trading one fp32 partition of device memory for no per-request host copy and restore. Test: residency per boundary and final weights match in both settings. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
…een_requests The new name states what the option controls: whether the fp32 gradient partition is freed from GPU HBM between requests. No behavior change. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
…ition_cache No behavior change. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Invert the flag: retain_hbm_grad_partition=true keeps the fp32 gradient partition in HBM until step(); the default false frees it between requests, as before. Behavior at the default is unchanged. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
|
why is there an 2.4e-3 error in row 5 of the pr description ( 1 | fp32 | fp32) and not 0, like z2? |
Per review: name the memory as GPU rather than HBM (LPDDR-attached GPUs exist) and state the span: the partition is kept until step(), not for each step. Behavior is unchanged. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This is a pre-existing divergence of z1 and z2 gas implementations z1: a single bf16 reduction of fp32 accumulated grads |
so should this be fixed? I suppose not many use z1, so it was never a priority? |
Yes, low priority for that reason. Also, selecting fp32 reduction type avoids the issue. |
The user is unlikely to know that though. I'd say a short term solution could be having fp32 reduction as a default for Z1 (albeit not sure about BC). |
Since this requires a bit of coding, so I will follow up with a different PR for the issue. |
tohtana
added
the
needs-gpu-test
tohtana
removed
the
needs-gpu-test
Summary
Follow-up to #9018 (merged as
a4c989b7). With CPU offload, the compact partition-sized accumulator (used for fp32 accumulation on a 16-bit model) now publishes to the host partition gradient at each request boundary and is freed between requests when the host copy can restore it exactly.step()._publish_unpublished_offload_grads).config.pyis unchanged: fp32 accumulation with fp16/bf16 master weights is not rejected, as before.zero_optimization.retain_gpu_grad_partition_until_step(defaultfalse). The default frees the fp32 gradient partition from GPU memory between requests. Set it totrueto keep the partition in GPU memory untilstep(); GPU memory use between requests then rises by one partition per rank. Withtrue, memory returns to the base level and the backward returns to base timing (Qwen3-4B stage 2, 1 GPU: 455 ms vs base 474 ms; synthetic: 14.4 ms vs 14.1 ms). Step time is unchanged. Test:test_zero_offload_accum_precision.py::TestZeroOffloadReleaseOptionchecks residency per boundary and that the final weights match in both settings.Unsupported: bf16 master weights with fp32 gradient accumulation. The host gradient is bf16 and cannot hold the fp32 sum, so the GPU partition is not released between requests in that case. It is also rounded once to bf16 when copied to the host (see below).
Test plan
test_zero2_gradient_safety.py,test_zero_cpu_offload_grad_accum.py,test_zero2_offload_multi_backward.py: 88 passed, 12 skipped. Newtest_zero_offload_accum_precision.py: 10 passed.9960ecb7: targeted offload suites (gradient safety, grad-accum norm, multi-backward, coalesced grad reduction, config dict) gave 227 passed, 28 warnings; the autoep gradient-parity file gave 18 passed, 2 skipped.test_multiple_request_boundaries_host_accumulatorandtest_fp16_host_gradient_overflow_skips_step: pass. With the pre-fix overflow check restored in a scratch copy, the fp32 case of the overflow test fails.step()per window). Values are MB oftorch.cuda.memory_allocatedafter each request backward, and the per-windowmax_memory_allocated:a4c989b7: between requests9960ecb7: between requestsThe 512 MB saving between requests is the fp32 partition of this model. The in-window peak is unchanged, because the partition is still built within a window. Stage 1 and other model sizes were not measured.
Test changes
test_zero2_gradient_safety.py: stubs pin the legacy host path with a comment; they bypass the reduction that fills the compact buffer.test_zero_cpu_offload_grad_accum.py: the norm assertion reads the device partition sum, or the host gradient once the partition has been freed, and runs with and without fp32 accumulation. It checks every norm is above zero.test_zero2_offload_multi_backward.py: adds a host-accumulator request-boundary test (exact gradient checks; fails if the publish call is disabled) andtest_fp16_host_gradient_overflow_skips_step.test_zero_offload_accum_precision.py(new): window-sum precision for stage x master x accumulation using micro-step gradients that are exact in bf16, andsafe_set_full_gradresidency behavior. See below.Measured precision and residency (not claims about production)
Window sum over 16 micro-steps, relative max error against the exact sum, read before
step():The base commit gives the same table, so this PR does not change precision for these cells.
Why stage 1 fp32/fp32 is 2.4e-3 and stage 2 fp32/fp32 is 0: stage 2 reduces each micro-step gradient into the fp32 partition. These micro-step values are exact in bf16, so every reduction is exact and the fp32 sum is exact. Stage 1 keeps the window sum in fp32 (
grad_accum), but at the boundary it copies that sum into the reduction bucket, which has the communication dtype (bf16 in this test, since the model is bf16). The sum is therefore rounded once to bf16 before the all-reduce. Stage 1 is not changed by this PR. A fp32 reduction for stage 1 would remove the rounding at twice the communication volume; it is not in this PR.safe_set_full_gradbetween backwards in one window:Stage-1 memory, step time, multi-rank, and Qwen3-4B
Same request-boundary probe as above (every backward is a boundary, then
step()), comparing basea4c989b7to this PR. Times are median after the first window. Memory istorch.cuda.memory_allocatedafter a boundary backward, on rank 0.Synthetic 134M-parameter stack (eight 4096x4096 linears), 1 GPU. Between-request memory for fp32 masters + fp32 accumulation drops from 816 MB to 304 MB on both stage 1 and stage 2. Window peak is unchanged (stage 1: 3492 MB; stage 2: 2090 MB). Default bf16 accumulation stays at 304 MB. Unsupported bf16 masters + fp32 accumulation stays at 816 MB. Median
step()is 26 ms on both trees. Median backward for the released fp32-accumulation path is about 10 ms slower on this stack (15→25 ms on stage 1, 14→24 ms on stage 2), which is the host copy plus restore on the next request. On Qwen3-4B the same path is much slower per backward: 474→733 ms (stage 2, 1 GPU), 496→785 ms (stage 1, 1 GPU), 483→687 ms (stage 2, 2 GPUs). This is the cost of releasing the partition.zero_optimization.retain_gpu_grad_partition_until_step: trueremoves it, at the price of keeping one fp32 partition resident (see the Summary).Same stack, 2 GPUs, fp32 masters + fp32 accumulation. Each rank's partition is half the model (256 MB of fp32). Stage 2 between-request memory drops from 560 MB to 304 MB; stage 1 from 2468 MB to 2212 MB. Window peaks are unchanged. Median
step()is 127–135 ms on both trees.Qwen3-4B (random-init from the published config, 4.02B parameters, seq 512, 4 request-boundary backwards per window), fp32 masters.
On 1 GPU the released fp32 partition is about 15.3 GB (4.02e9 × 4 bytes). On 2 GPUs it is about 7.7 GB per rank. Default bf16 accumulation is unchanged.
step()is within a few percent of base. This is a request-boundary microbenchmark, not a full SFT job.Remaining limits
Known follow-ups (not in this PR)