DeepSpeed

f9fe4df3 - ZeRO-2 CPU offload: accumulate in the compact partition between requests (#9044)

ZeRO-2 CPU offload: accumulate in the compact partition between requests (#9044) ## 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. - The device partition is freed after the boundary copy only if the host gradient restores the running sum bit for bit (same dtype, or an fp32 host holding bf16/fp16 values). The next backward in the same optimizer window restores it. Otherwise (for example bf16 masters with fp32 accumulation) it stays resident until `step()`. - The gate for the compact buffer is unchanged from before this series: fp32 accumulation on a 16-bit model, in the stages and offload modes that already used it (including stage 1 with offload). Muon groups keep the host accumulator. The default bf16/fp16 accumulation keeps the per-parameter path, so no partition-sized device buffer is added for it. - Overflow is now checked on the value cast to the host gradient dtype. With fp16 master weights and fp32 accumulation, an fp32 sum above the fp16 range used to pass the check and become inf on the host; it now skips the step. - Legacy host path: gradients that a non-boundary backward left only in the host accumulator are published at the next boundary (`_publish_unpublished_offload_grads`). - `config.py` is unchanged: fp32 accumulation with fp16/bf16 master weights is not rejected, as before. - New option `zero_optimization.retain_gpu_grad_partition_until_step` (default `false`). The default frees the fp32 gradient partition from GPU memory between requests. Set it to `true` to keep the partition in GPU memory until `step()`; GPU memory use between requests then rises by one partition per rank. With `true`, 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::TestZeroOffloadReleaseOption` checks 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 - CPU: `test_zero2_gradient_safety.py`, `test_zero_cpu_offload_grad_accum.py`, `test_zero2_offload_multi_backward.py`: 88 passed, 12 skipped. New `test_zero_offload_accum_precision.py`: 10 passed. - B200 (4 GPUs), final commit `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. - B200 (4 GPUs), new `test_multiple_request_boundaries_host_accumulator` and `test_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. - B200 device-memory probe (single GPU, ZeRO-2 CPU offload, synthetic 134M-parameter bf16 stack of eight 4096x4096 linears, 3 optimizer windows, 4 backwards per window each marked as a request boundary, one `step()` per window). Values are MB of `torch.cuda.memory_allocated` after each request backward, and the per-window `max_memory_allocated`: | masters / accumulation | base `a4c989b7`: between requests | this PR `9960ecb7`: between requests | window peak, base / this PR | | --- | --- | --- | --- | | fp32 / fp32 | 816 | 304 | 2090 / 2090 | | bf16 / fp32 (unsupported) | 816 | 816 | 2090 / 2090 | | fp32 / bf16 (default) | 304 | 304 | 1386 / 1386 | The 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) and `test_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, and `safe_set_full_grad` residency 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()`: | stage | masters | accumulation | error | notes | | --- | --- | --- | --- | --- | | 2 | fp32 | fp32 | 0 | the supported exact case | | 2 | fp32 | bf16 | 6.2e-3 | per-step bf16 accumulation | | 2 | bf16 | fp32 | 2.4e-3 | unsupported: one rounding to the bf16 host gradient | | 2 | bf16 | bf16 | 6.2e-3 | | | 1 | fp32 | fp32 | 2.4e-3 | the window sum is copied into the reduction bucket of the communication dtype (bf16 here) and rounded once before the all-reduce; identical on base | | 1 | fp32 | bf16 | 6.2e-3 | | | 1 | bf16 | fp32 | 2.4e-3 | | | 1 | bf16 | bf16 | 6.2e-3 | | 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_grad` between backwards in one window: - A value set while the GPU partition is resident (after the first backward of a window) is overwritten at the boundary; it is not kept. - A value set after a boundary has freed the partition (fp32 masters) is kept, and the next backward accumulates onto it exactly. ## Stage-1 memory, step time, multi-rank, and Qwen3-4B Same request-boundary probe as above (every backward is a boundary, then `step()`), comparing base `a4c989b7` to this PR. Times are median after the first window. Memory is `torch.cuda.memory_allocated` after 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: true` removes 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.** | stage / GPUs / accumulation | base between (MB) | this PR between (MB) | window peak base / this PR (MB) | median step ms base / this PR | | --- | --- | --- | --- | --- | | 2 / 1 / bf16 (default) | 8431 | 8431 | 13251 / 13251 | 3923 / 3782 | | 2 / 1 / fp32 | 23775 | 8431 | 29187 / 29187 | 3966 / 3984 | | 1 / 1 / bf16 | 8431 | 8431 | 19142 / 19142 | 3899 / 3793 | | 1 / 1 / fp32 | 23775 | 8431 | 42182 / 42017 | 3889 / 3929 | | 2 / 2 / bf16 | 8431 | 8431 | 13251 / 13251 | 16767 / 16375 | | 2 / 2 / fp32 | 16103 | 8431 | 20923 / 20024 | 16176 / 15781 | 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 - No full production training job (loss curve, tokens/s over many steps) was run. - ZeRO-1 offload with bf16/fp16 masters: fp32 accumulation does accumulate in fp32 on the device, but the window sum is rounded once when copied to the bf16 host gradient (table above; identical on base). ## Known follow-ups (not in this PR) - PCIe traffic optimization (deferred deliberately). - Stage 1 rounds the window sum to 16 bits at reduction (pre-existing). --------- Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>