[stas00]

Additionally to #8749 another precision issue

Porting this from @seanbell's Snowflake-AI-Research/Arctic-Platform#371

An accumulation boundary is not necessarily an optimizer step. A caller can finish several requests with engine.set_gradient_accumulation_boundary(True) while delaying engine.step() until all requests have contributed. ZeRO-1/2 must retain every request's gradients until that update and must not carry a previous update's gradients into the next one.

This fixes both non-offloaded and CPU-offloaded optimizer paths in deepspeed/runtime/zero/stage_1_and_2.py, without changing configuration defaults or adding an application-side workaround.

Root Cause And Fix

• Non-offloaded optimizers cleared all_grad_tensors at every boundary, discarding earlier requests before step(). Retain the accumulated tensors until the existing step cleanup. When the separate param.grad_accum tensor is cumulative, keep the existing reread behavior so reentrant checkpoint phases do not double-count gradients.
• CPU-offloaded optimizers with FP32 master weights skipped refreshing the persistent host accumulator on a step's first backward when that backward was also a boundary. A subsequent request could then add the previous step's stored gradients. Track each parameter's participation in the current update window instead of relying on the global micro_step_id. A first non-boundary contribution refreshes the host accumulator; a first boundary contribution uses only the existing optimizer-partition copy. If another backward follows before step(), recover the previous locally owned partition fragment in the accumulation dtype, then allocate or refresh the auxiliary host accumulator. Never reuse a prior update's sum.

A single-backward update needs no auxiliary allocation or copy. _offload_partition_only_param_ids records FP32-master parameters whose current sum lives only in the optimizer partition. Explicit finalization leaves those already-current gradients intact. step() clears both participation sets before either a successful update or an overflow return. Multiple backwards still need an auxiliary accumulator once the second contribution arrives; this change does not claim a speedup for that workload.

Low-precision masters use a shared buffer instead of the separate FP32-master accumulator. Their CPU accumulator is a view of the optimizer partition's gradient buffer. On a parameter's first use at a boundary, the existing boundary copy initializes that shared buffer directly, so a second host copy is unnecessary. Register the view and the parameter's participation even when skipping the extra copy; later backwards accumulate into the current window and explicit finalization can restore the shared gradient. A first non-boundary use still initializes the accumulator through the offload helper.

These changes address the stale-gradient first-use review, the low-precision boundary-handling clarification, and the subsequent single-backward allocation/copy concern. The code names first_use and boundary explicitly and distinguishes shared-buffer accumulation from the lazily allocated FP32-master accumulator.

This complements DeepSpeed #8749, rather than replacing its precision fix. #8749 preserves FP32 partition gradients, including CPU offload. The remaining defect concerns the lifetime of accumulated gradients when model and accumulation dtypes match. Arctic-Platform #371 motivated the native DeepSpeed repair; AP-specific hooks are not needed by this change.

Regression Coverage

The existing tests/unit/v1/zero/test_zero2_offload_multi_backward.py now includes a 16-case, two-rank integration matrix:

Axis	Cases
ZeRO stage	1, 2
Optimizer placement	Non-offloaded, CPU offload
BF16 model gradient accumulation	BF16, FP32
Activation checkpointing	Disabled, reentrant

Each case checks independently derived gradients and actual parameter updates. Two independent linear projections exercise partition ownership; analytical input-column sums supply the gradient reference and ordinary PyTorch SGD supplies the update reference. Binary-exact fixture values allow zero-tolerance assertions without loosening numerical thresholds.

Eight successive update windows cover a multi-call request, a zero-loss request followed by useful work, consecutive single-call requests, mixed request lengths, a boundary-first request followed by a multi-backward request, a fresh single backward after earlier updates, and two conditional-participation windows. Each projection is absent from the first backward of one window, then contributes repeatedly after having participated in the preceding update. The integration tests compare actual consecutive-single-backward updates; a separate storage regression checks the resource contract.

The existing gradient-safety suite also checks first use after another parameter's backward, with FP32 and BF16 master buffers and boundary/non-boundary first contributions. It checks the summed gradient and actual Adam update against an independent PyTorch optimizer, including a boundary-only overflow after an earlier update has allocated the auxiliary buffer. Additional controls check repeated finalization, partial partition ownership, BF16 accumulation rounding, pending pageable-source copies, and no additional storage or copy for consecutive single-boundary updates. These storage-level tests bypass the ZeRO constructor; they do not claim CPUAdam integration coverage for BF16 master weights.

Validation

All reported execution was on CPU, not on GPU. Hardware: one visible Intel Xeon Platinum 8488C socket, 96 logical CPUs; two CPU/Gloo ranks for distributed cases; no GPU or accelerator driver used. Software: PyTorch 2.9.1+cu130, DeepSpeed 0.19.8+0fe2dfcc7, Python from /home/yak/miniconda3/envs/dev/bin/python. The CUDA-enabled PyTorch build does not constitute CUDA hardware validation.

Test	Result
Plain request-boundary matrix before the original source fix	4 BF16-accumulation failures, 4 FP32-accumulation passes
Delayed-participation offload regressions before the review fix	2 BF16-accumulation failures, 2 FP32-accumulation passes
First-use storage regressions before the review fix	4 failures (FP32/BF16 masters, boundary/non-boundary first use); 4 ordinary-accumulation controls passed
Single-boundary allocation/copy regression before the lazy fix	1 FP32-master failure; 9 other storage controls passed
Complete request-boundary/checkpoint matrix after the lazy fix	16 passed
Existing multi-backward/offload controls after the lazy fix	10 passed
Gradient-safety suite after the lazy fix, including overflow	52 passed, 2 CUDA-only skips
Temporary CPU storage-lifetime and allocation/copy probes	6 passed

The repository suites total 78 passes and 2 CUDA-only skips; the temporary probes add 6 passes. The integration matrix completed in 88.91 seconds, native optimizer controls in 105.38 seconds, the final gradient-safety rerun in 8.63 seconds, and the final temporary probes in 2.51 seconds. These are test durations, not training-throughput measurements.

New matrix command, from the DeepSpeed checkout:

DS_ACCELERATOR=cpu PYTHONDONTWRITEBYTECODE=1 OMP_NUM_THREADS=1 \
  /home/yak/miniconda3/envs/dev/bin/python -m pytest -sv \
  tests/unit/v1/zero/test_zero2_offload_multi_backward.py \
  -k multiple_request_boundaries --tb=short --color=no

The existing optimizer controls initially timed out with optional shared-memory communication enabled. They passed using real Gloo with only that optional communication extension disabled. The temporary test launcher contains:

import sys
import deepspeed.ops
import pytest

deepspeed.ops.__compatible_ops__["deepspeed_shm

DS_ACCELERATOR=cpu PYTHONDONTWRITEBYTECODE=1 OMP_NUM_THREADS=1 \
  /home/yak/miniconda3/envs/dev/bin/python \
  /tmp/codex-01a13767-deepspeed-gloo-tests.py -sv \
  tests/unit/v1/zero/test_zero2_offload_multi_backward.py \
  -k TestZeroOffloadMultiBackward --tb=short --color=no

DS_ACCELERATOR=cpu PYTHONDONTWRITEBYTECODE=1 OMP_NUM_THREADS=1 \
  /home/yak/miniconda3/envs/dev/bin/python -m pytest -sv \
  tests/unit/v1/zero/test_zero2_gradient_safety.py --tb=short --color=no

The configured pre-commit hooks passed for all three changed Python files and this local description, including YAPF, flake8, license, and accelerator/distributed-import checks. The first run applied YAPF's required wrapping to the new tests; the second run passed without modifications. git diff --check passed.

Single-Backward Allocation And Copy Control

The lazy fix restores the original single-boundary resource behavior in a CPU-only fixture. For a fresh FP32-master optimizer with 1,048,576 BF16 gradient values, 12 single-boundary calls produced the following counts. The original base is 8d4f29c75; the eager review-fix head is 7e00e1a8e. The fixture resets per-update participation between calls and measures persistent auxiliary tensor storage and copy_ operations, not process peak memory or GPU transfers.

Revision	Auxiliary host bytes	Copies per boundary call
Original base	0	1
Eager review fix	2,097,152	2
Lazy fix	0	1

BF16-master controls remain at zero auxiliary bytes and one copy per call in all three revisions. A buffer allocated by an earlier multi-backward update remains available for reuse, but a later single-boundary update does not refresh that stale buffer. The repository regression test_single_boundary_offload_avoids_auxiliary_storage_and_copy checks the no-extra-storage/copy contract and actual consecutive Adam updates.

The larger allocation/copy comparison and four ordinary success/overflow storage-lifetime controls were run with the temporary probe /tmp/codex-01a13767-pr8949-additional-feedback.py:

DS_ACCELERATOR=cpu PYTHONDONTWRITEBYTECODE=1 OMP_NUM_THREADS=1 PYTHONPATH=.:tests \
  /home/yak/miniconda3/envs/dev/bin/python \
  /tmp/codex-01a13767-pr8949-additional-feedback.py

Synthetic CPU timings are not evidence of accelerator throughput. No end-to-end speedup is claimed.

Compatibility And Remaining Validation

The CPU runs validate optimizer-placement logic, not physical GPU optimizer placement or GPU-to-CPU transfers. The regression is also enabled on hosts with at least two CUDA devices, but GPU execution remains required before claiming on-device coverage. FP16, CPUAdam integration with low-precision master weights, NVMe offload, vendor accelerators, and tensor/pipeline/expert/sequence parallelism were not validated by the new matrix. BF16 master-buffer behavior was validated only in the gradient-safety fixture described above.

Static sibling inspection found that ZeRO-3 refreshes the first microstep's partition buffer and that the BF16 optimizer retains accumulation until its update or explicit zeroing. This PR changes neither implementation. Conditional participation after the first microstep was not validated for ZeRO-3; the inspected code is not evidence that every sibling accumulation case is correct.

Full-model convergence, throughput, and peak-memory measurements were not performed. Only the fixture's auxiliary-allocation and copy-count reductions above are established. GPU transfer cost and performance for repeated-backward workloads remain unmeasured.

Related Work And Submission

Public open-PR searches for gradient accumulation, stale gradients, and 8749 in:body, plus the #8749 discussion, were checked on October 29, 2026. No competing request-boundary fix was identified. Related #8728 clears unowned offload scratch for low-precision masters; #8856 repairs the BF16 full-gradient setter. Neither supplies this per-parameter request-lifetime fix. Repeat the duplicate check before opening another PR.

AI assistance was used. A human submitter must review every changed line, run the relevant accelerator tests, and provide any required sign-off. This is the local description for existing PR #8949, excluded from the source commit; GitHub PR text is not edited by this task.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-29T23:56:49.795422Z	bfa3c3d	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to 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" or "@codex security review".

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]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: bfa3c3dac2

ℹ️ 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".

[sfc-gh-truwase]

It is unclear how not self.low_precision_master_weights_and_grads conditional affects gradient accumulation boundary handling

[stas00]

addressed

[stas00]

@sfc-gh-truwase, Trino brought up some performance concerns:

• Extra memory/copy: the legacy FP32-master path now uses a persistent host accumulator, even for one backward. Measured with 1,048,576 BF16 values: 2 MiB additional storage and one extra copy per call. Allocation happens once, not every step. The compact FP32 accumulation path bypasses this.

I'm going to try to fix it, so the one bwd per update path isn't paying the cost it doesn't need to.

[stas00]

@sfc-gh-truwase, ok, now the fix is only for multiple updates, so the normal path isn't impacted memory-wise.