Mixed Precision Training › FP16 Training and Its Challenges

The term "FP16 training" is slightly misleading. Truly training entirely in FP16 does not work reliably for most models. What practitioners mean is a mixed precision approach: use FP16 where it is safe and beneficial, and FP32 where it is necessary. The art is in knowing where to draw that line.

The foundational observation that made mixed precision training practical was established empirically by Micikevicius et al. at NVIDIA in their 2018 paper "Mixed Precision Training." They tested a broad range of model types, including convolutional image classifiers, RNNs for speech recognition, and language models, and found a consistent pattern: the operations dominating training time (large matrix multiplications) are less sensitive to precision than the operations affecting training stability (gradient accumulation, optimizer state updates). This asymmetry allows you to use the cheaper format where speed matters and the expensive format where accuracy matters.

The NVIDIA team also identified the specific mechanisms by which naive FP16 training fails and proposed targeted solutions for each. Their framework is essentially the one still in use today in PyTorch, JAX, and every other major deep learning library. Understanding the problem they were solving clarifies why the solution looks the way it does.

The primary beneficiaries of FP16 are the large matrix multiplications that dominate compute time in transformer training: the attention score computations, the feed-forward layer projections, and the embedding lookups. These operations are memory-bandwidth-bound and compute-bound, and using FP16 halves the data movement and doubles the throughput on hardware with dedicated FP16 units.

When an operation runs in FP16, both the inputs and the outputs are in FP16. The CUDA kernels for these operations internally accumulate in FP32 to avoid precision loss in the accumulation itself, then convert the result back to FP16 before writing to memory. This is a hardware-level behavior for tensor cores, and you get it automatically. The intermediate FP32 accumulation prevents summing thousands of FP16 values from accumulating so much rounding error that the result becomes meaningless. The hardware designers anticipated this and built the accumulation into the tensor core itself.

Modern NVIDIA tensor cores operate in a specific way: they take FP16 inputs, compute the matrix multiply-accumulate in FP32, and write an FP16 result. This is called mixed precision multiplication at the hardware level, and it is the key reason why FP16 matrix multiplications are accurate despite the low-precision inputs. The precision loss occurs only when writing the result back to FP16, not during the arithmetic itself.

Certain operations are sensitive to precision and should remain in FP32:

• The master copy of model weights: Stored in FP32, then cast to FP16 before each forward pass.
• Gradient accumulation: Gradients are accumulated in FP32 to avoid the vanishing update problem.
• The optimizer state: Momentum terms and Adam's second-moment estimates should stay in FP32.
• Batch normalization and layer normalization statistics: Mean and variance computations that reduce over large sets of values benefit from FP32 accumulation.
• Loss computation: The scalar loss itself, especially with cross-entropy over large vocabularies, can involve large intermediate sums.

This "keep a master copy in FP32" pattern is the defining feature of the NVIDIA mixed precision training recipe from their 2018 paper "Mixed Precision Training" by Micikevicius et al. The key insight in that paper was empirical: they tested a wide range of models (convolutional networks, RNNs, language models) and found that the FP32 master weight copy, combined with loss scaling, was sufficient to match FP32 training accuracy in virtually all cases.

The optimizer state represents a large chunk of the memory saved by not needing FP16 for these values. For AdamW, the optimizer state includes the first moment (momentum estimate) and second moment (variance estimate) for every parameter, each stored as FP32. This is already 8N bytes for N parameters. Storing them in FP16 would halve this to 4N bytes, but the precision loss would cause optimizer updates to be inaccurate, ultimately harming model quality. In practice, the optimizer state is kept in FP32 in essentially all serious mixed precision training setups.

To understand intuitively why the optimizer state must stay in FP32: Adam's second moment estimate $v_t$ tracks the running average of squared gradients. For a parameter that receives occasional large gradients and frequent small ones, $v_t$ might be a small positive number like $10^{-8}$. The Adam update divides the gradient by $\sqrt{v_t + \epsilon}$, which is sensitive to the exact value of $v_t$. If $v_t$ rounds from $10^{-8}$ to $10^{-7}$ due to FP16 precision, the effective step size for that parameter changes by a factor of $\sqrt{10}$. Over millions of steps, this corrupts the adaptive learning rate signal that Adam depends on.

Even with master weights in FP32, a persistent problem remains: the FP16 gradients can overflow. Because the FP16 maximum is only 65504, gradient values that legitimately arise during training can become Inf or NaN when cast to FP16. This causes the entire training step to be corrupted.

The key point is that gradient overflow is not a sign of a broken model or bad hyperparameters. Gradient values that are perfectly valid in FP32, and that would produce a correct optimizer step, can simply exceed FP16's representable range. The gradient is not intrinsically wrong; FP16 just cannot hold it. This distinction matters because it explains why scaling the gradient, rather than clipping or discarding it, is the right solution. The gradient carries accurate information about the loss surface; we just need to shift its magnitude into a range where FP16 can represent it.

Loss scaling directly solves this problem.