Add configurable sum gradient reduction - #8232
Conversation
Signed-off-by: Wang Zupeng <zupenwang@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 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".
| self._unmanaged_backward_count = 0 | ||
| self.skipped_steps = 0 | ||
| self.gradient_average = True | ||
| self.gradient_average = config_class.gradient_allreduce_op == GRADIENT_ALLREDUCE_OP_MEAN |
There was a problem hiding this comment.
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 👍 / 👎.
| return get_scalar_param(param_dict, GRADIENT_PREDIVIDE_FACTOR, GRADIENT_PREDIVIDE_FACTOR_DEFAULT) | ||
|
|
||
|
|
||
| def get_gradient_allreduce_op(param_dict): |
There was a problem hiding this comment.
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 👍 / 👎.
| <= ZeroStageEnum.max_stage), "DeepSpeedConfig: Maximum supported ZeRO stage is {}".format( | ||
| ZeroStageEnum.max_stage) | ||
|
|
||
| if (self.gradient_allreduce_op == GRADIENT_ALLREDUCE_OP_SUM |
There was a problem hiding this comment.
any reason why zero3 is not supported at this moment?
| else: | ||
| values.mul_(1. / (dp_world_size)) | ||
| else: | ||
| values.mul_(1. / (dp_world_size)) |
There was a problem hiding this comment.
is this where we divide by dp world size to make sure the gradients used for optimizer is still scaled?
2 participants
Summary
gradient_allreduce_opconfiguration with"mean"as the default and"sum"as the new optionAddresses #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 passedtests/unit/runtime/test_ds_config_dict.py— 6 passedpytest --forked -q tests/unit/v1/zero/test_zero_coalesce_grad_reduction.py::TestCoalesceCombinations— 12 passedpytest --forked -q tests/unit/runtime/sparse_tensor/test_averaging_sparse_gradients.py— 1 passedpre-commithooks, including YAPF, flake8, codespell, license,check-torchdist, andcheck-torchcuda— passedThe distributed test matrix covers ZeRO stages 0/1/2, mean and sum reductions, reduce-scatter, gradient predivide, prescale, and non-contiguous gradient fallback.
Limitations
gradient_allreduce_op="sum"is intentionally not supported with ZeRO stage 3 or ZenFlow. These combinations fail during configuration instead of silently applying mean semantics.