Mixed Precision Training

Training a neural network involves an enormous number of matrix multiplications and gradient calculations. Mixed precision training speeds up that work by using lower-precision numbers where they are safe, while preserving higher precision where small numerical errors would damage learning.

Two number formats, one training run
Traditional training stores activations, gradients, and weights as 32-bit floating-point values (FP32). Mixed precision combines FP32 with a 16-bit format such as FP16 or bfloat16:

Large matrix operations run in 16-bit precision, which modern GPUs accelerate using specialized hardware such as Tensor Cores.
Numerically sensitive operations—such as some reductions, normalization calculations, and parameter updates—remain in FP32.
Optimizers commonly keep an FP32 “master copy” of parameters, even when the model uses 16-bit copies during forward and backward passes.

Why loss scaling is needed
FP16 has a limited range: very small gradients can round down to zero, especially in deep networks where gradients already weaken as they flow backward. Loss scaling multiplies the loss by a large scale factor before backpropagation, making gradients large enough to represent in FP16. Before the optimizer updates parameters, those gradients are divided by the same factor, so the intended update is unchanged. Dynamic loss scaling reduces the scale after an overflow—detected as inf or NaN gradients—and raises it when training remains stable. bfloat16 has a much wider range than FP16, so it usually needs less scaling.

Practical training behavior
PyTorch’s automatic mixed precision uses torch.autocast to select suitable precisions per operation and GradScaler for FP16 loss scaling. The payoff is lower activation-memory use and substantially faster training, enabling larger batches or models. But mixed precision does not cure a bad learning rate: a loss curve that suddenly becomes NaN can still signal diverging optimization. It also adds care around custom operations, gradient clipping (unscale gradients first), and debugging numerical instability. Used correctly, it preserves near-FP32 training quality while making hardware work far more efficiently.

Mixed precision training trains a neural network using both lower-precision floating-point formats, typically FP16 or BF16, and FP32. Compute-intensive operations use lower precision for speed and memory savings, while numerically sensitive values such as master weights, gradient accumulation, and selected reductions remain in FP32. Loss scaling prevents small FP16 gradients from underflowing. It enables larger models and batches with near-FP32 training accuracy.

Imagine doing a huge shopping list with two kinds of notes: quick shorthand for everyday items, and careful full-detail notes for anything where a tiny mistake would matter. Mixed precision training uses a similar idea when teaching an AI system.

Most of the work is done with less detailed numbers, which are faster to process and use less memory. For the delicate parts—where rounding errors could make learning go wrong—it keeps more detailed numbers. This lets researchers train larger AI models faster, often without reducing their quality. It matters because modern AI can require enormous amounts of computing power.