# Diagnose and fix NaN/inf loss with PyTorch mixed precision (AMP)
Use this when training with `torch.amp` produces NaN or inf loss (or NaN
gradients) and plain FP32 training is clean. AMP bugs have a specific taxonomy:
the fix depends on whether the NaN originates in the forward pass (an fp16
overflow in your math) or in the scaled gradients (a loss-scale problem), and
on whether the AMP recipe itself is misordered. Classify first, using FP32 as
the control.
## 1. Establish the control: is AMP actually the cause?
Run the same training in pure FP32 (drop `autocast` and the `GradScaler`):
- **NaN in FP32 too:** this is not an AMP bug. Check the data for inf/NaN
(`torch.isfinite(batch).all()`), the learning rate, and loss stability
before touching anything below.
- **Clean in FP32, NaN with AMP:** continue with this procedure.
Also note WHEN the NaN appears. NaN at iteration 0–few → forward-pass fp16
overflow or a broken recipe (sections 2–3). NaN after many clean steps → a
learning-rate schedule spike or an over-aggressive loss scale later in training
(section 4). Sudden NaN right after an LR warmup ends is a schedule problem,
not a scaler problem.
## 2. Verify the AMP recipe is ordered correctly
The canonical loop (PyTorch ≥ 2.1 device-agnostic API; `torch.cuda.amp` is the
older equivalent) is:
```python
import torch
scaler = torch.amp.GradScaler("cuda")
for data, target in loader:
data, target = data.cuda(), target.cuda()
optimizer.zero_grad()
# autocast must wrap BOTH the forward pass AND the loss computation
with torch.amp.autocast("cuda", dtype=torch.float16):
output = model(data)
loss = criterion(output, target)
scaler.scale(loss).backward()
# unscale BEFORE any gradient inspection or clipping
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
```
Check these recipe violations — each one is a known NaN source:
1. **Loss computed outside `autocast`.** Reductions, softmax, and log/exp in the
loss are the most overflow-sensitive ops in the network. If `criterion` runs
in FP32 outside the context while the model runs in fp16, you get the worst
of both. Wrap both together as above.
2. **Clipping before `unscale_`.** `clip_grad_norm_` on still-scaled gradients
compares against a threshold that is off by the current scale factor (default
`init_scale=65536`, i.e. 2**16), which effectively disables or distorts
clipping. Always call `scaler.unscale_(optimizer)` first.
3. **Calling `optimizer.step()` directly.** Any step taken outside
`scaler.step(optimizer)` applies unscaled-overflowed or double-scaled
gradients. The only step call should be `scaler.step(optimizer)` followed by
`scaler.update()`.
4. **Manually scaling the loss before `scaler.scale()`.** Do not multiply the
loss yourself for accumulation and then also call `scaler.scale(loss)` — the
scaler multiplies by the dynamic scale on top of whatever you did. If you
divide the loss by accumulation steps, do it inside the autocast block and
let the scaler handle scaling.
## 3. Classify forward overflow vs gradient overflow
**Forward overflow** — the printed `loss` itself is NaN/inf. Something in fp16
math exceeded the fp16 max (~65504) or divided by ~zero: large logits through
softmax/exp, `log` of a tiny probability, sums over long sequences, or custom
layers with exp/pow. Fixes, in order:
1. Keep the fragile op in FP32 inside the autocast region with an explicit
cast, e.g. compute softmax/log-softmax in `float32` and cast back. Autocast
already keeps reductions and softmax in FP32 for built-in ops; custom code
does not get that protection automatically.
2. Check the input data range — fp16 overflows on values FP32 training never
noticed (e.g. targets or features with magnitude > 1e4).
3. If the model is an LLM/transformer and the platform allows it, switch to
bfloat16 instead (section 5).
**Gradient overflow** — the loss prints fine but training stalls or diverges
with the scaler misbehaving. `GradScaler` skips the optimizer step whenever the
scaled gradients contain inf/NaN, so weights stay clean but no learning happens.
Diagnose with the scale value:
```python
print(f"scale={scaler.get_scale():.0f}")
```
- Scale collapses toward its minimum (1.0) and stays there with most steps
skipped → persistent overflow. The gradients are too large for fp16 at this
scale.
- Scale oscillates at a healthy high value with occasional skipped steps → this
is normal dynamic scaling doing its job; the NaN you saw may be a one-off
from a bad batch. Log the scale over time before acting.
Gradient-overflow fixes:
1. Lower the starting scale: `torch.amp.GradScaler("cuda", init_scale=2**10)`.
The default 2**16 assumes well-behaved gradients; models with naturally
large gradients (some detection/segmentation heads) overflow immediately.
2. Make growth less aggressive: `growth_interval=2000` is the default number of
consecutive clean steps before the scale doubles. If overflow returns right
after each growth step, raise `growth_interval` or lower
`growth_factor=2.0`.
3. Add/strengthen gradient clipping — but only after `unscale_`, per section 2.
## 4. Check the learning rate before blaming the scaler
A large fraction of "AMP NaN" reports are LR problems that fp16 merely
amplifies: the loss scale rides the edge of overflow, and an LR spike pushes
activations over it. If NaN coincides with the end of LR warmup, a schedule
restart, or an LR increase, reduce the peak LR (or extend warmup) and re-test
before tuning scaler knobs. AMP does not fix an unstable optimization; it only
changes the arithmetic.
## 5. Consider bfloat16 instead of float16
On Ampere and newer NVIDIA GPUs (A100, H100, RTX 30-series and later),
`bfloat16` autocast needs no loss scaling at all, because bf16 has the same
dynamic range as FP32 (it trades mantissa precision for range):
```python
with torch.amp.autocast("cuda", dtype=torch.bfloat16):
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
```
No `GradScaler` — and a whole class of overflow NaNs disappears. The trade is
reduced mantissa precision, which transformers and most modern architectures
tolerate well; for models that need fp16's precision but not its range issues,
stay with fp16 + scaler. On pre-Ampere hardware bf16 is not accelerated, so
this option does not apply.
## 6. Checklist for the NaN
1. Control run: pure FP32 clean → AMP bug (continue); FP32 NaN too → data/LR
problem, stop here.
2. Verify recipe: loss inside autocast, `unscale_` before clipping,
`scaler.step`/`scaler.update()` only, no manual pre-scaling.
3. Classify: NaN in printed loss → forward fp16 overflow (cast fragile ops to
FP32, check input magnitudes); clean loss but stalled training → read
`scaler.get_scale()`.
4. Persistent gradient overflow: lower `init_scale`, raise `growth_interval`,
clip after unscale.
5. Check LR schedule alignment with the NaN onset before tuning the scaler.
6. On Ampere+: try `dtype=torch.bfloat16` and drop the scaler entirely.