# Diagnose and fix NaN or stalled loss with fastai mixed precision (`learn.to_fp16()`)
Use this when training with `learn.to_fp16()` produces NaN loss, diverges,
or runs but plateaus at a clearly worse loss than the same model in fp32.
Classify the failure mode first — overflow, underflow, and silent step
skipping each have a different fix.
## 0. How fastai mixed precision actually works
```python
learn = learn.to_fp16() # returns the learner with the MixedPrecision callback added
learn.fit_one_cycle(5, 3e-3)
learn = learn.to_fp32() # removes the callback; back to full precision
```
`to_fp16(**kwargs)` adds the `MixedPrecision` callback, which on current
fastai uses PyTorch's native AMP:
- Each batch's forward pass runs under `torch.amp.autocast` in fp16.
- **Predictions are cast back to fp32 before the loss is computed**
(`after_pred: self.learn.pred = to_float(self.pred)`), so the loss
itself is computed in fp32 — a NaN *loss* means fp16 produced Inf/NaN
*activations* upstream of that cast, or the loss overflows after it.
- The loss is multiplied by a scale factor (`torch.amp.GradScaler`) before
backward, so tiny gradients survive fp16's limited range. Keyword
arguments to `to_fp16()` are passed straight to `GradScaler`:
`init_scale` (default `2**16`), `growth_factor`, `backoff_factor`,
`growth_interval`.
- **On gradient overflow the optimizer step is silently skipped**
(`CancelStepException`) and the scale is halved; after enough clean steps
it doubles again (dynamic loss scaling). Skipped steps print nothing —
training just learns less per epoch than you think.
- `learn.to_bf16()` switches to bfloat16 instead. BF16 has the same
exponent range as fp32, so **no GradScaler is used** and overflow is
far less likely — but it raises
`ValueError("Unsupported GPU for bfloat16 mixed precision training")`
unless `torch.cuda.is_bf16_supported()` (Ampere and newer).
## 1. Classify the failure from observable evidence
### A. Loss goes NaN within the first epoch
The dynamic scaler cannot cause this: it only scales gradients *after* a
successful forward/backward, and it halves the scale on overflow. NaN this
early means **forward-pass overflow in fp16** (fp16 tops out at 65504) or
poisoned inputs. Diagnose:
1. Confirm it's fp16-specific, not a data/model bug: `learn.to_fp32()`,
then `learn.fit_one_cycle(1, same_lr)`. If fp32 also NaNs, the model
or data is broken — fix that first (unnormalized inputs are the usual
culprit).
2. If fp32 is healthy, find the overflowing op. The usual suspects, in
order: unnormalized inputs (raw 0–255 pixels, unscaled tabular data);
a custom loss or metric that overflows in fp16; large logits from an
unscaled final layer; `exp`/`pow` in custom code. Note the loss is
computed in fp32, so a custom loss that internally casts its inputs to
half will still overflow.
3. Fixes in order: normalize/scale the inputs; fix the overflowing op;
only then consider `learn.to_fp16(init_scale=2.**8)` to start the
dynamic scaler lower (this helps when the *scaled loss itself*
overflows, not when the forward pass does).
### B. Loss is fine for a while, then goes NaN mid-training
- Check `learn.scales` — the `MixedPrecision` callback appends
`scaler.get_scale()` after every optimizer step, so this list is the
history of the dynamic loss scale. A scale that collapses (repeated
halving down to tiny values) right before the NaN means chronic
gradient overflow the backoff couldn't outrun: some layer's gradients
keep overflowing in fp16. Common with very deep nets at high lr or
with gradient spikes (e.g. RNNs/transformers without clipping).
- Fixes: lower the learning rate; add gradient clipping; or accept that
this architecture is not fp16-safe and switch to `learn.to_bf16()` on
a supported GPU (no scaler, no overflow cliff) or back to fp32.
### C. Training runs to completion but loss plateaus above the fp32 baseline
- Likely **gradient underflow**: small gradients flush to zero in fp16
faster than the scaler can compensate, or so many steps are being
skipped on overflow that effective training is a fraction of what you
think. Again, `learn.scales` is the diagnostic: a scale pinned at its
minimum while loss is flat = underflow/backoff loop.
- Sanity check the comparison properly: same seed, same `fit_one_cycle`
schedule, fp32 via `learn.to_fp32()`. A small gap (a few percent) can be
normal fp16 noise; a large gap is a real problem.
- Fix: this model genuinely needs fp32's range. Prefer `to_bf16()`
(same range as fp32, still half the memory) over fighting the scaler.
### D. `to_fp16()` errors immediately
- `ValueError: Unsupported GPU for bfloat16 mixed precision training` —
you called `to_bf16()` on a pre-Ampere GPU. Use `to_fp16()` or fp32.
- AMP requires a CUDA build of PyTorch; on CPU-only machines mixed
precision is a no-op at best — don't use it there.
## 2. The reliable comparison protocol
Don't eyeball two different runs. This is the minimum honest A/B:
```python
learn_fp32 = learn.to_fp32()
learn_fp32.fit_one_cycle(3, 3e-3)
loss_fp32 = learn_fp32.recorder.values[-1]
learn_fp16 = learn.to_fp16() # same model, same schedule
learn_fp16.fit_one_cycle(3, 3e-3)
loss_fp16 = learn_fp16.recorder.values[-1]
print(learn_fp16.scales[-10:]) # loss-scale history: collapsing = overflow trouble
```
`learn.scales` only exists while the `MixedPrecision` callback is attached
(it's torn down in `after_fit`), so read it right after `fit`, before
`to_fp32()`.
## 3. Checklist
1. fp32 baseline first: `learn.to_fp32()` + one epoch tells you whether
the NaN is fp16-specific.
2. Read `learn.scales` after the run — collapsing scale = overflow,
pinned-minimum scale + flat loss = underflow.
3. Fix the data/op (normalization, custom loss internals) before touching
the scaler; `init_scale` is the last knob, not the first.
4. If the model needs fp32's dynamic range, `to_bf16()` on Ampere+ beats
fighting fp16 overflow — same memory savings, no GradScaler cliff.