Mixed Precision Training › Debugging Mixed Precision Training
Language AI HandbookAll chapters Mixed Precision Training › Debugging Mixed Precision TrainingOne section · Sections in this chapter8 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Mixed precision training introduces new failure modes that require specific diagnostic strategies. Understanding what can go wrong and how to observe it helps you resolve problems quickly.
Diagnosing NaN and Inf Gradients
If your training run produces NaN losses or diverges unexpectedly, the first step is to determine whether the problem originates from overflow in FP16 gradients or from an underlying numerical problem in the model.
With FP16 and GradScaler, monitor the scale value during training. If you see the scale rapidly decreasing from its initial value and continuing to decrease until it reaches very small values (around 1 or below), your gradients are consistently overflowing. This usually means the model architecture or hyperparameters produce very large gradient values. Common causes include a learning rate that is too high, missing or improperly scaled residual connections, or operations that produce large intermediate values (such as logit scales in attention that are not divided by ).
If the scale remains stable but you still see NaN values, the problem is not overflow in the loss-scaled gradients. Instead, look for custom operations that are not registered with autocast and may receive FP16 inputs they cannot handle, for operations that produce NaN in FP32 as well (like taking the log of zero, or dividing by zero), or for numerical instability in rarely-executed code paths (such as the mask filling in attention for padding tokens).
With BF16, there is no loss scaling, so any NaN values indicate an underlying numerical problem rather than a range issue. BF16's wider range means overflow is essentially impossible. If you see NaN values in BF16 training, you can rule out FP16-style overflow and focus on the model architecture.
Checking for Silently Stalled Weights
The vanishing update problem is harder to detect because it is silent. A useful diagnostic is to track the fraction of model parameters whose gradient norms fall below a threshold. You can log gradient norms per layer using PyTorch's model.named_parameters() loop after the backward pass and before the optimizer step.
If you observe that certain layers consistently show near-zero gradient norms while others show normal magnitudes, those layers may have stopped updating. Check whether the problematic layers are early in the network (where gradients can vanish due to multiplicative propagation through many layers) or whether they are layers that produce very small-magnitude features (such as layers operating on near-zero residuals).
Debugging Autocast Scope
A common mistake is placing the loss computation outside the autocast context. The loss function should generally be inside the autocast block for correct dtype handling, even though cross_entropy will internally use FP32 for its computation:
# Correct: loss inside autocast context
with autocast(device_type="cuda", dtype=torch.float16):
logits = model(x)
loss = criterion(logits, y)
# Incorrect: loss outside autocast may produce dtype mismatch warnings
with autocast(device_type="cuda", dtype=torch.float16):
logits = model(x)
loss = criterion(logits, y) # logits is FP16, criterion expects FP32The second pattern often works because PyTorch will implicitly upcast, but it produces efficiency warnings and can cause subtle issues with custom loss functions that do not handle mixed-dtype inputs.