Diagnose and fix NaN/inf loss with PyTorch mixed precision (AMP)

[{"constraint": ">=2.1 for the torch.amp device-agnostic API; torch.cuda.amp recipe identical in 1.6+", "technology": "PyTorch", "version_scheme": "semver"}]

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:

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:

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):

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.

Built from the PyTorch AMP documentation (autocast, GradScaler defaults initscale=2**16/growthinterval=2000, the documented unscale-before-clipping pattern) and the NVIDIA mixed-precision user guide's dynamic loss scaling FAQ, plus recurring NaN signatures from discuss.pytorch.org.