Add configurable sum gradient reduction (#8232)
## 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.
<img width="2496" height="1572" alt="loss_comparison"
src="https://github.com/user-attachments/assets/3532ad11-5297-42db-baca-8a632bf793ad"
/>
## 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.
---------
Signed-off-by: Wang Zupeng <zupenwang@gmail.com>
