The loss turns NaN in the middle of training. What are the causes, and the fix for each?
A run trains normally for about two epochs, then the loss prints nan and never recovers. The same data trains fine in fp32 with a smaller learning rate.
Constraints: keep mixed precision on, keep the architecture, and make the failure detectable at the step it occurs.
scaler = torch.amp.GradScaler()
for x, y in loader:
opt.zero_grad()
with torch.autocast("cuda", dtype=torch.float16):
logits = model(x)
loss = -(y * torch.log(torch.softmax(logits, dim=-1))).sum(-1).mean()
scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()Find every reason the loss can turn NaN here and fix them.
Exploding gradients — clip by global norm. An fp16 overflow — use a loss scaler or bf16. A hand-written log(0) — use the fused logits loss or add an epsilon. A rate that is too high — lower it or add warmup. NaN already in the batch — assert on inputs. Find the first bad step and log gradient norms.
- ✗Blaming the data and never checking gradient norms or the loss implementation
- ✗Writing log(softmax(x)) by hand instead of using the fused logits loss
- ✗Assuming a GradScaler makes fp16 overflow impossible
- →Why does bf16 tolerate overflow better than fp16 at the same bit width?
- →How would you catch the exact step and tensor where the first NaN appears?
There are three separate causes here, and GradScaler only covers one of them.
1. The hand-written log(softmax(x)). On a confident prediction softmax underflows to exactly 0 in fp16, so log(0) is -inf, and -inf * 0 then gives nan. The fused cross_entropy computes the same quantity through a stable log_softmax and never takes the log of zero.
2. No gradient clipping. Before scaler.step the gradients must be unscaled (scaler.unscale_) and clipped by global norm — otherwise a rare outlier flies off to inf.
3. No assertion at the point of failure. The break must be caught at the step where it happens, not an epoch later.
scaler = torch.amp.GradScaler()
for step, (x, y) in enumerate(loader):
assert torch.isfinite(x).all(), f"step {step}: NaN/inf in inputs"
opt.zero_grad()
with torch.autocast("cuda", dtype=torch.float16):
logits = model(x)
loss = torch.nn.functional.cross_entropy(logits, y) # fused, stable
assert torch.isfinite(loss), f"step {step}: loss {loss.item()}"
scaler.scale(loss).backward()
scaler.unscale_(opt) # unscale BEFORE clipping
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(opt)
scaler.update()If nan survives all of that, switch dtype to torch.bfloat16: its exponent range matches fp32, so overflow all but disappears and the loss scaler becomes unnecessary.
