Scale ZeRO-3 __avg_scatter_grads by sequence_parallel_size

[deep1blue]

Summary

Fixes #8692.

In DeepSpeedZeroOptimizer_Stage3, __avg_scatter_contiguous_grads (stage3.py:1877) and allreduce_bucket (stage3.py:2278) both divide the gradient averaging world size by float(self.sequence_parallel_size), matching ZeRO-0 (engine.py:3884) and ZeRO-1/2 (stage_1_and_2.py:1505, 1687, 2379).

However, __avg_scatter_grads (stage3.py:1936-1940) delegates to reduce_scatter_coalesced (which divides by partition_world_size) and only rescales when averaging_world_size != partition_world_size without accounting for self.sequence_parallel_size:

partition_world_size = dist.get_world_size(group=process_group)
averaging_world_size = self._gradient_averaging_world_size(
    params_to_reduce, partition_world_size) / float(self.sequence_parallel_size)
if averaging_world_size != partition_world_size:
    scale = partition_world_size / float(averaging_world_size)
    grad_partitions_for_rank = [g.mul(scale) for g in grad_partitions_for_rank]

Without this division by float(self.sequence_parallel_size):

1. Under sequence_parallel_size > 1, ZeRO-3 with reduce_scatter=True or contiguous_gradients=False produces gradients smaller by a factor of sequence_parallel_size than reduce_scatter=False, contiguous_gradients=True or ZeRO-1/2.
2. Even within a single ZeRO-3 model using reduce_scatter=False, contiguous_gradients=True, parameters whose size exceeds reduce_bucket_size fall back to __avg_scatter_grads (stage3.py:1686), while smaller parameters use __avg_scatter_contiguous_grads (stage3.py:1683), causing different layers of the same model to receive gradients scaled by 1 / sp_size vs 1.0 in the same backward step.

Non-duplication Check

• Checked Gradient scale under sequence parallelism differs between ZeRO reduction paths #8692 (closed_by_pull_requests: 0) and searched open/closed PRs touching stage3.py, __avg_scatter_grads, and sequence_parallel_size.
• Related PRs do not touch this path:
  • Apply the ZeRO-3 predivide postscale when the factor equals world size #8668 fixes gradient_predivide_factor == partition_world_size post-scaling in stage3.py.
  • Apply the ZeRO-1/2 predivide postscale under sequence parallelism #8687 fixes gradient_predivide_factor post-scaling in stage_1_and_2.py.
  • [AutoSP] Aggregate the causal LM loss and gradients over the SP group #8457 and [AutoSP] Align cross entropy targets after an in-model causal label shift #8811 only touch compile/ AutoSP passes and groups.py.

Validation & Test Results

Added TestZero3SequenceParallelGradientScale in tests/unit/v1/zero/test_zero.py covering all four ZeRO-3 reduction configurations under sequence_parallel_size=2 ((reduce_scatter=False, contiguous_gradients=True, reduce_bucket_size=500), (False, False, 500), (True, True, 500), and mixed-bucket (False, True, 16) where small takes __avg_scatter_contiguous_grads and large takes __avg_scatter_grads).

Commands Run

1. New regression test (TestZero3SequenceParallelGradientScale):

   PYTHONPATH=.:"$PWD/tests" DS_ACCELERATOR=cpu LOCAL_SIZE=2 pytest tests/unit/v1/zero/test_zero.py -k "TestZero3SequenceParallelGradientScale" -v

   • Before fix: 3 failed (test[False-False-500], test[True-True-500], test[False-True-16]), 1 passed (test[False-True-500]).
   • After fix: 4 passed in 23.18s.

2. Existing ZeRO gradient reduction tests (TestGradientAllreduceOp):

   PYTHONPATH=.:"$PWD/tests" DS_ACCELERATOR=cpu LOCAL_SIZE=2 pytest tests/unit/v1/zero/test_zero.py -k "TestGradientAllreduceOp and not Training" -v

   • Result: 18 passed in 90.64s.

Compatibility Matrix

Feature / Path	Validated Locally	Notes
ZeRO-3 (reduce_scatter=False, contiguous_gradients=True)	Yes (2-rank CPU/Gloo)	Matches 1.0 scale
ZeRO-3 (reduce_scatter=False, contiguous_gradients=False)	Yes (2-rank CPU/Gloo)	Fixed from 1/sp to 1.0
ZeRO-3 (reduce_scatter=True, contiguous_gradients=True)	Yes (2-rank CPU/Gloo)	Fixed from 1/sp to 1.0
ZeRO-3 mixed bucket (param.ds_numel > reduce_bucket_size)	Yes (2-rank CPU/Gloo)	Both small and large params get 1.0 scale
ZeRO-0 / ZeRO-1 / ZeRO-2	Yes (2-rank CPU/Gloo)	Unchanged (18 passed)
Multi-node all_to_all_quant_reduce	Not run locally (single-node host)	Uses same post-reduction averaging_world_size rescale

Environment (ds_report)

• OS: Linux aarch64 (Debian 13)
• Python: 3.13.5, PyTorch: 2.14.1+cpu (DS_ACCELERATOR=cpu, Gloo backend, 2 processes)

AI Assistance Disclosure

AI assistance (GitHub Copilot) was used to help draft the regression test and PR description; all root-cause analysis and test executions were run and verified locally.

Credits & Acknowledgments

• Full credit to @alanhuangyoo for reporting Gradient scale under sequence parallelism differs between ZeRO reduction paths #8692, diagnosing the exact gradient-scaling discrepancy between __avg_scatter_grads and __avg_scatter_contiguous_grads / allreduce_bucket under sequence parallelism (including the mixed-bucket single-model case), and providing the 2-rank reproduction scripts; and to @Eliaaazzz for citing Fix the sequence-parallelism for the dense model architecture #4530 and _reduce_non_expert_gradients confirming the DP-world averaging convention.
• If @alanhuangyoo would like to submit their own PR as offered in Gradient scale under sequence parallelism differs between ZeRO reduction paths #8692, we are more than happy to yield/close this PR in favor of theirs or add Co-authored-by:.

[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-18T03:39:27.833856Z	086cf17	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.