Mixed Precision Training › Mixed Precision Training

Training large language models demands enormous computational resources. A single forward and backward pass through GPT-3 moves billions of numbers through GPU registers, each consuming memory and time. One of the most practical ways to reduce both is to change which number format those values are stored in. Mixed precision training does exactly this: it keeps some tensors in lower-precision floating point to save memory and speed up computation, while preserving enough numerical accuracy for the model to learn effectively.

The gain is substantial. Modern GPUs execute FP16 matrix multiplications roughly two to four times faster than FP32, and each tensor takes half the memory. For a training run that would otherwise require 80 GB of GPU memory in full precision, mixed precision may bring it down to 40 GB or less. That difference can mean fitting a larger batch, training a bigger model, or using fewer GPUs entirely.

This chapter builds the conceptual foundation for understanding mixed precision training. We start with how floating point numbers work and why precision matters, then walk through the FP16 format and its limitations, introduce BF16 as a more numerically stable alternative, explain the loss scaling technique that makes FP16 training stable, and show how to implement mixed precision training in PyTorch using the torch.amp module. By the end, you will understand how to toggle a setting in your training loop, why the whole system works, and when it might fail.

This chapter sits within the broader Part XXXI: Training Infrastructure. The chapters on GPU Architecture and Memory Management earlier in this part explained why memory is scarce and how the GPU operates. Mixed precision training is one of the most effective answers to that scarcity. Later chapters on Communication Optimization and Checkpointing and Recovery will address other bottlenecks that emerge as training scales further.