[VenusTZZ]

Stacked on #8962 (GPTBigCode AutoTP correctness fixes), which should be merged first. Only the last two commits are new. Child-only diff: VenusTZZ/DeepSpeed@1a19b1a...218cbf2. This PR stays Draft until #8962 merges; it will then be rebased onto master and marked ready.

Summary

AutoTP multi-query GPTBigCode keeps a full copy of the key/value (K/V) block on every tensor-parallel (TP) rank, so its gradient needs a TP all-reduce on every micro-batch. With gradient accumulation this repeats the same small collective for every micro-batch, gradient_accumulation_steps times per optimizer step. For ZeRO 0/1 the full window stays in param.grad, so the K/V rows can be all-reduced once at the accumulation boundary instead. This is a performance optimization follow-up to #8962, split out at the reviewer's request; #8962 keeps the correctness fixes only.

Motivation

The per-micro-batch reduce was introduced as a correctness fix in #8962. It is necessary for ZeRO 2/3 and for configurations that move earlier micro-batches out of param.grad, but for ZeRO 0/1 with gradient accumulation it adds 2 × layers all-reduces per micro-batch for a few hundred rows. Measured on a GPTBigCode multi-query attention (MQA) model (24 layers, hidden=2048, 268.8M parameters), TP=4, bf16, ZeRO-1, seq=1024, micro-batch=1. In the table, "accumulation steps" is the number of micro-batches per optimizer step (gradient_accumulation_steps):

Platform	Accumulation steps	Baseline: per micro-batch	Deferred to boundary	Step	Tokens/s
4×910B4	8	1,467 ms / 5,585 tok/s	1,271 ms / 6,443 tok/s	−13.3%	+15.4%
4×910B4	32	5,408 ms / 6,060 tok/s	4,508 ms / 7,269 tok/s	−16.6%	+20.0%
4×A100-SXM4-80GB	8	750.2 ms / 10,920 tok/s	676.1 ms / 12,117 tok/s	−9.9%	+11.0%
4×A100-SXM4-80GB	32	3,302.5 ms / 9,922 tok/s	2,939.3 ms / 11,148 tok/s	−11.0%	+12.4%

K/V all-reduce calls per optimizer step drop from 1,536 to 48 at 32 accumulation steps, and from 384 to 48 at 8. The gain is workload-dependent: 4 layers/hidden=4096 was within noise on NPU, and a single accumulation step has no gain by construction.

Design

• _ReplicatedKvGradReducer keeps its existing tensor hook and post-accumulate fallback. When the engine allows deferral, both hooks leave the local K/V partial in place; at the boundary the post-accumulate hook all-reduces the rows accumulated during the window.
• DeepSpeedEngine._backward_prologue() publishes the accumulation state to the reducers. Deferral is allowed only when the whole window stays in param.grad: ZeRO 0 (non-BF16 optimizer, no CPU offload) or ZeRO-1 with DeepSpeedZeroOptimizer and no separate gradient-accumulation buffer, more than one accumulation step (gradient_accumulation_steps() > 1), managed gradient accumulation, not inside no_sync/coalesced reduction, and not pipeline parallel. Every other configuration keeps the per-micro-batch path.
• The state is stored per reducer, so two engines in one process do not overwrite each other.
• The decision is latched per accumulation window. Deciding per backward would let one window mix the two paths when no_sync covers only the non-boundary micro-batches, and the boundary's whole-window reduce would count already-reduced micro-batches again.
• A reentrant checkpoint may run the boundary hook more than once; later calls reduce only the rows added since the previous call.

The fallback and the deferral share reducer state, so this branch depends on the reducer added in #8962. It is stacked and should be retargeted after #8962 lands.

Files changed

File	Why it changes	Diff
deepspeed/module_inject/layers.py	Deferral state on the K/V reducer and a delta-based boundary reduce.	+54/-11
deepspeed/runtime/engine.py	Publish the per-window accumulation state from the backward prologue.	+42/-1
tests/unit/v1/autotp/test_autotp_training.py	Reducer unit tests and two-micro-batch accumulation parity tests for ZeRO 0/1/2/3, low precision, and no_sync.	+236/-1

Testing

Environment

• NPU: 4×Ascend 910B4, CANN 8.5.1, torch 2.10.0+cpu, torch_npu 2.10.0.post7, transformers 5.19.0
• CUDA: 4×A100-SXM4-80GB, CUDA 12.6, torch 2.14.1+cu126, transformers 5.20.2

NPU

cd tests/unit
PYTHONPATH=$(git rev-parse --show-toplevel) HCCL_NPU_SOCKET_PORT_RANGE=auto \
python -m pytest -q --forked \
  v1/autotp/test_autotp_training.py::TestGPTBigCodeMultiQueryTraining \
  v1/autotp/test_autotp_training.py::TestGPTBigCodeMultiQueryAccumulation \
  v1/autotp/test_autotp_training.py::test_replicated_kv_gradient_defers_to_accumulation_boundary \
  v1/autotp/test_autotp_training.py::test_replicated_kv_gradient_deferred_boundary_reduces_new_rows

Result: 15 passed in 358.76 s. Coverage:

• TestGPTBigCodeMultiQueryAccumulation: ZeRO 0/1/2/3 with and without the canonical distributed data parallel (DDP) no_sync pattern, plus fp16/bf16 with fp32 gradient accumulation on ZeRO-1. The test asserts both loss parity against the unsharded reference and the per-step K/V all-reduce count (2/2 per step for deferred windows, 4/4 for fallback windows), so it fails if the feature silently falls back or if the two paths mix.
• Reducer unit tests: deferral to the boundary, the reentrant boundary delta, and both hook paths.
• TestGPTBigCodeMultiQueryTraining: single-micro-batch regression (no gradient accumulation), 5 passed.

The tree is unchanged from the implementation before the split (218cbf242^{tree} equals the pre-split e02f9f3c0^{tree}), so the earlier NPU and A100 performance validation still applies.

CUDA

cd tests/unit
CUDA_VISIBLE_DEVICES=3,4,5,7 \
PYTHONPATH=$(git rev-parse --show-toplevel) \
python -m pytest -q --forked \
  v1/autotp/test_autotp_training.py::TestGPTBigCodeMultiQueryTraining \
  v1/autotp/test_autotp_training.py::TestGPTBigCodeMultiQueryAccumulation \
  v1/autotp/test_autotp_training.py::test_replicated_kv_gradient_defers_to_accumulation_boundary \
  v1/autotp/test_autotp_training.py::test_replicated_kv_gradient_deferred_boundary_reduces_new_rows

Result: 15 passed, 1 warning in 186.35 s (0:03:06) at split tip 218cbf242. Same 15-test set as NPU, including the loss-parity and per-step all-reduce-count assertions. This closes the previous CUDA test gap; the split only rewrote history, so the A100 performance numbers measured pre-split apply to this branch.

Not tested

• Full CUDA AutoTP suite: only the targeted GPTBigCode training/accumulation tests above were rerun after the split.
• ZeRO-2/3 deferral: they keep the per-micro-batch path by design.
• torch.compile, pipeline parallelism, and windows that use coalesce_grad_reduction: these take the fallback path and are not benchmarked.
• NVMe offload: excluded from deferral by the gate; not run.

Dependencies

• Stacked on Fix AutoTP multi-query GPTBigCode training at TP > 1 #8962; merge it first, then this PR. Afterward: git rebase --onto upstream/master 1a19b1aab, force-push, mark ready.

Notes for review

• No new config flag; the optimization is automatic for the safe cases and falls back otherwise.
• The K/V access pattern and the collective shape are unchanged; only the number of calls changes.
• AI assistance: the change was drafted with Claude Code; the author reviewed the diff and ran the NPU and CUDA validation above.

Checklist

• NPU validation (4×910B4)
• CUDA validation after split (4×A100)
• Performance measurement (NPU + A100)
• pre-commit (implementation before the split; the split only rewrote history)