Add configurable sum gradient reduction

[zupengwang]

Summary

• add a gradient_allreduce_op configuration with "mean" as the default and "sum" as the new option
• support unscaled gradient sums for ZeRO stages 0, 1, and 2 across reduce-scatter, allreduce, and non-contiguous fallback paths
• reject unsupported ZeRO stage 3, ZenFlow, and DeepCompile combinations with clear configuration errors
• document the option and preserve existing mean-reduction behavior

Addresses #7107.

Motivation

Some distributed objectives, including contrastive learning over globally gathered embeddings, require summing data-parallel gradients rather than averaging them. Today users need to rescale the loss manually to cancel DeepSpeed's world-size normalization.

This change makes the reduction semantics explicit while keeping the current behavior as the default.

Validation

Tested on two NVIDIA GeForce RTX 3090 GPUs with PyTorch 2.13.0+cu130:

• pytest --forked -q tests/unit/v1/zero/test_zero.py::TestGradientAllreduceOp — 18 passed
• targeted configuration tests in tests/unit/runtime/test_ds_config_dict.py — 8 passed; the new ZeRO-1/2 DeepCompile cases both failed against the prior head because no error was raised
• pytest --forked -q tests/unit/v1/zero/test_zero_coalesce_grad_reduction.py::TestCoalesceCombinations — 12 passed
• pytest --forked -q tests/unit/runtime/sparse_tensor/test_averaging_sparse_gradients.py — 1 passed
• changed-file pre-commit hooks, including YAPF, flake8, codespell, license, check-torchdist, and check-torchcuda — passed

The distributed test matrix covers ZeRO stages 0/1/2, mean and sum reductions, reduce-scatter, gradient predivide, prescale, and non-contiguous gradient fallback.

Real-training equivalence

A deterministic two-rank, five-step SimpleModel regression compares the default MEAN reduction with SUM while normalizing only the SUM backward loss by world_size, so both modes provide identical gradients to the optimizer.

• AdamW: ZeRO-0/1/2
• Muon: ZeRO-1/2
• targeted gradient-reduction tests: 23 passed
• complete tests/unit/v1/zero/test_zero.py: 98 passed, 1 skipped

Across all five optimizer/stage configurations and all five training steps, the observed loss, full-gradient, and full-parameter differences were zero.

Limitations

gradient_allreduce_op="sum" is intentionally not supported with ZeRO stage 3, ZenFlow, or DeepCompile. These combinations fail during configuration instead of silently applying mean semantics.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1a3b05de83

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

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Propagate sum mode into DeepCompile reductions

When ZeRO stage 1 or 2 runs with an active DeepCompile backend, allreduce_gradients() returns at engine.py:2842 before any of the new sum-aware paths execute. DeepCompile instead keeps pre_div_reduce = true (csrc/compile/deepcompile.cpp:21) and unconditionally divides every bucket by the process-group size in CustomOpExecutor::applyPreDivision() (csrc/includes/deepcompile.h:590-600), so gradient_allreduce_op="sum" silently still produces a mean and scales updates down by the data-parallel world size. Pass this mode into the compiled reducer or reject the combination during configuration.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

The reviewed commit is a non-merge commit, but its message contains no Signed-off-by trailer. Add the trailer using the configured Git identity before merging so the commit satisfies the repository's commit and CI requirements.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.

[pengdurice]

any reason why zero3 is not supported at this moment?

[zupengwang]

Yes. I limited this PR to ZeRO stages 0-2 because ZeRO-3 uses a separate reduction and partition path: reduce_scatter_coalesced() pre-divides by world size, and the Stage 3 optimizer currently does not receive gradient_average. Supporting sum safely would require plumbing the mode through the Stage 3 optimizer and auditing its regular, contiguous, and quantized/all-to-all reduction paths, with dedicated ZeRO-3 tests. For this PR I chose to reject the combination during config validation instead of silently returning mean semantics. I can handle ZeRO-3 in a follow-up if that scope is desired.

[pengdurice]

is this where we divide by dp world size to make sure the gradients used for optimizer is still scaled?

[zupengwang]

Yes—this is the sparse-gradient path. With the default mean mode, self.gradient_average is true, so the values are divided by the data-parallel world size before all-gather. With gradient_allreduce_op="sum", self.gradient_average is false, this block is skipped, and the optimizer receives the unnormalized sum. The dense and ZeRO stages 1/2 paths use the same guard.

[pengdurice]

is this overall simplistic? maybe add a real training as well and compare the losses are close for each step with the same model with or without this knob on?

[pengdurice]

at the same time, maybe a loss comparison figure with our without enabling this will be useful.

[pengdurice]

this change should be optimizer agnostic. should we test that using different optimizers, e.g. AdamW and Muon, things all work fine?

[zupengwang]

@pengdurice Thanks — I agree the original single-step SGD test was too narrow.
I added a deterministic real-training regression using two data-parallel ranks and five steps. It compares the raw forward loss at every step, the full gradients before the optimizer step, and the full FP32 parameters after every step.
One normalization detail is important here: SUM produces world_size times the MEAN gradient, so an unnormalized SUM run is not generally expected to follow the same trajectory at the same learning rate. Applying lr / world_size is also not optimizer agnostic for adaptive or nonlinear optimizers. The new test therefore leaves the forward loss unchanged and divides only the SUM backward loss by world_size, which gives both runs identical optimizer inputs. The existing low-level test remains in place and independently verifies the raw unnormalized SUM = world_size × MEAN contract.
The matrix covers AdamW with ZeRO-0/1/2 and Muon with ZeRO-1/2. Muon ZeRO-0 is intentionally excluded because the current documented/runtime-supported Muon path starts at ZeRO-1.
On two RTX 3090s, the combined targeted tests passed 23/23 and the complete changed test file passed 98 tests with one existing conditional skip. The five training configurations had zero observed loss, full-gradient, and full-parameter differences across all five steps. I also prepared the attached loss-comparison figure showing the DP-averaged raw forward losses for MEAN and normalized-backward SUM.

[pengdurice]

@pengdurice Thanks — I agree the original single-step SGD test was too narrow. I added a deterministic real-training regression using two data-parallel ranks and five steps. It compares the raw forward loss at every step, the full gradients before the optimizer step, and the full FP32 parameters after every step. One normalization detail is important here: SUM produces world_size times the MEAN gradient, so an unnormalized SUM run is not generally expected to follow the same trajectory at the same learning rate. Applying lr / world_size is also not optimizer agnostic for adaptive or nonlinear optimizers. The new test therefore leaves the forward loss unchanged and divides only the SUM backward loss by world_size, which gives both runs identical optimizer inputs. The existing low-level test remains in place and independently verifies the raw unnormalized SUM = world_size × MEAN contract. The matrix covers AdamW with ZeRO-0/1/2 and Muon with ZeRO-1/2. Muon ZeRO-0 is intentionally excluded because the current documented/runtime-supported Muon path starts at ZeRO-1. On two RTX 3090s, the combined targeted tests passed 23/23 and the complete changed test file passed 98 tests with one existing conditional skip. The five training configurations had zero observed loss, full-gradient, and full-parameter differences across all five steps. I also prepared the attached loss-comparison figure showing the DP-averaged raw forward losses for MEAN and normalized-backward SUM.

Thank you for the reply!
Can we do this?

1. maybe change the unit test to use SimpleModel for multiple steps of training similar to

   DeepSpeed/tests/unit/v1/zero/test_zero_cpu_offload_grad_accum.py

   Line 62 in 7904603

   model = SimpleModel(hidden_dim, nlayers=2)

   then assert parameter equals with our without the knob.
2. Add the figures above to the PR description.
   After that, lgtm.

[zupengwang]

@pengdurice Done — the follow-up now uses the repository SimpleModel for deterministic two-rank, five-step training and compares the full FP32 parameters after every step with the knob set to MEAN versus SUM (with SUM backward normalized by world_size). The matrix covers AdamW with ZeRO-0/1/2 and Muon with ZeRO-1/2.

I also added the loss-comparison figure and the updated validation results to the PR description. The targeted tests pass 23/23, and the complete changed test file passes 98 tests with one conditional skip. Thanks!

[pengdurice]

@pengdurice Done — the follow-up now uses the repository SimpleModel for deterministic two-rank, five-step training and compares the full FP32 parameters after every step with the knob set to MEAN versus SUM (with SUM backward normalized by world_size). The matrix covers AdamW with ZeRO-0/1/2 and Muon with ZeRO-1/2.

I also added the loss-comparison figure and the updated validation results to the PR description. The targeted tests pass 23/23, and the complete changed test file passes 98 tests with one conditional skip. Thanks!

would you mind push your change to remote? I cannot find new commit yet

[zupengwang]

Thanks! The change has been pushed to the PR's fork branch (zupengwang:feature/gradient-allreduce-sum) and is available at zupengwang@f612ed9. The PR head is now f612ed9605bce88b89c575b4b59d4fb344c1e1c4, and all CI checks are green.