Mixed Precision Training › Limitations and Practical Considerations

Mixed precision training is not without tradeoffs and failure modes. Understanding when it can go wrong helps you diagnose training instability and choose the right configuration.

Even with dynamic loss scaling, FP16 accumulates more rounding error than FP32. For most models, this error is negligible compared to the stochasticity of SGD. But for models that are numerically sensitive, such as very deep networks, models trained at very high precision for downstream tasks, or models with operations that involve large sums of small values, the reduced precision can measurably affect final performance.

The BF16 trade-off is slightly different: its wider range prevents overflow, but its 7-bit mantissa means each gradient computation is less precise than FP16. In practice, the stability advantage of BF16 dominates, and models trained in BF16 generally match FP32 quality. But if you observe that your BF16 model plateaus earlier or reaches slightly worse perplexity than your FP32 baseline, the precision reduction may be a factor worth investigating. Running a short FP32 ablation over the same hyperparameters for a few thousand steps is the most direct way to isolate precision as a variable.

The precision degradation from FP16 is most likely to surface in specific architectural patterns: softmax operations with very large or very small logits, layer normalizations over very wide hidden dimensions, and attention mechanisms with very long sequence lengths where many values are summed. These operations involve reducing a large number of values into a smaller result, and the accumulated error from each FP16 multiply-accumulate compounds over the length of the reduction. When training models with very long context windows (tens of thousands of tokens), the cumulative rounding error in attention can become noticeable. This is one reason why modern long-context models almost universally use BF16 rather than FP16: the wider range of BF16 tolerates the larger intermediate values that arise in attention over long sequences.

The speedup from mixed precision depends heavily on the GPU. On older Volta GPUs (V100), FP16 tensor cores provide a 2x throughput advantage for matrix multiplications, but only when matrix dimensions are multiples of 8. On Ampere GPUs (A100), BF16 and FP16 tensor cores provide similar throughput, and dimensions should be multiples of 16 for peak efficiency. On consumer cards without tensor cores, mixed precision may be slower than FP32 due to the casting overhead.

This means that benchmark results transfer only between similar hardware. An FP16 speedup of 2.5x on an A100 does not guarantee any speedup on a GTX 1080. When planning training infrastructure, always benchmark the specific model architecture on the target hardware before committing to a precision strategy. The speedup also depends on the batch size and sequence length: operations that are already memory-bandwidth-limited see a larger benefit from the reduced memory footprint of FP16.

Achieving peak tensor core utilization requires attention to matrix dimensions. The tensor core units on NVIDIA GPUs operate most efficiently when the matrix dimensions are multiples of specific tile sizes: 8 for Volta FP16, 16 for Ampere BF16/FP16, and 32 for Hopper FP8. If your model's hidden dimension, vocabulary size, or feed-forward dimension is not aligned to these tile sizes, you lose a fraction of the potential speedup. This is why large language model architectures often choose hidden dimensions like 4096, 8192, or 16384 (all large powers of 2) rather than arbitrary round numbers like 4000 or 8000. The alignment is not accidental; it is a deliberate choice to maximize hardware utilization.

Some operations cannot run in FP16 because their outputs would overflow or lose too much precision. PyTorch's autocast dispatch table handles the most common cases, but custom CUDA kernels or certain third-party operations may not be registered. If you see NaN values appearing in unusual places, inspect whether a custom operation is receiving FP16 inputs when it expects FP32.

You can force specific operations out of autocast with with torch.amp.autocast(enabled=False): inside the autocast context. This is useful for custom attention implementations, specialized normalization layers, or any operation where you have determined empirically that FP16 produces incorrect results.

Gradient accumulation (accumulating gradients over multiple small batches before an optimizer step) requires care with loss scaling. The scaled gradients should be divided by the number of accumulation steps when computing the loss, before scaling, to ensure the effective loss scale is consistent. Alternatively, use the scaler.step() and scaler.update() pattern only at the true optimizer step, not at each accumulation step. Calling scaler.update() at every accumulation step would incorrectly adjust the scale based on partially-accumulated gradients.

A correct pattern for gradient accumulation with loss scaling looks like this:

The division by accumulation_steps before the backward pass ensures that the accumulated gradient matches the gradient you would have gotten from a single step with the full effective batch. The scaler.update() call happens only once per true optimizer step, so the scale adapts at the correct granularity.

Mixed precision reduces memory for activations, which are the dominant consumer during training. For a model with $N$ parameters, the approximate memory breakdown is:

The total parameter and optimizer overhead is 16N bytes for FP32-only (4N parameters + 8N Adam state + 4N gradients) and approximately 18N for mixed precision (2N working + 4N master + 8N Adam + 4N for FP16 gradients). The parameter memory difference is small. The activation memory, which scales with batch size and sequence length, is halved.

To put concrete numbers on this: a 7-billion parameter model like LLaMA-7B has parameters requiring 28 GB in FP32. With AdamW, the optimizer state adds another 56 GB, for a total of about 84 GB before counting activations. On a single A100 with 80 GB of memory, pure FP32 training is impossible for this model at any batch size. With BF16 mixed precision, parameters drop to 14 GB, and the optimizer state stays at 56 GB (maintained in FP32), totaling 70 GB before activations. You can then fit a small batch. For even larger models, combining mixed precision with ZeRO optimization (covered in the ZeRO Optimization chapter) partitions the optimizer state across GPUs, enabling training that would be otherwise infeasible.

For practitioners facing this decision today, the guidance is relatively clear. If you are training on NVIDIA Ampere or newer hardware (A100, A10, RTX 30/40 series), H100, or Google TPUs, use BF16. The simpler training loop, the elimination of loss scaling overhead, and the consistently stable behavior outweigh BF16's slightly lower fraction precision. The major open-source models released in the past several years (LLaMA, Mistral, Falcon, and their variants) all use BF16, and the training codebases for those models reflect this as the default.

If you are training on older hardware that lacks BF16 support (V100, older RTX 20 series, any Pascal-era GPU), FP16 with dynamic loss scaling is the correct choice. The GradScaler machinery handles the overflow detection reliably in practice, and the training stability of FP16 with loss scaling is well-established for transformer architectures.

If you are running on hardware without any dedicated low-precision units (which is increasingly rare for serious LLM training), mixed precision may offer little speedup and can introduce complexity without benefit. In this case, FP32 training is simpler and avoids the precision management overhead entirely.