Skip to content

Fix-FP16-ZeRO-Stage-0-loss-Scaling-In-engine.backward() #8393

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
etemadiamd wants to merge 1 commit into
base: master
Choose a base branch
from
+10 −0
Open

Fix-FP16-ZeRO-Stage-0-loss-Scaling-In-engine.backward() #8393

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions deepspeed/runtime/engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3325,6 +3325,16 @@ def backward(self, loss, retain_graph=False, scale_wrt_gas=True):
# 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():

@chatgpt-codex-connector chatgpt-codex-connector Bot Sep 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Add the mandatory sign-off trailer

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 👍 / 👎.

# fp16 + ZeRO stage 0: FP16_Optimizer.step() always divides gradients
# by cur_scale, so loss must be pre-scaled before backward. Calling
# 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)

@chatgpt-codex-connector chatgpt-codex-connector Bot Sep 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Preserve graph registration before FP16 backward

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 👍 / 👎.

self._backward_epilogue()
self._running_engine_backward = False
return gas_scaled_loss
Comment on lines +3334 to +3337

@chatgpt-codex-connector chatgpt-codex-connector Bot Sep 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Keep FP16 backward inside the compiled-autograd context

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 👍 / 👎.

elif self.torch_autocast_z0_gradscaler:
loss = self.torch_autocast_z0_gradscaler.scale(loss)

Expand Down
Loading