fix(training): EMA updates on skipped optimizer step with AMP by nrenvo · Pull Request #15173 · huggingface/diffusers

nrenvo · 2026-11-03T11:46:01Z

Problem

When AMP skips an optimizer step due to gradient overflow, EMA still updated with unchanged weights. Also, train_loss was reset inside the skipped-step guard, so the accumulator wasn't cleared on overflow.

An overflow at the last batch of an epoch also caused the epoch loop to exit with global_step < max_train_steps.

Fix

Guard EMA update and step bookkeeping with not accelerator.optimizer_step_was_skipped. Move train_loss = 0.0 outside the guard so it resets on every gradient sync.

Replace for epoch in range(...) with while global_step < args.max_train_steps to ensure training always completes the requested steps.

Test

Added a parametrized test that simulates an overflow on step 2 and asserts EMA is never called for a skipped step.

Added test_text_to_image_ema_skipped_step_at_epoch_boundary to cover overflow at epoch boundary.

Fixes #15112

@sayakpaul @linoytsaban @ZihaoW123

…uggingface#15112) When AMP causes a gradient overflow, accelerate skips the optimizer step but the previous code still updated EMA, progressed the step counter, and logged metrics. Guard the EMA update (and all per-step bookkeeping) with `not accelerator.optimizer_step_was_skipped`. Also move `train_loss = 0.0` outside that guard so the loss accumulator is still cleared on any `sync_gradients` boundary, including overflowed ones. Adds a parametrized unit test that monkeypatches AcceleratedOptimizer to simulate an overflow on step 2 and asserts EMA is never called for a skipped step while the applied-step count stays correct.

…izer step When AMP skips an optimizer step at the last batch of an epoch, the static `num_train_epochs` ceiling (computed before any overflow occurs) caused the outer epoch loop to exit with `global_step < max_train_steps`. Replace `for epoch in range(first_epoch, args.num_train_epochs)` with a `while global_step < args.max_train_steps` loop so the training script always completes the requested number of successful optimizer steps, regardless of where in an epoch the overflow happens. Add `test_text_to_image_ema_skipped_step_at_epoch_boundary` to pin this behaviour: 2 steps/epoch, overflow on step 2 (epoch boundary), training must still reach global_step == 2 by spilling into a third epoch.

sayakpaul

Please explain the changes and how they relate to EMA updates.

sayakpaul · 2026-11-05T07:35:10Z

@@ -35,6 +46,174 @@


 class TestTextToImage(ExamplesTestsAccelerate):


Could the tests go to https://github.com/huggingface/diffusers/blob/main/tests/others/test_ema.py?

Also, do the other training scripts need to propagate this update too?

The tests here launch the full training script end-to-end via runpy, monkeypatching Accelerator, AcceleratedOptimizer, and EMAModel together. test_ema.py covers EMAModel in isolation (shadow params, decay, serialization) so they feel like different things. Keeping them next to the script they exercise seemed more natural.

Propagated the fix to all other scripts using EMAModel in this PR: train_text_to_image_sdxl, train_instruct_pix2pix, train_instruct_pix2pix_sdxl, train_unconditional, train_amused, train_vqgan.

sayakpaul · 2026-11-05T07:35:43Z


-            # Checks if the accelerator has performed an optimization step behind the scenes
-            if accelerator.sync_gradients:
+            if accelerator.sync_gradients and not accelerator.optimizer_step_was_skipped:


optimizer_step_was_skipped -- does it warrant any minimum version requirement checks?

Present since at least accelerate v0.10.0, so no version bump needed — the examples requirements.txt already requires >=0.16.0.

sayakpaul · 2026-11-05T07:35:57Z

                progress_bar.update(1)
                global_step += 1
                accelerator.log({"train_loss": train_loss}, step=global_step)
-                train_loss = 0.0


Could you explain this change?

train_loss should reset on every sync_gradients, including skipped steps — otherwise the loss from a skipped step bleeds into the next logged value. Moving it outside the skipped-step guard ensures it always resets on a gradient sync.

sayakpaul · 2026-11-05T07:36:50Z

-    for epoch in range(first_epoch, args.num_train_epochs):
+    epoch = first_epoch
+    while global_step < args.max_train_steps:


How do these changes relate to EMA updates?

Not directly related to EMA. It fixes a separate issue: if an overflow hits the last batch of an epoch, global_step is not incremented and the for-loop exits before reaching max_train_steps. The while loop ensures training always completes the requested number of steps.

nrenvo · 2026-11-05T08:41:41Z

Two independent fixes:

EMA skip guard: when AMP skips an optimizer step due to gradient overflow, weights are unchanged. Updating EMA with unchanged weights biases the moving average, so EMA should only step when the optimizer actually applied gradients.
Epoch boundary fix: if overflow hits the last batch of an epoch, global_step is not incremented and the for-loop exits before reaching max_train_steps. The while loop ensures training always completes the requested steps. Unrelated to EMA correctness but discovered in the same context.

sayakpaul · 2026-11-05T08:50:52Z

Two independent fixes:

Can we split it into two PRs then?

nrenvo · 2026-11-05T10:04:43Z

Split into two PRs as requested:

fix(training): guard EMA updates on skipped optimizer step with AMP #15189 — EMA updates guarded on skipped optimizer step (all affected training scripts)
fix(training): continue training past epoch boundary on skipped optimizer step #15190 — epoch boundary fix (while loop replacing for-epoch loop)

nrenvo · 2026-11-09T14:06:39Z

@sayakpaul hi, following up on your suggestion from a few days ago — split #15173 into two separate PRs: #15189 (EMA guard) and #15190 (epoch boundary fix).

sayakpaul

Let's remove the test. It's verbose and tests something that is not that much used by the community through our scripts.

nrenvo · 2026-11-10T11:22:54Z

@sayakpaul thanks for the approval! Removed the tests from #15189 and #15190 as suggested. Just wanted to check — would you like to proceed with the two split PRs, or merge #15173 directly?

nrenvo added 2 commits November 2, 2026 02:07

github-actions Bot added examples fixes-issue size/M PR with diff < 200 LOC labels Nov 3, 2026

sayakpaul reviewed Nov 5, 2026

View reviewed changes

This was referenced Nov 5, 2026

fix(training): guard EMA updates on skipped optimizer step with AMP #15189

Open

fix(training): continue training past epoch boundary on skipped optimizer step #15190

Open

sayakpaul approved these changes Nov 10, 2026

View reviewed changes

		@@ -35,6 +46,174 @@


		class TestTextToImage(ExamplesTestsAccelerate):

Conversation

nrenvo commented Nov 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Fix

Test

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nrenvo commented Nov 5, 2026

Uh oh!

sayakpaul commented Nov 5, 2026

Uh oh!

nrenvo commented Nov 5, 2026

Uh oh!

nrenvo commented Nov 9, 2026

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

nrenvo commented Nov 10, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

nrenvo commented Nov 3, 2026 •

edited

Loading