DeepSpeed

a4c989b7 - ZeRO-1/2: one CPU-offload accumulator per group; zero idle gradients (#9018)

ZeRO-1/2: one CPU-offload accumulator per group; zero idle gradients (#9018) Replacement for #8964, which was closed after an accidental force-push of an unrelated commit onto this branch. GitHub will not reopen #8964 after that rewrite. Follow-up to #8949 (merged as dc046b15). Fixes #8991. This replaces the earlier content of this PR (accumulating every backward in the optimizer's FP32 partition gradient). That approach was dropped: its extra complexity was not worth its gains. This version keeps #8949's design and makes two narrow changes: preallocate the CPU-offload accumulator per partition group instead of lazily per parameter, and fix the stale-gradient leak. ## Problem - **Lazy accumulator complexity.** #8949 stores CPU-offload gradient sums in `accumulated_grads_in_cpu`, a dict keyed by parameter id and filled on first use. FP32 masters get owned buffers sized to the whole parameter; low-precision masters get aliased views of the partition gradient, created in two places. Accumulate, copy-back, and restore each branch on which kind they hold (about ten branches across four functions). The pinned buffers are never unpinned, and the low-precision restore leaves stale data in a shared scratch tensor that the overflow scan reads. - **Stale-gradient leak (#8991).** `reset_cpu_buffers` did not zero the FP32 partition gradient. A trainable parameter that receives no gradient in the next window kept its previous slice, and Adam applied it again. ## Fix - One host accumulator per partition group, sized to the partition, in the bit16 dtype, pinned when `cpu_offload_pin_memory` is set. Each parameter uses a fixed slice at its `dest_offset`, the same layout as the FP32 partition. The buffer is allocated on the first non-boundary accumulation, so boundary-only training (GAS=1) pins no extra host memory. Low-precision masters keep accumulating in the partition gradient. - `_get_cpu_offload_accumulator_slice()` is the only place that picks the buffer. The per-parameter dict, both alias sites, and the full-parameter versus fragment branching are removed. - Restore clears the scratch tensor before copying this rank's fragment. - The group buffers are unpinned with the other offload buffers. - `reset_cpu_buffers` zeros each FP32 partition gradient after the step, after overflow, and on reset. - `set_norm_for_param_grad` is removed; it has no callers. No config or API changes. ## Overflow detection Restore now zeros the non-owned part of the restored gradient, where #8949 kept this rank's local pre-reduction values. Skip-step behavior does not change: each reduced fragment is checked by the rank that owns it, and `has_overflow()` MAX-reduces `local_overflow` across the data-parallel and model-parallel groups. A multi-rank test with an inf in a non-owned region was not run. ## Not a duplicate - #8949 (merged): this builds on it. - #8991 (open issue): the leak this fixes. - #8728 (merged as 0d0aa30b): clears stale offload scratch before restoring low-precision gradients. Its change is the same `zero_()` in `_restore_cpu_offload_grad_to_gpu` that this PR includes. This PR is rebased onto master (b62afa82) after #8728 landed. The one conflict was in `_restore_cpu_offload_grad_to_gpu`; the resolution keeps this PR's version (unconditional `zero_()` on the scratch, then the fragment copy from `_get_cpu_offload_accumulator_slice`). - #8920 (open): FP32 folded-gradient accumulation. Unrelated. - #8957 (open issue): remaining test coverage. Not required for this change. ## Validation ### CPU unit tests Hardware: 1x Intel Xeon Platinum 8488C (96 logical CPUs), no GPU. Python 3.12.12, PyTorch 2.9.1+cu130. ```bash python -m pytest -q tests/unit/v1/zero/test_zero2_gradient_safety.py \ tests/unit/v1/zero/test_zero_cpu_offload_grad_accum.py ``` 57 passed, 2 skipped (the two CUDA stream tests). After rebasing onto master b62afa82, `test_zero2_gradient_safety.py`, `test_zero_cpu_offload_grad_accum.py`, and `test_zero2_offload_multi_backward.py` (includes #8728's test) run with `LOCAL_SIZE=2`: 87 passed, 2 skipped. #8728's regression test, `TestZero2OffloadOverflowRecovery::test_finite_step_after_overflow` (two CPU ranks via `LOCAL_SIZE=2`, BF16 and FP16, low-precision and FP32 masters), passes 4/4 on this branch. With the `zero_()` in `_restore_cpu_offload_grad_to_gpu` removed, all 4 cases fail with `finite step did not update large`. The new leak tests (`test_idle_parameter_window_does_not_reuse_previous_gradient`, `test_idle_parameter_stays_zero_beside_active_neighbor`, FP32 and BF16 masters) compare against `torch.optim.Adam`; all four fail with the `reset_cpu_buffers` zeroing removed. `test_single_boundary_offload_avoids_auxiliary_storage_and_copy` checks that a boundary-only step allocates no host accumulator. ```bash python -m pytest -q tests/unit/runtime/zenflow \ tests/unit/runtime/zero/test_zero1_optimizer_fastpath.py \ tests/unit/runtime/half_precision/test_zero_grad_norm_overflow.py ``` 97 passed, 42 skipped. `pre-commit run --files` passed on the three changed files. ### GPU memory and throughput Hardware: 8x NVIDIA H200 (141 GiB), PyTorch 2.14.0+cu130, CUDA 13.0 (driver version not recorded). Qwen3-4B and Qwen3.5-4B, ZeRO-2, bf16, CPU optimizer offload with pinned memory, `torch.optim.Adam`, GAS 1 and 4, sequence length 512 and 2048, 1 warmup and the median of 2 measured iterations, at 1, 2, 4, and 8 GPUs. Baseline is #8949 plus the same `reset_cpu_buffers` zeroing, so these numbers isolate the accumulator change. The measured tree differs from this commit only in comments and tests. All 16 runs passed (64 rows). Change vs. baseline (#8949 plus the same partition-gradient zeroing), in percent. Negative means faster. `bwd` is total backward time per optimizer step; `iter` is forward + backward + step. Across all rows, iteration time moves between -4.0% and +2.0%. Backward moves within ±2.5% in every row except Qwen3.5-4B at 8 GPUs with GAS=4, which is 9.0% faster at sequence length 512 and 8.8% faster at 2048. Both sequence lengths show it, but this sweep does not explain the cause. Sequence length 512: | Model | GPUs | GAS | bwd | iter | |---|---|---|---|---| | Qwen3-4B | 1 | 1 | +4.4% | +0.2% | | Qwen3-4B | 1 | 4 | +2.1% | +1.2% | | Qwen3-4B | 2 | 1 | -0.4% | +2.0% | | Qwen3-4B | 2 | 4 | -0.1% | +1.4% | | Qwen3-4B | 4 | 1 | -0.7% | -0.3% | | Qwen3-4B | 4 | 4 | -0.7% | -3.5% | | Qwen3-4B | 8 | 1 | -1.2% | -1.6% | | Qwen3-4B | 8 | 4 | -1.3% | -0.4% | | Qwen3.5-4B | 1 | 1 | +0.6% | +0.7% | | Qwen3.5-4B | 1 | 4 | -0.3% | -1.0% | | Qwen3.5-4B | 2 | 1 | +2.5% | +0.4% | | Qwen3.5-4B | 2 | 4 | +1.4% | +0.9% | | Qwen3.5-4B | 4 | 1 | +1.0% | +0.1% | | Qwen3.5-4B | 4 | 4 | -0.4% | -1.8% | | Qwen3.5-4B | 8 | 1 | +1.3% | -0.0% | | Qwen3.5-4B | 8 | 4 | -9.0% | +0.0% | Sequence length 2048: | Model | GPUs | GAS | bwd | iter | |---|---|---|---|---| | Qwen3-4B | 1 | 1 | +4.6% | +1.8% | | Qwen3-4B | 1 | 4 | +2.4% | +1.3% | | Qwen3-4B | 2 | 1 | +0.7% | +1.9% | | Qwen3-4B | 2 | 4 | -0.4% | +0.1% | | Qwen3-4B | 4 | 1 | -1.1% | +0.4% | | Qwen3-4B | 4 | 4 | -0.7% | -2.1% | | Qwen3-4B | 8 | 1 | -0.5% | -3.0% | | Qwen3-4B | 8 | 4 | -0.8% | -4.0% | | Qwen3.5-4B | 1 | 1 | -0.5% | -0.1% | | Qwen3.5-4B | 1 | 4 | -0.6% | -1.3% | | Qwen3.5-4B | 2 | 1 | +2.1% | -1.7% | | Qwen3.5-4B | 2 | 4 | +0.5% | -0.9% | | Qwen3.5-4B | 4 | 1 | +1.3% | +1.9% | | Qwen3.5-4B | 4 | 4 | -0.3% | -0.9% | | Qwen3.5-4B | 8 | 1 | +0.6% | -0.7% | | Qwen3.5-4B | 8 | 4 | -8.8% | -0.6% | - Peak GPU allocated and reserved memory: identical in all 32 pairs. - Pinned host memory (peak GiB, baseline → this PR, sequence length 512): | Model | GPUs | GAS=1 | GAS=4 | |---|---|---|---| | Qwen3-4B | 1 | 41.0 → 41.0 | 51.6 → 49.0 | | Qwen3-4B | 2 | 21.0 → 21.0 | 26.3 → 25.0 | | Qwen3-4B | 4 | 11.0 → 11.0 | 13.6 → 13.0 | | Qwen3-4B | 8 | 6.0 → 6.0 | 7.3 → 7.0 | | Qwen3.5-4B | 1 | 42.0 → 42.0 | 53.9 → 50.0 | | Qwen3.5-4B | 2 | 22.0 → 22.0 | 28.1 → 26.0 | | Qwen3.5-4B | 4 | 12.0 → 12.0 | 15.2 → 14.0 | | Qwen3.5-4B | 8 | 7.0 → 7.0 | 9.0 → 8.0 | Host RSS moves with pinned memory. No convergence run was performed. ### Compatibility matrix Validated: ZeRO-1 and ZeRO-2 with CPU optimizer offload, FP32 and BF16 (low-precision) masters, bf16 parameters, managed gradient accumulation (GAS 1 and 4), 1 to 8 GPUs data parallel (GPU sweep is ZeRO-2 only), and the ZenFlow and ZeRO-1 fast-path unit tests above. Not validated: FP16 loss scaling, NVMe offload, ZeRO-3, pipeline, tensor, sequence, or expert parallelism, gradient clipping, checkpoint load in the middle of a window, and convergence. The GPU sweep never takes a non-boundary step at GAS=1, so the first-use allocation in unmanaged mode is covered only by CPU tests. The cost of zeroing the full FP32 partition each step was not measured on its own, because the baseline includes the same zeroing. ZeRO-3 is unchanged: its `reset_cpu_buffers` only clears norms, and it was not checked for the same leak. ZenFlow inherits `reset_cpu_buffers` and so gets the new zeroing; its unit tests pass. AI assistance was used to write, review, and measure this change. The submitting human needs to review every changed line before sign-off. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>