Skip to content

Fix repeated gradient scaling across pipeline stages - #8154

Merged
tohtana merged 1 commit into
deepspeedai:masterfrom
Functionhx:fix/pipeline-gradient-scaling
Jul 20, 2026
Merged

Fix repeated gradient scaling across pipeline stages#8154
tohtana merged 1 commit into
deepspeedai:masterfrom
Functionhx:fix/pipeline-gradient-scaling

Conversation

@Functionhx

Copy link
Copy Markdown
Contributor

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

@tohtana tohtana left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @Functionhx,
Thank you for submitting this PR! This looks good to me.
I know this type of bug is very hard to find and fix. I appreciate your continuous contribution.

@tohtana
tohtana enabled auto-merge July 20, 2026 22:20

@PKUWZP PKUWZP left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks for the PR!

@tohtana
tohtana added this pull request to the merge queue Jul 20, 2026
Merged via the queue into deepspeedai:master with commit d99a642 Jul 20, 2026
13 of 15 checks passed
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

3 participants