Fix shared loss gradient accumulation

[tohtana]

Problem

Fixes #8224.

With gradient accumulation enabled, a loss can combine a forward pass through the DeepSpeed engine with a forward pass through the model inside the engine. The engine-output hook divides only the gradient from the engine forward by the accumulation count. The gradient from the forward pass on the model inside the engine remains unscaled, producing an incorrect gradient for that parameter.

Approach

Apply gradient-accumulation scaling to the complete loss passed through managed engine.backward. During that call, mark the loss graph as already scaled so engine-output hooks do not apply the scaling again, and restore all managed-backward state on every exit.

If a managed backward is interrupted, preserve the existing ZeRO retry behavior by not running a direct-backward epilogue over incomplete reduction state. Direct tensor backward, pipeline output hooks, and scale_wrt_gas=False otherwise retain their existing behavior.

Testing

• pytest -q tests/unit/runtime/zero/test_zero_shared_loss_gradient.py
• Verified the existing ZeRO-3 exception-retry lifecycle behavior.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 92302850b6

ℹ️ 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".

[chatgpt-codex-connector]

Prevent double-scaling branches from other engines

When engine_a.backward(loss_a + loss_b) is used with multiple DeepSpeed engines and gradient accumulation greater than one, this divides the entire combined loss by engine A's accumulation count, but only engine A has _is_engine_backward_loss_scaled set. Engine B's output hook therefore still divides its incoming gradient in _backward_prologue_per_tensor, silently scaling engine B's gradient twice (by an extra 1 / GAS when the engines use the same count). Combined multi-engine autograd graphs are already supported by TestZeroUserBackwardMultipleEngines, so the suppression needs to cover every participating engine or preserve branch-specific scaling.

Useful? React with 👍 / 👎.

[pengdurice]

Does it work for zero3? Or does zero3 has this problem? the original issue uses only zero2.