Mixed Precision Training › Mixed Precision Architecture in Practice

The full mixed precision training setup combines the observations above into a coherent system. Let us walk through how it works end-to-end, then examine the memory implications carefully.

The model maintains two copies of its parameters:

1. FP16 (or BF16) working copies: Used for the forward pass and backward pass. The matrix multiplications and attention operations run on these copies using fast low-precision arithmetic.
2. FP32 master copies: The authoritative version of the weights. Optimizer states (momentum, variance) are also kept in FP32 here.

At each training step, the working copies are created by casting the FP32 masters to the low-precision format. After gradients are computed and scaled back to FP32, they are applied to the FP32 masters. The working copies are discarded and recomputed at the next step.

This two-copy scheme has a memory cost: the FP32 masters exist alongside the FP16 working copies. For a model with $N$ parameters, you store $2N$ bytes for the FP16 working copies and $4N$ bytes for the FP32 masters, totaling $6N$ bytes rather than the $4N$ bytes of pure FP32 training. The gain comes from the activations during the forward and backward passes, which are the dominant memory consumer for large models. Activations in FP16 use half the memory of FP32 activations, and this saving more than compensates for the extra master copy.

The reason activations dominate memory rather than parameters is a consequence of the size of modern models and the batch sizes used during training. A transformer layer with $d_{\text{model}} = 4096$ and a sequence of length 2048 produces activation tensors with millions of elements per layer, multiplied by the batch size. The parameter count for that same layer is fixed at roughly $4 \times d_{\text{model}}^2 \approx 67$ million parameters, while the activations scale with both sequence length and batch size. At large batches and long sequences, activations dwarf parameters in memory, making FP16 activations the dominant source of memory savings.

Tracing a single training step in mixed precision clarifies which operations happen in which precision and where the memory savings and costs accumulate.

First, at the start of the step, the FP32 master weights are cast to FP16 working copies. This casting is fast (a bit-level operation) and the FP16 copies live in GPU memory alongside the FP32 masters. The extra memory for the FP16 copies is $2N$ bytes.

Second, the forward pass runs with the FP16 working copies. Activations, which include all the intermediate tensors produced by each layer's computation, are stored in FP16. For a batch size of $B$ and sequence length $T$, the activations at each of the $L$ transformer layers occupy roughly $O(B \times T \times d_{\text{model}})$ memory in FP16. These activations must be retained in memory until the corresponding backward pass computes their gradients, which for a $L$-layer model means the first layer's activations must survive until the last backward step. This is why activation memory scales with model depth and batch size, and why the FP16 savings here are so valuable.

Third, the backward pass computes gradients in FP16 (or with FP32 accumulation for tensor cores, as described earlier). If using loss scaling, the gradients are in scaled FP16. The GradScaler.unscale_() call converts these to FP32 by dividing by the scale factor, in-place, overwriting the FP16 gradient storage with FP32 values. This temporarily increases gradient memory by 2x for the brief duration of the unscaling, but the FP16 gradient memory is freed immediately after.

Fourth, the FP32 gradients are applied to the FP32 master weights by the optimizer. Adam updates its FP32 momentum and variance estimates and computes the updated master weights. The FP16 working copies are now stale and can be freed (they will be recast from the updated masters at the next step).

The net effect is that the peak memory consumption is lower than pure FP32 because the large activation tensors are stored in FP16, and this saving is much larger than the overhead of the FP32 master weight copy for typical transformer workloads.

PyTorch's automatic mixed precision (AMP) framework handles the casting decisions automatically. The torch.autocast context manager wraps code regions where operations should run in lower precision. Within this context, PyTorch consults an internal dispatch table that maps each operation to its appropriate precision:

• Matrix multiplications (torch.mm, torch.bmm, torch.linear): FP16/BF16
• Attention operations: FP16/BF16
• Convolutions: FP16/BF16
• Loss functions: FP32
• Reduction operations (sum, mean): FP32
• Batch norm, layer norm: FP32 (for statistics, though parameters can be in FP16)

The autocast mechanism inserts casts at operation boundaries automatically. You do not need to call .half() or .to(torch.float16) anywhere in your model code. Operations that benefit from low precision run in low precision; operations that need FP32 stay in FP32.

The dispatch table is based on empirical evidence from the PyTorch and NVIDIA teams about which operations are safe in lower precision. Operations that perform large reductions (softmax, layernorm) are kept in FP32 because their outputs are sensitive to accumulation error. Operations that perform independent, parallel computations (linear layers) are safe in FP16 because the per-element error does not compound within a single operation.

One practical consequence of the autocast dispatch table is that operations can change precision across calls even within the same model. A linear layer will run in FP16 under autocast, but a subsequent layer normalization will cast its inputs to FP32. This cast has a small overhead, typically negligible compared to the compute savings, but it means that the model's internal data flow involves more dtype transitions than a pure FP32 model. If you profile a mixed precision model with PyTorch's profiler, you will see autocast_cpu_cast or _cast operations interspersed between the main computations. These are expected and not a sign of a problem.

Using BF16 instead of FP16 in PyTorch AMP is a single argument change: dtype=torch.bfloat16 instead of the default FP16. When using BF16, you should omit the GradScaler entirely, since loss scaling is unnecessary. The training loop becomes visually simpler and conceptually cleaner. Many practitioners who have switched from FP16 to BF16 report that the transition required only a handful of lines of code and improved training stability, particularly for large models.

The BF16 path through autocast is slightly different from FP16 at the hardware level: the CUDA kernels for BF16 tensor core operations use different instruction variants than FP16. The throughput is similar on Ampere and newer hardware, so the practical speedup is comparable. On some operations, BF16 may be slightly faster because the simpler exponent alignment (BF16 shares the FP32 exponent format) reduces conversion overhead.