Text-to-image example updates EMA when AMP skips the optimizer step

Describe the bug

In train_text_to_image.py, EMA and global_step advance when AMP skips an optimizer update. The model parameters stay unchanged at that step, but the EMA parameters move toward them again and the EMA counter increments.

This is a controlled overflow test of the actual training example with a tiny public Stable Diffusion model. One gradient is set to infinity after backward pass 4. It does not establish how often this happens in normal training or its effect on model quality. All three controls below finished training, constructed the final pipeline, saved it, and exited with code 0.

The block after optimizer.step() checks only accelerator.sync_gradients. That remains true at the accumulation boundary when the prepared optimizer reports a skipped update.

I expected EMA and the training-step counter to advance only after an optimizer update was applied. Is counting skipped attempts here intentional, or should this block also check accelerator.optimizer_step_was_skipped?

Reproduction

The test uses Diffusers 48cdbaef5ee3d718aaee392b44f2099ae511c4a5, fp16, EMA, batch size 1, gradient accumulation 1, six requested steps, seed 13, and eight generated 64×64 RGB images. The loss scale starts at 16. Finite injects nothing; overflow and gated each inject one infinite gradient on backward pass 4. The recorder observes the real prepared optimizer and EMA calls.

Finite and overflow use the unmodified upstream training source. The gated comparison changes only the condition around the EMA/progress block:

-if accelerator.sync_gradients:
+if accelerator.sync_gradients and not accelerator.optimizer_step_was_skipped:

The model is hf-internal-testing/tiny-stable-diffusion-pipe at 3ee6c9f225f088ad5d35b624b6514b091e6a4849. Its model_index.json names Flax classes. The setup helper makes a separate local copy and changes those five class references to their PyTorch equivalents. It checks that every other file, including all weights, has the original hash, then constructs and saves the fixture on CPU before training. The original download stays unchanged. Training receives this local path without --revision.

Save the two files in the expandable sections below, then run these commands in a CUDA environment with PyTorch and Torchvision installed. Use a fresh results directory:

git init diffusers
git -C diffusers remote add origin https://github.com/huggingface/diffusers.git
git -C diffusers fetch --depth=1 origin 48cdbaef5ee3d718aaee392b44f2099ae511c4a5
git -C diffusers checkout --detach FETCH_HEAD
python - <<'PY'
import importlib.metadata
from pathlib import Path
Path('gpu_constraints.txt').write_text(''.join(
    name + '==' + importlib.metadata.version(name) + '\n'
    for name in ['torch', 'torchvision']))
PY
python -m pip install --no-build-isolation -c gpu_constraints.txt -e ./diffusers -r ./diffusers/examples/text_to_image/requirements.txt
mkdir -p fixture/results
python prepare_model.py --work ./fixture
python - <<'PY'
import subprocess
import sys
for case in ['finite', 'overflow', 'gated']:
    run = subprocess.run([
        sys.executable, 'reproduce_diffusers.py',
        '--checkout', './diffusers', '--work', './results/' + case,
        '--case', case,
        '--model-revision', '3ee6c9f225f088ad5d35b624b6514b091e6a4849',
        '--model-path', './fixture/pytorch_model',
    ])
    print(case, 'exit code:', run.returncode)
    if run.returncode:
        raise SystemExit(run.returncode)
PY

The recorder writes optimizer events, parameter-equality checks, EMA calls/counters, global_step, and final pipeline file sizes/hashes to results/<case>/result.json. Validation prompts are not enabled, so this does not test inference output.

Executed file: prepare_model.py

"""Copy the pinned tiny fixture and correct only its PyTorch pipeline metadata."""
import argparse
import hashlib
import json
from pathlib import Path
import shutil

REVISION = '3ee6c9f225f088ad5d35b624b6514b091e6a4849'
INDEX_SHA256 = 'fd74f494317246ac18f606a9bd77d8f7097d5d0a6c08cbb9ff569a264fc70928'
CLASSES = {
    'safety_checker': ('stable_diffusion', 'FlaxStableDiffusionSafetyChecker', 'StableDiffusionSafetyChecker'),
    'scheduler': ('diffusers', 'FlaxDDIMScheduler', 'DDIMScheduler'),
    'text_encoder': ('transformers', 'FlaxCLIPTextModel', 'CLIPTextModel'),
    'unet': ('diffusers', 'FlaxUNet2DConditionModel', 'UNet2DConditionModel'),
    'vae': ('diffusers', 'FlaxAutoencoderKL', 'AutoencoderKL'),
}


def hashes(path):
    return {str(p.relative_to(path)): hashlib.sha256(p.read_bytes()).hexdigest()
            for p in path.rglob('*') if p.is_file()}


def overlay(source, destination):
    before = hashes(source)
    assert before['model_index.json'] == INDEX_SHA256, 'Pinned metadata changed'
    shutil.copytree(source, destination, symlinks=False)
    original = json.loads((source / 'model_index.json').read_text())
    corrected = dict(original)
    changes = {}
    for name, (library, old, new) in CLASSES.items():
        assert original[name] == [library, old]
        corrected[name] = [library, new]
        changes[name] = {'before': original[name], 'after': corrected[name]}
    (destination / 'model_index.json').write_text(json.dumps(corrected, indent=2) + '\n')
    after = hashes(destination)
    assert before.keys() == after.keys()
    assert [name for name in before if before[name] != after[name]] == ['model_index.json']
    assert hashes(source) == before, 'Original cache must remain unchanged'
    return {'model': 'hf-internal-testing/tiny-stable-diffusion-pipe', 'revision': REVISION,
            'source_path': str(source.resolve()), 'path': str(destination.resolve()),
            'original_files': before, 'files': after, 'metadata_changes': changes,
            'original_metadata': original, 'corrected_metadata': corrected,
            'only_model_index_changed': True}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--work', type=Path, required=True)
    args = parser.parse_args()
    from huggingface_hub import snapshot_download
    from diffusers import StableDiffusionPipeline
    source = Path(snapshot_download('hf-internal-testing/tiny-stable-diffusion-pipe', revision=REVISION,
                  ignore_patterns=['*.msgpack', '*.onnx', '*.pb', '*.tflite', '*.h5', '*.ckpt']))
    receipt = overlay(source, args.work / 'pytorch_model')
    (args.work / 'results/MODEL_SNAPSHOT.json').write_text(json.dumps(receipt, indent=2) + '\n')
    # Construction and save on CPU catch fixture errors before any training arm.
    pipeline = StableDiffusionPipeline.from_pretrained(receipt['path'], local_files_only=True)
    destination = args.work / 'preflight_saved'
    pipeline.save_pretrained(destination)
    receipt['preflight'] = {'status': 'CONSTRUCT_AND_SAVE_PASS', 'device': str(pipeline.device),
                            'saved_files': hashes(destination)}
    assert receipt['preflight']['device'] == 'cpu'
    assert 'model_index.json' in receipt['preflight']['saved_files']
    (args.work / 'results/MODEL_SNAPSHOT.json').write_text(json.dumps(receipt, indent=2) + '\n')
    print(json.dumps(receipt), flush=True)


if __name__ == '__main__':
    main()

Executed file: reproduce_diffusers.py

"""Instrument the actual pinned text-to-image example; run only on cloud CUDA."""
import argparse
import functools
import hashlib
import importlib.metadata
import json
import os
from pathlib import Path
import runpy
import subprocess
import sys
import traceback


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--checkout', type=Path, required=True)
    p.add_argument('--work', type=Path, required=True)
    p.add_argument('--case', choices=['finite', 'overflow', 'gated'], required=True)
    p.add_argument('--model-revision', required=True)
    p.add_argument('--model-path', type=Path, required=True)
    args = p.parse_args()
    args.work.mkdir(parents=True, exist_ok=True)
    result = {'case': args.case, 'status': 'STARTED', 'optimizer_events': [], 'ema_events': [], 'global_steps': []}
    output = args.work / 'result.json'
    try:
        import torch
        from accelerate import Accelerator
        from accelerate.optimizer import AcceleratedOptimizer
        from diffusers.training_utils import EMAModel
        from PIL import Image

        if not torch.cuda.is_available():
            raise RuntimeError('CUDA required: this is not a CPU substitute for the fp16 example')
        result['environment'] = {
            'gpu': torch.cuda.get_device_name(), 'capability': list(torch.cuda.get_device_capability()),
            'python': sys.version, 'cuda': torch.version.cuda,
            'packages': {n: importlib.metadata.version(n) for n in ['torch', 'diffusers', 'accelerate', 'transformers', 'datasets']},
            'model': 'hf-internal-testing/tiny-stable-diffusion-pipe', 'model_revision': args.model_revision,
            'model_path': str(args.model_path.resolve()),
        }
        assert args.model_path.is_dir() and (args.model_path/'model_index.json').is_file()
        target = args.checkout / 'examples/text_to_image/train_text_to_image.py'
        assert subprocess.check_output(['git', '-C', str(args.checkout), 'rev-parse', 'HEAD'], text=True).strip() == '48cdbaef5ee3d718aaee392b44f2099ae511c4a5'
        original_source = target.read_text()
        result['upstream_source_sha256'] = hashlib.sha256(target.read_bytes()).hexdigest()
        assert result['upstream_source_sha256'] == 'aa6c87e80c859c1448896d678c9bda48fe36f44cc95f86ef995b37f468acf34e'
        assert not subprocess.check_output(['git', '-C', str(args.checkout), 'status', '--porcelain', '--untracked-files=no'], text=True).strip()
        result['source_pin'] = '48cdbaef5ee3d718aaee392b44f2099ae511c4a5'
        if args.case == 'gated':
            old = '# Checks if the accelerator has performed an optimization step behind the scenes\n            if accelerator.sync_gradients:'
            new = '# Checks if the accelerator has performed an optimization step behind the scenes\n            if accelerator.sync_gradients and not accelerator.optimizer_step_was_skipped:'
            assert original_source.count(old) == 1, 'The exact intended source site changed'
            target = args.work / 'gated_train_text_to_image.py'
            target.write_text(original_source.replace(old, new))
        result['executed_source_sha256'] = hashlib.sha256(target.read_bytes()).hexdigest()

        data = args.work / 'images'
        data.mkdir(exist_ok=True)
        metadata = []
        for i in range(8):
            pixels = bytes((x * 7 + y * 11 + i * 13 + c * 29) % 256 for y in range(64) for x in range(64) for c in range(3))
            Image.frombytes('RGB', (64, 64), pixels).save(data / f'{i}.png')
            metadata.append(json.dumps({'file_name': f'{i}.png', 'text': 'a small patterned image'}))
        (data / 'metadata.jsonl').write_text('\n'.join(metadata) + '\n')

        state = {'backwards': 0, 'model': None, 'accelerator': None, 'raw_steps': 0, 'optimizer_event': None}
        def snapshot(parameters):
            return [x.detach().clone() for x in parameters]
        def equal(a, b):
            return len(a) == len(b) and all(torch.equal(x, y) for x, y in zip(a, b))
        old_init, old_prepare = Accelerator.__init__, Accelerator.prepare
        old_backward, old_log = Accelerator.backward, Accelerator.log
        old_step, old_ema, old_raw_step = AcceleratedOptimizer.step, EMAModel.step, torch.optim.AdamW.step

        @functools.wraps(old_init)
        def init(self, *a, **kw):
            old_init(self, *a, **kw)
            assert self.device.type == 'cuda' and self.mixed_precision == 'fp16'
            self.scaler = torch.amp.GradScaler('cuda', init_scale=16., growth_interval=1000000)
            state['accelerator'] = self

        @functools.wraps(old_prepare)
        def prepare(self, *a, **kw):
            for obj in a:
                if isinstance(obj, torch.nn.Module) and any(x.requires_grad for x in obj.parameters()):
                    state['model'] = obj
            return old_prepare(self, *a, **kw)

        @functools.wraps(old_backward)
        def backward(self, loss, *a, **kw):
            old_backward(self, loss, *a, **kw)
            state['backwards'] += 1
            if args.case != 'finite' and state['backwards'] == 4:
                grad = next(x.grad for x in state['model'].parameters() if x.grad is not None)
                grad[(0,) * grad.ndim] = float('inf')

        @functools.wraps(old_raw_step)
        def raw_step(self, *a, **kw):
            state['raw_steps'] += 1
            return old_raw_step(self, *a, **kw)

        @functools.wraps(old_step)
        def step(self, *a, **kw):
            params = list(state['model'].parameters())
            before = snapshot(params)
            count = state['raw_steps']
            scale = state['accelerator'].scaler.get_scale()
            value = old_step(self, *a, **kw)
            event = {'attempt': len(result['optimizer_events']) + 1, 'raw_step_called': state['raw_steps'] > count,
                     'skipped': bool(self.step_was_skipped), 'weights_unchanged': equal(before, params),
                     'scale_before': scale, 'scale_after': state['accelerator'].scaler.get_scale()}
            result['optimizer_events'].append(event)
            state['optimizer_event'] = event
            return value

        @functools.wraps(old_ema)
        def ema(self, parameters):
            parameters = list(parameters)
            before = snapshot(self.shadow_params)
            differs = not equal(before, parameters)
            count = self.optimization_step
            old_ema(self, parameters)
            event = state['optimizer_event']
            result['ema_events'].append({'attempt': event['attempt'], 'optimizer_skipped': event['skipped'],
                'ema_differed_from_model_before': differs, 'ema_changed': not equal(before, self.shadow_params),
                'counter_before': count, 'counter_after': self.optimization_step})

        @functools.wraps(old_log)
        def log(self, values, *a, **kw):
            if 'train_loss' in values:
                result['global_steps'].append({'attempt': state['optimizer_event']['attempt'], 'global_step': kw.get('step')})
            return old_log(self, values, *a, **kw)

        Accelerator.__init__, Accelerator.prepare = init, prepare
        Accelerator.backward, Accelerator.log = backward, log
        AcceleratedOptimizer.step, EMAModel.step, torch.optim.AdamW.step = step, ema, raw_step
        ns = runpy.run_path(str(target), run_name='upstream_repro_target')
        globals_ = ns['main'].__globals__
        original_parse = globals_['parse_args']
        def parse():
            parsed = original_parse()
            parsed.report_to = []  # Avoid external loggers; optimization and EMA paths are unchanged.
            return parsed
        globals_['parse_args'] = parse
        sys.argv = [str(target), '--pretrained_model_name_or_path=' + str(args.model_path.resolve()),
                    '--train_data_dir=' + str(data), '--resolution=64',
                    '--center_crop', '--train_batch_size=1', '--gradient_accumulation_steps=1', '--max_train_steps=6',
                    '--learning_rate=0.001', '--lr_scheduler=constant', '--lr_warmup_steps=0', '--use_ema',
                    '--mixed_precision=fp16', '--seed=13', '--dataloader_num_workers=0',
                    '--checkpointing_steps=100000', '--output_dir=' + str(args.work / 'trained')]
        result['training_argv'] = list(sys.argv)
        ns['main']()
        trained = args.work / 'trained'
        result['pipeline_export'] = {str(f.relative_to(trained)): {
            'sha256': hashlib.sha256(f.read_bytes()).hexdigest(), 'bytes': f.stat().st_size}
            for f in trained.rglob('*') if f.is_file()}
        assert 'model_index.json' in result['pipeline_export']
        for component in ['unet', 'vae', 'text_encoder', 'safety_checker']:
            assert any(name.startswith(component + '/') and name.endswith(('.bin', '.safetensors'))
                       for name in result['pipeline_export']), component
        result['final_pipeline_save'] = 'PASS'
        skipped = [e for e in result['optimizer_events'] if e['skipped']]
        if args.case == 'finite':
            assert not skipped and state['raw_steps'] == 6
            result['status'] = 'FINITE_CONTROL_PASS'
        else:
            assert len(skipped) == 1 and skipped[0]['attempt'] == 4
            assert skipped[0]['weights_unchanged'] and not skipped[0]['raw_step_called']
            ema_on_skip = [e for e in result['ema_events'] if e['optimizer_skipped']]
            if args.case == 'overflow':
                assert len(ema_on_skip) == 1 and ema_on_skip[0]['ema_differed_from_model_before']
                assert ema_on_skip[0]['ema_changed']
                assert ema_on_skip[0]['counter_after'] == ema_on_skip[0]['counter_before'] + 1
                assert state['raw_steps'] == 5 and result['global_steps'][-1]['global_step'] == 6
                result['status'] = 'CONFIRMED_ACTUAL_EXAMPLE_EMA_ON_SKIPPED_UPDATE'
            else:
                assert not ema_on_skip and state['raw_steps'] == 6
                assert len(result['optimizer_events']) == 7 and result['global_steps'][-1]['global_step'] == 6
                result['status'] = 'GATED_CONTROL_PASS'
    except Exception:
        result['status'] = 'ERROR_OR_INCONCLUSIVE'
        result['traceback'] = traceback.format_exc()
    output.write_text(json.dumps(result, indent=2) + '\n')
    print(json.dumps(result, indent=2), flush=True)
    return 0 if result['status'] != 'ERROR_OR_INCONCLUSIVE' else 1


if __name__ == '__main__':
    raise SystemExit(main())

Logs

Control	Optimizer attempts	Applied updates	EMA calls	Final `global_step`	Pipeline saved / exit
Finite	6	6	6	6	Yes / 0
Overflow at attempt 4	6	5	6	6	Yes / 0
Same overflow, gated block	7	6	6	6	Yes / 0

At overflow attempt 4, step_was_skipped=True, the raw AdamW step was not called, and model parameters were unchanged. The scale fell from 16 to 8. EMA differed from the model before that call, changed during the call, and advanced its counter from 3 to 4. global_step also advanced from 3 to 4.

flowchart LR
    A["Attempt 4: infinite gradient"] --> B["GradScaler skips optimizer update"]
    B --> C["Current block: EMA changes; global_step 3 → 4"]
    B --> D["Gated block: no EMA call; global_step stays 3"]

The gated control used one additional attempt to reach six applied updates. This is a narrow comparison of the EMA/progress condition, not a claim that all training bookkeeping is fixed.

System Info

Diffusers version: 0.42.0.dev0
Platform: Linux-6.6.122+-x86_64-with-glibc2.39
Running on Google Colab?: Yes
Python version: 3.13.15
PyTorch version (GPU?): 2.11.0+cu130 (True)
Huggingface_hub version: 1.33.0
Transformers version: 5.20.1
Accelerate version: 1.15.0
PEFT version: 0.21.2
torchao version: 0.10.0
Safetensors version: 0.8.0
xFormers version: not installed
Accelerator: Tesla T4, 15360 MiB
Using GPU in script?: Yes, Tesla T4; fp16
Using distributed or parallel set-up in script?: No, one process / one GPU
Source commit: 48cdbaef5ee3d718aaee392b44f2099ae511c4a5
Datasets version: 4.8.5

Who can help?

I’d fix this by guarding EMA, global_step, progress, training-loss logging, and checkpoints with accelerator.sync_gradients and not accelerator.optimizer_step_was_skipped. Accumulated loss would still reset at every sync boundary.

Before opening a PR, two scope questions:

Should this also address the fixed epoch budget, which can finish below --max_train_steps applied updates after skips?
Similar blocks appear in SDXL, instruct-pix2pix, and other examples. Should those be included, or should this start with train_text_to_image.py?

Navigation Menu

Repository navigation

Text-to-image example updates EMA when AMP skips the optimizer step #15112

Describe the bug

Reproduction

Logs

System Info

Who can help?

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Text-to-image example updates EMA when AMP skips the optimizer step #15112

Description

Describe the bug

Reproduction

Logs

System Info

Who can help?

Activity

SeverinVisionary commented on Oct 23, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions