Skip to content

Fix repeated gradient scaling across pipeline stages#8154

Open
Functionhx wants to merge 1 commit into
deepspeedai:masterfrom
Functionhx:fix/pipeline-gradient-scaling
Open

Fix repeated gradient scaling across pipeline stages#8154
Functionhx wants to merge 1 commit into
deepspeedai:masterfrom
Functionhx:fix/pipeline-gradient-scaling

Conversation

@Functionhx

Copy link
Copy Markdown

Summary

Fix repeated gradient accumulation scaling when pipeline parallelism is combined with gradient accumulation.

Every pipeline stage registers an output backward hook that scales gradients by gradient_accumulation_steps. The last stage therefore scales the gradient correctly before sending it upstream, but each preceding stage scales the already-scaled gradient again.

Fixes #8152

Changes

  • Apply gradient accumulation scaling only at the last pipeline stage.
  • Preserve the existing behavior for single-stage pipelines and gradient_accumulation_steps=1.
  • Add a deterministic two-stage regression test that verifies both stages receive correctly averaged gradients.
  • Disable gradient clipping in the regression test so it cannot mask the scaling difference.

Testing

  • DS_ACCELERATOR=cpu LOCAL_SIZE=2 OMP_NUM_THREADS=1 pytest -q tests/unit/runtime/pipe/test_pipe.py::TestPipeGradientAccumulationScaling::test_gradients_are_scaled_once
  • pre-commit run --files deepspeed/runtime/pipe/engine.py tests/unit/runtime/pipe/test_pipe.py
  • Two-process CPU/Gloo numerical comparison for gradient_accumulation_steps=1,2,8

Signed-off-by: Yuchen Fan <functionhx@gmail.com>
@Functionhx
Functionhx marked this pull request as ready for review July 18, 2026 04:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Gradient accumulation scaling is applied repeatedly across pipeline stages

1 participant