Mixed Precision Training › Summary

Mixed precision training exploits the difference between numerical range and precision in floating point formats to achieve faster training with less memory, without materially harming model quality.

The key concepts are:

• Floating point formats: FP32 (32 bits, 8 exponent, 23 mantissa), FP16 (16 bits, 5 exponent, 10 mantissa), BF16 (16 bits, 8 exponent, 7 mantissa). Exponent bits control range; mantissa bits control precision.
• FP16 limitations: The narrow range (max ~65504) risks overflow; the limited precision risks underflow for small gradients. Both cause training instability.
• Loss scaling: Multiply the loss by a large constant $S$ before the backward pass to shift gradient magnitudes into FP16's representable range. Dynamic loss scaling adjusts $S$ automatically by detecting overflow and adapting.
• BF16 advantages: Same numerical range as FP32 (8 exponent bits), making overflow essentially impossible. No loss scaling required. First-class support on modern GPUs and TPUs. Has become the default for large-scale transformer training.
• Mixed precision implementation: Use torch.autocast to run matrix multiplications and attention in low precision. Use torch.amp.GradScaler with FP16, omit it with BF16. Keep FP32 master weights and optimizer states.
• The two-copy pattern: FP16 working copies for the forward and backward passes, FP32 master copies for the optimizer update. This preserves weight update precision while gaining activation memory savings.
• Practical considerations: Tensor core efficiency requires matrix dimension alignment; gradient accumulation must account for loss scaling; custom operations may need explicit FP32 coercion.

The next chapter on Communication Optimization extends these infrastructure ideas to distributed training, where the precision and format of gradients communicated between GPUs adds another layer of efficiency to consider.