The Promise and the Pain of FP16
Mixed precision training sounds easy on paper. Cast your model weights to float16, train faster, use less memory, and get the same accuracy. In practice, I've spent more hours debugging FP16 instabilities than I'd like to admit. The issues are subtle, intermittent, and always seem to show up at 80% through a week-long training run.
Here's the basic setup with PyTorch's AMP. It works beautifully for most standard architectures.
import torch
from torch.cuda.amp import autocast, GradScaler
model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler()
for batch in dataloader:
optimizer.zero_grad()
with autocast(dtype=torch.float16):
output = model(batch['input'].cuda())
loss = criterion(output, batch['target'].cuda())
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
The GradScaler handles loss scaling — multiplying the loss by a large factor before backward to push small gradients into FP16's representable range, then unscaling before the optimizer step. When gradients overflow to inf, the scaler skips that optimizer step and reduces the scale factor. Simple enough in theory.
Where FP16 Actually Breaks
The Underflow Problem
FP16 can represent values down to about 6e-8. That sounds small, but gradient magnitudes in deep networks routinely drop below this. Early layers in a 100+ layer network often produce gradients around 1e-9 or smaller. Those get rounded to zero in FP16, and your early layers stop learning entirely. You won't see this in the loss curve right away — the model appears to train fine for thousands of steps until you notice the frozen early layers.
The fix is loss scaling, but the default dynamic loss scaling isn't always aggressive enough. I've had to manually set initial scale factors as high as 2^16 or 2^20 for deep architectures.
Related reading: GPU Resource Management for Shared ML Training Clusters.
scaler = GradScaler(init_scale=2**16, growth_interval=100)
Reduction Operations
Summing many FP16 values loses precision catastrophically. Consider summing 1000 values around 0.001 — the true sum is 1.0, but accumulation errors in FP16 can give you 0.85 or 1.2. This hits loss computation, gradient all-reduce in distributed training, and any custom reduction in your model.
PyTorch's autocast keeps reductions in FP32 by default, but if you've written custom CUDA kernels or use non-standard reduction ops, you need to be explicit about precision. I've seen a team lose two weeks of training because a custom attention kernel accumulated softmax scores in FP16.
Batch Normalization Statistics
Running mean and variance in batch norm must stay in FP32. If they drift into FP16, the variance can underflow to zero, which makes the normalization divide by zero and produces NaN that cascades through the entire network. PyTorch's built-in BatchNorm handles this correctly, but custom normalization layers need explicit FP32 buffers.
class SafeLayerNorm(torch.nn.Module):
def __init__(self, dim):
super().__init__()
self.weight = torch.nn.Parameter(torch.ones(dim))
self.bias = torch.nn.Parameter(torch.zeros(dim))
self.eps = 1e-5
def forward(self, x):
# Force FP32 for mean/variance computation
x_float = x.float()
mean = x_float.mean(-1, keepdim=True)
var = x_float.var(-1, keepdim=True, unbiased=False)
x_normed = (x_float - mean) / torch.sqrt(var + self.eps)
return (x_normed * self.weight + self.bias).to(x.dtype)
BF16: The Better Float16
BFloat16 uses the same exponent range as FP32 (8 exponent bits vs FP16's 5) but with reduced mantissa precision (7 bits vs FP16's 10). This means BF16 can represent the same range of magnitudes as FP32, which eliminates most underflow and overflow problems. The tradeoff is less precision per value, but in my experience, neural network training tolerates mantissa noise far better than range limitations.
Related reading: Embedding Space Quality Metrics and Debugging Techniques.
On A100 and newer GPUs, BF16 runs at the same speed as FP16 on tensor cores. There's genuinely no reason to use FP16 on these chips unless you need the extra mantissa bits for inference accuracy. Training with BF16 doesn't even need loss scaling — the gradient magnitudes stay in range without it.
# BF16 training — no scaler needed
for batch in dataloader:
optimizer.zero_grad()
with autocast(dtype=torch.bfloat16):
output = model(batch['input'].cuda())
loss = criterion(output, batch['target'].cuda())
loss.backward()
optimizer.step()
Diagnosing Precision Issues
When a training run goes sideways after enabling mixed precision, I follow this debugging checklist. First, check for NaN in the loss. If NaN appears immediately (first 100 steps), it's usually an overflow — the loss or an intermediate activation exceeds FP16's max of 65504. If NaN appears after thousands of steps, it's more likely underflow in gradients or a slow accumulation of precision errors.
Second, compare a short training run in FP32 versus mixed precision. If they diverge within 500 steps, you've got a precision bug. If they match for 500 steps but diverge later, the issue is likely in gradient accumulation or optimizer state precision.
Third, enable anomaly detection to catch the exact operation that produces NaN.
See also: Sentence Embedding Models: Contrastive Learning and Evaluati.
torch.autograd.set_detect_anomaly(True)
# Run a few training steps — this will pinpoint the NaN source
# WARNING: this is extremely slow, only use for debugging
Optimizer State Precision
Adam and AdamW maintain first and second moment estimates for every parameter. These must remain in FP32, even when model parameters are in FP16. The second moment (squared gradient running average) involves very small numbers that underflow easily in FP16. PyTorch's optimizers keep states in FP32 by default, but some custom or memory-optimized implementations might cut this corner. Always verify.
The memory overhead matters. For a 7B parameter model, FP32 optimizer states for Adam consume about 56GB (two FP32 copies per parameter — 8 bytes × 2 × 7B). This is often the binding constraint on single-GPU training, not the model weights or activations. It's why optimizers like 8-bit Adam exist — they quantize the optimizer states to save memory while keeping updates numerically stable.
Practical Recommendations
If you're on A100 or newer GPUs, use BF16. Full stop. It eliminates an entire class of precision bugs with zero performance penalty. Keep your optimizer states in FP32, and don't bother with loss scaling.
If you're on V100 or older GPUs that lack BF16 tensor cores, use FP16 with AMP and the GradScaler. Start with a high initial scale (2^16) and a conservative growth interval (200 steps). If you see gradient overflow more than once per 100 steps, something in your model is producing unusually large activations — investigate rather than just reducing the scale.
For inference, FP16 works reliably on all architectures I've tested. The precision requirements during inference are much lower since there's no gradient computation. You can safely quantize further to INT8 for inference, which I'll cover in a separate article on quantization.