Loss becomes NaN after switching to fp16 mixed precision

A ranking model trained in fp32 took 11 hours per run. We enabled automatic mixed precision with fp16 to speed it up. The first run printed:

step 250  loss 0.4821
step 275  loss 0.4790
step 300  loss nan
step 325  loss nan

and every metric after that was NaN.

Reproducing it

We enabled anomaly detection for a short debugging run:

torch.autograd.set_detect_anomaly(True)

RuntimeError: Function 'SoftmaxBackward0' returned nan values in its 0th output.

Logging activation statistics before the softmax showed attention scores reaching 71,000 in a few heads. fp16's largest finite value is 65,504.

What we ruled out

• Bad data. The same batches trained fine in fp32, and we checked inputs for NaN and infinities.
• Learning rate too high. The fp32 run with the same learning rate was stable.

Root cause

Three problems combined:

1. A custom attention block computed q @ k.transpose(-1, -2) and applied its own scaling after the matmul. Under autocast, the matmul ran in fp16, and before scaling some scores exceeded fp16's maximum, becoming inf. Softmax of inf produces NaN.
2. The loss was computed on fp16 logits with a hand-written log-sum-exp that also overflowed for large logits.
3. We used autocast without GradScaler. Small gradients underflowed to zero in fp16, and when overflow occurred there was no mechanism to skip the step, so NaN parameters propagated.

What is happening under the hood

fp16 has a narrow range: large values overflow to infinity and small gradients underflow to zero. Autocast runs some ops in fp16 (matmuls, convolutions) and keeps others in fp32 (most reductions, softmax, losses), but only for operations it knows; custom code between them runs in whatever dtype its inputs have.

GradScaler multiplies the loss by a large factor so gradients stay representable, unscales them before the optimizer step, skips steps where infinities or NaNs appear, and adjusts the factor over time.

bf16 keeps fp32's exponent range, so overflow and underflow are rarely a problem, and loss scaling is not needed. It has less precision, which matters little for most training.

Environment

Minimal reproduction

x = torch.full((4,), 300.0, device="cuda", dtype=torch.float16)
print((x @ x))                        # 360000 > 65504 -> inf
print(torch.softmax(torch.tensor([float("inf"), 1.0]), 0))   # nan

The fix

1. Use bf16 on hardware that supports it. Our A10G GPUs do:

use_bf16 = torch.cuda.is_bf16_supported()
amp_dtype = torch.bfloat16 if use_bf16 else torch.float16
scaler = torch.amp.GradScaler("cuda", enabled=not use_bf16)

for batch in loader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast("cuda", dtype=amp_dtype):
        logits = model(batch.inputs)
    loss = loss_fn(logits.float(), batch.targets)        # loss in fp32
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)
    scaler.update()

With enabled=False, the scaler calls are pass-throughs, so the same loop works for bf16 and fp16.

2. Scale before the matmul and use the fused attention. We replaced the custom block with torch.nn.functional.scaled_dot_product_attention, which applies scaling inside a numerically careful kernel. Where custom code remained, we scale queries before the matmul, not after.

3. Use library losses. F.cross_entropy on fp32 logits replaced the hand-written log-sum-exp.

4. A NaN guard. Training stops with a clear error and saves the batch and state if the loss is not finite, instead of running for hours producing NaN:

if not torch.isfinite(loss):
    save_debug_snapshot(batch, model, step)
    raise FloatingPointError(f"non-finite loss at step {step}")

Alternatives and trade-offs

• fp16 with GradScaler only. Works on GPUs without bf16 support; custom ops still need care.
• fp32 for the whole model. Stable and slow; we were at 11 hours per run.
• TF32 matmuls (on by default for some operations on Ampere and newer). A smaller speedup with almost no stability risk.

How to test the fix

Five runs of 20,000 steps in bf16 finished without NaN, with validation metrics within noise of fp32 runs. A unit test runs the attention block in fp16 with large inputs and asserts finite outputs.

Confirming in production

• Training time dropped from 11 hours to 6 hours 40 minutes in bf16.
• Validation NDCG matched the fp32 baseline within 0.1 points across three seeds.
• The NaN guard has not triggered in production training runs.

Regression guard

• The fp16 overflow unit test for custom layers.

• Training dashboards plot loss scale (when fp16 is used), gradient norm and fraction of skipped steps.

• Code review rule: custom numerics must state the dtype they run in.

• New model architectures run a short mixed precision soak test (2,000 steps) before long training runs are scheduled.

Security implications

None. Debug snapshots saved by the NaN guard contain training batches, so they are written to the same restricted bucket as training data and expire after 14 days.

Prevention checklist

• Prefer bf16 on GPUs that support it.
• Use GradScaler with fp16.
• Compute losses and reductions in fp32.
• Use fused, library implementations for attention and losses.
• Stop training on the first non-finite loss and save context.
• Log activation and gradient statistics for custom layers during the first runs of a new architecture.

References

• PyTorch automatic mixed precision — https://pytorch.org/docs/stable/amp.html
• AMP examples with GradScaler — https://pytorch.org/docs/stable/notes/amp_examples.html
• scaled_dot_product_attention — https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
• NVIDIA mixed precision training guide — https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html