Mixed Precision and GPU Memory Budgets
Training a model in a mix of 16-bit and 32-bit numbers roughly halves the memory it needs and speeds up matrix multiplies, at the cost of a few numerical traps to guard against.
Every number a model stores — weights, activations, gradients — normally sits in 32-bit floating point (fp32). Mixed precision keeps most of that arithmetic in 16-bit (fp16 or bfloat16) instead, which uses half the memory and runs faster on GPU tensor cores built specifically to accelerate 16-bit matrix multiplication, while keeping a small number of sensitive operations — like the master copy of weights and certain sums — in fp32 so the model still trains stably.
Mixed precision is not "train in low precision and hope" — it is a specific recipe: fp16 for the bulk of the compute, an fp32 master weight copy for updates, and loss scaling to stop small gradients from rounding to zero.
The memory saving is what actually matters operationally. A model that needs 40GB of GPU memory in pure fp32 might fit in 20-25GB under mixed precision, which can be the difference between fitting on one GPU or needing to shard across several. The catch is that fp16 has a much narrower exponent range than fp32, so gradients that are individually tiny can underflow to exact zero before they ever get applied, silently stalling learning in parts of the network.
Worked example. A gradient value of 0.00003 is well within fp32's range but underflows to zero in fp16. Loss scaling fixes this by multiplying the loss by a large constant (say, 1024) before backpropagation, which scales every gradient up by the same factor, keeping small values representable; the gradients are then divided back down by 1024 before the optimizer step. Frameworks now do this automatically, growing the scale factor over training and backing off if it causes overflow (inf/nan) instead of underflow.
For a quant research team, the practical upshot is that mixed precision is close to a free speedup on modern GPUs and is usually enabled by default in training frameworks — the main job left to the practitioner is watching training logs for NaN losses, which usually means the automatic loss scale needs tuning down.
Related concepts
Practice in interviews
Further reading
- Micikevicius et al., 'Mixed Precision Training' (NVIDIA/Baidu)