thebench.bio #1 Ecosystem for the Bio-Tech Industry
All Deep Learning questions

Question 20 of 20 Advanced Troubleshooting

Loss becomes NaN during mixed-precision training. How would you debug it?

Short answer

Find the first step and operation that produce an inf or NaN. In float16 the usual causes are overflow beyond 65504, gradients scaled without a GradScaler, numerically unsafe operations such as BCELoss on probabilities, or a learning rate that is too high. Rerunning in float32 or bfloat16 tells you whether precision is the cause.

Why

Mixed precision runs most operations in a 16-bit type for speed. The PyTorch AMP documentation explains that small float16 gradients can flush to zero, which gradient scaling prevents by multiplying the loss before backward, and that float16 has a maximum of 65504, so models pretrained in bfloat16 can overflow instead. The AMP recipe notes that scaler.step skips the optimizer step when gradients contain infs or NaNs, so occasional skips are normal but a loss that is itself NaN is not.

NaNs propagate, so once one appears in the weights every later output is NaN, and the step where it first appeared is the only informative one. Causes split into three groups: data, such as a corrupt image or a division by a zero standard deviation, numerics, such as overflow in float16 or an unstable loss, and optimization, such as a learning rate that throws weights to extreme values.

How to work it through

  1. Log the loss, the gradient scale and per-layer gradient norms every step, and note the first step where anything becomes inf or NaN.
  2. Check the input batch at that step for NaNs, extreme values or division by zero in normalization.
  3. Rerun that step in float32 and in bfloat16 to see whether precision is the cause.
  4. Enable torch.autograd.detect_anomaly briefly to locate the backward operation producing the NaN, accepting that it slows training.
  5. Replace unsafe operations, for example binary_cross_entropy on probabilities with binary_cross_entropy_with_logits, which the AMP documentation recommends.
  6. Confirm gradients are unscaled before clipping, and lower the learning rate or add warmup if the loss spikes before the NaN appears.

An example

A segmentation network trains for 3,000 steps in float16 and then the loss turns NaN. Logs show the scale falling repeatedly and one layer's gradient norm reaching inf two steps earlier. The input batch contains an image with zero variance, and per-image normalization divides by zero.

Adding a small epsilon fixes that, and replacing a sigmoid followed by BCELoss with BCEWithLogitsLoss removes later overflow. Training then runs 20,000 steps without NaN, with results matching a float32 reference within 0.3 Dice points.

Where the answer stops holding

detect_anomaly adds heavy overhead and should be used only for short debugging runs. A NaN can originate in data, so fixing numerics alone may hide a data defect. Switching to bfloat16 avoids most overflow but has less precision, so check that results agree with a float32 reference.

Distributed training adds sources such as mismatched gradient reductions that need their own checks.

A strong answer covers

  • Locate the first inf or NaN step and operation
  • float16 range limit of 65504 and underflow
  • GradScaler skips steps with non-finite gradients
  • Use logits-based losses and safe normalization
  • Compare with float32 or bfloat16 runs

Common mistakes

  • Clipping gradients before unscaling them, so the threshold applies to scaled values
  • Using sigmoid plus BCELoss in autocast regions instead of BCEWithLogitsLoss
  • Lowering the learning rate repeatedly without checking the data batch that triggered the NaN