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

TL;DR

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. Classify the failure mode first, overflow, underflow, and silent step skipping each have a different fix.

When to use

You are seeing this: Keyword arguments to tofp16() are passed straight to GradScaler: initscale (default 216), growthfactor, backofffactor, growthinterval. - On gradient overflow the optimizer step is silently skipped (CancelStepException) and the scale is halved; after... Use this skill when you run into "Diagnose and fix NaN or stalled loss with fastai mixed precision (learn.tofp16())".

When not to use

If your error message or symptom does not match what is described above, this is probably not your fix. Search for your exact error text instead of forcing this one to fit.

Versions

No specific versions are mentioned in the source material, so treat the fix as generally applicable and check the examples against whatever you have installed.

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

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:

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.

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