Trainer counts fp16 steps skipped by GradScaler as completed update steps (`global_step`, `max_steps`, `on_optimizer_step`)

Author: SeverinVisionaryCreated Oct 1, 2026Updated Oct 2, 2026

System Info

transformers 5.17.0 and main @ 96331a9; accelerate 1.15.0 and main @ 421af8b; torch 2.14.0; Python 3.13; Tesla T4, CUDA 13.0, fp16=True. Also reproduces on CPU with a GradScaler installed on the Accelerator (on CPU, fp16=True creates no scaler).

Who can help?

@SunMarc (trainer)

Reproduction

With fp16=True, GradScaler skips optimizer.step() when the gradients contain inf/NaN. Trainer already knows about the skip: it checks self.accelerator.optimizer_step_was_skipped before lr_scheduler.step(). Everything after that check still treats the step as done (trainer.py L1896-L1912 on main):

self.state.global_step += 1 runs unconditionally;
on_optimizer_step and on_step_end fire, and logging/saving/evaluation scheduled on global_step run for a step that changed nothing;
max_steps is compared against global_step, so every skipped step shortens training by one real update.

TrainerState.global_step is documented as "the number of update steps completed".

repro_transformers.py below injects one inf gradient on the 3rd step, with max_steps=6 and optim="adamw_torch":

python

"""HF Trainer fp16: a GradScaler-skipped update is counted in state.global_step, fires on_optimizer_step /
on_step_end, and counts toward max_steps. On CUDA run as-is; on CPU a GradScaler is installed on the
Accelerator to emulate the CUDA fp16 path (CPU fp16 otherwise creates no scaler)."""
import tempfile, torch
from transformers import Trainer, TrainingArguments, TrainerCallback

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__(); torch.manual_seed(0)
        self.net = torch.nn.Sequential(torch.nn.Linear(4, 8), torch.nn.Tanh(), torch.nn.Linear(8, 1))
        calls = {"n": 0}
        def poison(g):  # inf in one gradient element on the 3rd backward
            calls["n"] += 1
            if calls["n"] == 3:
                g = g.clone(); g.view(-1)[0] = float("inf")
            return g
        self.net[0].weight.register_hook(poison)
    def forward(self, x, y):
        return {"loss": torch.nn.functional.mse_loss(self.net(x).float(), y)}

class DS(torch.utils.data.Dataset):
    def __init__(self):
        g = torch.Generator().manual_seed(0); self.x = torch.randn(12, 4, generator=g); self.y = self.x.sum(1, keepdim=True)
    def __len__(self): return 12
    def __getitem__(self, i): return {"x": self.x[i], "y": self.y[i]}

class Rec(TrainerCallback):
    def __init__(self): self.opt_step_events = 0
    def on_optimizer_step(self, args, state, control, **kw): self.opt_step_events += 1
    def on_step_end(self, args, state, control, **kw):
        opt = kw["optimizer"]; opt = getattr(opt, "optimizer", opt)
        adam_step = max((float(s["step"]) for s in opt.state.values() if "step" in s), default=0.0)
        print(f"on_step_end: global_step={state.global_step} optimizer_step_count={adam_step:.0f} "
              f"lr_scheduler.last_epoch={kw['lr_scheduler'].last_epoch} on_optimizer_step_calls={self.opt_step_events}")

cuda = torch.cuda.is_available()
args = TrainingArguments(output_dir=tempfile.mkdtemp(), per_device_train_batch_size=2, max_steps=6, learning_rate=1e-2,
                         fp16=True, use_cpu=not cuda, optim="adamw_torch", report_to=[], save_strategy="no",
                         logging_steps=1, disable_tqdm=True, remove_unused_columns=False)
rec = Rec()
tr = Trainer(model=Net(), args=args, train_dataset=DS(), callbacks=[rec],
             data_collator=lambda b: {k: torch.stack([e[k] for e in b]) for k in ("x", "y")})
# small init_scale so that only the poisoned step overflows (the default 65536 works too, with more skips)
if cuda:
    tr.accelerator.scaler = torch.amp.GradScaler("cuda", init_scale=16.0)
else:
    tr.accelerator.scaler = torch.amp.GradScaler("cpu", init_scale=16.0); tr.accelerator.native_amp = True
tr.train()
print(f"final: state.global_step={tr.state.global_step} (max_steps=6), updates actually applied="
      f"{max(float(s['step']) for s in tr.optimizer.optimizer.state.values()):.0f}")

on_step_end: global_step=2 optimizer_step_count=2 lr_scheduler.last_epoch=2 on_optimizer_step_calls=2
on_step_end: global_step=3 optimizer_step_count=2 lr_scheduler.last_epoch=2 on_optimizer_step_calls=3
on_step_end: global_step=4 optimizer_step_count=3 lr_scheduler.last_epoch=3 on_optimizer_step_calls=4
...
final: state.global_step=6 (max_steps=6), updates actually applied=5

The optimizer and the scheduler agree (5 updates, last_epoch=5), but global_step says 6 and training stops one update early. A linear schedule therefore never reaches its final value. Checkpoints written with save_steps record the inflated global_step, and on resume global_step also decides how many batches to skip. With the default optim (adamw_torch_fused), the scheduler advances on the skipped step as well. That part is huggingface/accelerate#4242 and is not what this issue is about.

The same happens without an injected fault. With the default scaler (init_scale=65536) and no injection, on the same T4, steps 2 and 3 overflow on their own and the run ends with state.global_step=6 (max_steps=6), updates actually applied=4.

Expected behavior

A step skipped by the scaler is not counted as a completed update step. global_step stays put, max_steps counts applied updates, and on_optimizer_step does not fire (or it reports that the step was skipped). If counting skipped steps is intended, the TrainerState docstring and the max_steps docs should say so, and a separate "applied updates" counter would let callbacks tell the two cases apart.

Source: huggingface/transformers

View original on GitHub View discussion on GitHub