Fix shared loss gradient accumulation - #8245
Conversation
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 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".
| backward_kwargs["retain_graph"] = True | ||
|
|
||
| self._running_engine_backward = False | ||
| loss = loss / self.gradient_accumulation_steps() if scale_wrt_gas else loss |
There was a problem hiding this comment.
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 👍 / 👎.
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Discard unfinished ZeRO hook, bucket, and unreduced gradient state at the next root forward when an aborted backward has not crossed a reduction boundary. Track actual bucket drains and fail closed before forward or step once an accumulation window may contain an irreversible contribution. Preserve step guards across runtime optimizer replacement and restore ZeRO-1/2 micro-step state for safe CPU-offload retry. Reject incomplete backward state at the public engine step entry, including non-boundary accumulation steps. Add lifecycle coverage for clean-control parity, oversized ZeRO-2 gradients, BF16 immediate updates, ZeRO-1 accumulation, coalesced ZeRO-2/3 gradients, multi-bucket ZeRO-2/3 failures, ZeRO-3 leaf reductions, step-before-forward rejection, CPU offload, and GAS>1 non-boundary step rejection. Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
1 participant
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.
This PR applies gradient-accumulation scaling to the complete loss passed through managed engine.backward. During that call, mark the loss as already scaled so engine-output hooks do not apply the scaling again, and restore the marker on every exit. The behavior of direct tensor backward, pipeline output hooks, and scale_wrt_gas=False is unchanged.