FP16 numerical issues — short summary:- FP16 has ~11 bits exponent range ±65,504 and 10 bits mantissa → much smaller dynamic range and precision than FP32. Underflows (gradients/activations become zero), overflows (large accumulations -> Inf/NaN), and rounding errors are common. Small gradients can be flushed to zero; sums/variance estimates lose precision. These cause training instability and sudden NaNs.How AMP (automatic mixed precision) manages loss scaling:- AMP runs forward in FP16 where safe (matrix multiplies, convs) and keeps some ops/weights in FP32 (master weights, BatchNorm). Because FP16 can underflow small gradients, AMP multiplies the scalar loss by a large scale factor before backprop (loss scaling) so gradients are larger and don’t underflow to zero.- After backprop, gradients are unscaled (divided by the same factor) before updating FP32 master weights.- Two loss-scaling strategies: - Static loss scaling: fixed scale chosen by user. Simple but brittle — too large -> overflow/NaN; too small -> underflow. - Dynamic loss scaling: start with a large scale, reduce it when inf/NaN detected in gradients, increase it slowly if no overflows for N steps. This is the default in modern AMP implementations (e.g., torch.cuda.amp.GradScaler).Concrete debugging & mitigation checklist (step-by-step):1) Reproduce and isolate- Run a few steps with amp but in a controlled setting: smaller batch, single GPU, deterministic seeds.- Enable gradient/activation checks every step.2) Detect where NaN/Inf appears- Add runtime checks after forward and after backward:python
def has_inf_nan(tensor):
return torch.isinf(tensor).any() or torch.isnan(tensor).any()
# check activations (example)
for name, act in model.named_buffers(): pass
# simpler: check gradients after backward
for p in model.parameters():
if p.grad is not None and has_inf_nan(p.grad):
print("NaN in grad:", p.name)
- Log layer outputs and gradients to find first layer producing Inf/NaN.3) Loss-scaling actions- Use dynamic loss scaling (recommended). In PyTorch: GradScaler() handles this; in TF-MixedPrecision use LossScaleOptimizer with "dynamic".- If using static scaling (e.g., for reproducibility), pick conservative smaller scale (e.g., 128, 256) and test. If you see frequent overflow, reduce the scale.4) If NaNs persist, enable framework option to skip NaN gradients- Some frameworks offer skip_nan_gradients / skip_if_overflow behavior that prevents optimizer.step when grads contain NaN/Inf and optionally zeros them. This avoids corrupting master weights. Use it as a temporary mitigation while you find root cause. Note: skipping updates can stall training for a while; monitor scale reductions and frequency.5) Verify custom ops and third-party layers- Custom CUDA kernels or ops often aren’t AMP-safe (might assume FP32 or use inadequate accumulation). Ensure: - Kernels support FP16 inputs and use FP32 accumulation where needed (e.g., in reduction sums). - Autograd implementations correctly propagate gradients and preserve dtype semantics. - Add dtype casts explicitly around unsafe ops: cast inputs to FP32, perform op, cast back to FP16 or keep FP32 outputs.- Run with AMP’s “debug/verbose” flags if available to see which ops are auto-promoted or excluded.6) Optimizer and hyperparameter adjustments- Use FP32 master weights (standard AMP practice). Ensure optimizer updates happen on FP32 params.- Reduce learning rate and/or grad clipping threshold — large updates are more likely to produce overflow in FP16 pathways.- Try switching optimizer (e.g., AdamW vs SGD) or reduce Adam’s betas / eps to more stable defaults. For Adam, increasing eps (e.g., 1e-6 -> 1e-4) can stabilize denomations.- Enable gradient clipping (clip by norm) before optimizer.step.- Use larger batch normalization eps or convert BatchNorm to FP32 (many frameworks already do this automatically).7) Numerical robustness practices- Keep certain ops in FP32 (LayerNorm/BatchNorm/softmax/logsumexp/normalization and loss computations).- Use mixed-precision safe implementations (cuBLAS, cuDNN, fused kernels from vendor libraries).- Accumulate gradients (gradient accumulation) with smaller micro-batches if batch size causes instability.8) Logging & metrics- Track frequency of scale reductions (if using dynamic scaling), number of skipped steps, and loss growth patterns. If scale quickly collapses to 1 and overflows persist, likely a bug in model or custom op.Example troubleshooting sequence (practical):1. Switch to AMP dynamic scaling (GradScaler). Run with smaller LR and batch size.2. Add checks to identify first NaN-producing layer.3. If NaN found in custom op, cast to FP32 inside op or rewrite op accumulation in FP32.4. If NaNs are sporadic, enable skip_nan_gradients to avoid corrupt updates while continuing training and monitoring scaler behavior.5. Adjust optimizer eps / weight decay / gradient clipping as additional mitigation.Why these help — reasoning:- Loss scaling addresses underflow by making gradients representable in FP16; dynamic scaling prevents overflow by adapting scale when overflows occur.- Checking activations/gradients pins down the faulty location — often a single op causes blowup.- Converting numerically sensitive ops to FP32 or using FP32 accumulation preserves precision and avoids catastrophic cancellation or overflow.- Optimizer parameter tweaks reduce update magnitude and denominator instability, reducing propensity for Inf/NaN.Edge cases & final notes:- Intermittent NaNs can come from data (labels, targets, invalid values) — validate dataset for extremes.- Mixed precision uncovers bugs that FP32 masks; treat NaNs as signals to inspect model math and custom kernels.- Always keep master FP32 copy of weights and prefer dynamic loss scaling as the first line of defense.