# Diagnose and fix NaN or stalled loss with fastai mixed precision (learn.to_fp16())

Diagnose fastai mixed-precision training failures: classify early-NaN (fp16 forward overflow) vs mid-training NaN (chronic gradient overflow) vs plateau above the fp32 baseline (gradient underflow / skipped steps) using the learn.scales loss-scale history, then fix the data/op, tune GradScaler init_scale, or switch to to_bf16()/to_fp32().

Exact reference: {"kind":"skill_version","skill_id":"skl_M2QcMkDSFRHkLThOAlHqzQ","version_id":"skv_Fj-KLtbVik8bhQDE5oGxHw"}

Applicability: [{"constraint":">=2.0","technology":"fastai","version_scheme":"semver"},{"constraint":">=2.0","technology":"PyTorch","version_scheme":"semver"}]

# 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.


## Supporting basis and limitations

Built from the fastai v2 source: MixedPrecision callback and Learner.to_fp16/to_bf16/to_fp32 in fastai/callback/fp16.py (torch.amp autocast + GradScaler, after_pred to_float, overflow skip via CancelStepException, learn.scales history, bf16 unsupported-GPU ValueError), and the mixed-precision explainer in nbs/18_callback.fp16.ipynb (the three fp16 problems: imprecise updates, gradient underflow, activation/loss overflow).

## Change and rationale

New skill: diagnose and fix NaN or stalled loss with fastai mixed precision (learn.to_fp16()).

Mixed-precision NaN/divergence questions recur on the fastai forums and most advice jumps straight to tweaking the loss scale, when early NaN is actually forward-pass overflow in fp16 (which the dynamic scaler cannot prevent) and plateaus are gradient underflow. This skill adds a decision procedure keyed on when the NaN appears plus the learn.scales diagnostic, and documents the real v2 mechanics (fp32 loss via to_float preds, silent CancelStepException skips, GradScaler kwargs, bf16 mode).