Computing in lower precision to save memory
Float32 (standard precision) uses 4 bytes per number; float16 (half precision) uses 2 bytes. Half the memory overhead is huge at scale: a 100-billion parameter model in fp32 requires 400GB of memory, while fp16 requires 200GB. This difference often determines whether training fits on available hardware.
The catch: fp16 has less numerical precision and smaller range. Very small gradients underflow to zero, and large values overflow. Mixed precision solves this: compute forward and backward passes in fp16 (fast, compact), but keep a master weight copy in fp32 for the optimizer step.
Maintaining numerical stability
Loss scaling is the key trick: multiply loss by a large scale factor (e.g., 2^16) before computing gradients in fp16. This shifts all gradient values to larger, representable numbers. After computing gradients, unscale them back. The scale factor is dynamic: if overflow occurs, reduce the scale; if no overflow, increase it.
With proper loss scaling, mixed precision achieves near-identical convergence curves to pure fp32, with the memory and speed benefits of fp16.
When to use which precision
Bfloat16 (brain float) offers better stability than fp16 because it preserves the full fp32 range, only losing precision in the mantissa. Modern GPUs (NVIDIA's A100+, Google TPUs) support bfloat16. If available, prefer it over fp16. Standard fp32 is reserved for the optimizer state and master weights, not the compute path.