Fix repeated gradient scaling across pipeline stages (#8154)

## 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>

Author: Functionhx

deepspeed/runtime/pipe/engine.py

@@ -534,6 +534,12 @@ def is_last_stage(self):
         """True if this process is in the last stage in the pipeline."""
         return self.stage_id == self.num_stages - 1
 
+    def _backward_prologue_per_tensor(self, grad):
+        # The last stage applies GAS scaling before sending gradients upstream.
+        if self.is_last_stage():
+            return super()._backward_prologue_per_tensor(grad)
+        return grad
+
     def get_pipeline_parallel_rank(self):
         return self.stage_id
 

tests/unit/runtime/pipe/test_pipe.py

@@ -13,6 +13,7 @@
 import deepspeed.comm as dist
 from deepspeed.runtime.pipe.topology import PipeDataParallelTopology
 from deepspeed.runtime.pipe.module import PipelineModule
+from deepspeed.utils import RepeatingLoader
 from unit.alexnet_model import AlexNetPipe, train_cifar
 from unit.common import DistributedTest
 from unit.util import skip_on_arch, no_child_process_in_deepspeed_io
@@ -49,6 +50,44 @@ def rel_diff(A, B):
     return abs(A - B) / abs(A)
 
 
+class TestPipeGradientAccumulationScaling(DistributedTest):
+    world_size = 2
+
+    def test_gradients_are_scaled_once(self):
+        gas = 2
+        learning_rate = 0.1
+        config = {
+            "train_batch_size": gas,
+            "train_micro_batch_size_per_gpu": 1,
+            "gradient_accumulation_steps": gas,
+            "gradient_clipping": 0.0,
+            "optimizer": {
+                "type": "SGD",
+                "params": {
+                    "lr": learning_rate
+                }
+            },
+            "zero_allow_untested_optimizer": True,
+            "pipeline": {
+                "activation_checkpoint_interval": 0
+            },
+        }
+
+        layers = [nn.Linear(1, 1, bias=False), nn.Linear(1, 1, bias=False)]
+        for layer in layers:
+            layer.weight.data.fill_(1.0)
+
+        model = PipelineModule(layers=layers, num_stages=2, loss_fn=nn.MSELoss())
+        engine, _, _, _ = deepspeed.initialize(config=config, model=model, model_parameters=model.parameters())
+        engine.set_dataiterator(RepeatingLoader([(torch.ones(1, 1), torch.zeros(1, 1))]))
+
+        engine.train_batch()
+
+        expected_weight = torch.tensor([[1.0 - 2.0 * learning_rate]], device=engine.device)
+        actual_weight = next(engine.module.parameters())
+        assert torch.allclose(actual_weight, expected_weight)
+
+
 @pytest.mark.parametrize('topo_config', [
     {
         "num_pp": 1,