Fix-FP16-ZeRO-Stage-0-loss-Scaling-In-engine.backward() - #8393
Conversation
![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: 857e98cbee
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # loss.backward() directly makes effective gradient updates ~cur_scale | ||
| # times too small, stalling training. Route through | ||
| # FP16_Optimizer.backward() which applies the scaling correctly. | ||
| self.optimizer.backward(loss, **backward_kwargs) |
There was a problem hiding this comment.
When gradient_accumulation_steps > 1, this early backward call and return bypass the graph registration at lines 3341–3344. The output hooks installed by forward() therefore make _backward_prologue_per_tensor() divide each output gradient by the accumulation count, even though the loss was already divided at line 3322, attenuating FP16 ZeRO-0 updates by an additional factor of gradient_accumulation_steps. Keep the scaled backward on the common registered path, or register the graph before invoking the optimizer.
Useful? React with 👍 / 👎.
| # TODO: handle these scaling with direct calls to loss.backward() | ||
| if isinstance(self.optimizer, ZeROOptimizer): | ||
| loss = self.optimizer.scale_if_loss(loss) | ||
| elif self.fp16_enabled() and not self.zero_optimization(): |
There was a problem hiding this comment.
This is a single-parent, non-merge commit, but its commit message has no Signed-off-by trailer, so it does not meet the repository's mandatory commit requirement and should be recreated with --signoff.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
| self.optimizer.backward(loss, **backward_kwargs) | ||
| self._backward_epilogue() | ||
| self._running_engine_backward = False | ||
| return gas_scaled_loss |
There was a problem hiding this comment.
When callers enable compiled_autograd_enabled=True through DeepSpeedEngine.compile(), every other supported backward path executes inside compiled_autograd(...), but this FP16 ZeRO-0 branch performs the actual backward and returns before entering that context. Training remains eager despite the explicitly enabled option, losing the expected compiled-backward behavior and performance for this configuration.
Useful? React with 👍 / 👎.
When fp16 is enabled with ZeRO stage 0, engine.backward() falls through to loss.backward() directly. FP16_Optimizer.step() always divides gradients by cur_scale (typically 32768), so gradients that were never pre-scaled are unscaled by the same factor, making effective parameter updates ~cur_scale times too small and stalling training. Add an explicit elif branch for fp16 + ZeRO-0 that routes through FP16_Optimizer.backward(), which pre-scales the loss by cur_scale before calling backward(), matching what step() expects. The existing # TODO comment on line 3325 acknowledges this path was left incomplete. The ZeRO stages > 0 and AMP paths are handled correctly; only the fp16 non-ZeRO path was missing. Tested with a CIFAR-10 MoE benchmark (fp16, ZeRO stage 0, 8 experts): - Before: accuracy 22-28%, loss stalls at ~1.8 - After: accuracy 49%, loss converges to ~0.7 Signed-off-by: Roohollah Etemadi <Roohollah.Etemadi@amd.com>
857e98c to
e614064
Compare
1 participant
When fp16 is enabled with ZeRO stage 0, engine.backward() falls through to loss.backward() directly. FP16_Optimizer.step() always divides gradients by cur_scale (typically 32768), so gradients that were never pre-scaled are unscaled by the same factor, making effective parameter updates ~cur_scale times too small and stalling training.
Add an explicit elif branch for fp16 + ZeRO-0 that routes through FP16_Optimizer.backward(), which pre-scales the loss by cur_scale before calling backward(), matching what step() expects.
The existing # TODO comment on line 3325 acknowledges this path was left incomplete. The ZeRO stages > 0 and AMP paths are handled correctly; only the fp16 non-ZeRO path was missing.
Tested with a CIFAR-10 MoE benchmark (fp16, ZeRO stage 0, 8 experts):