Mixed-precision training uses lower precision where it is safe, usually FP16 or BF16, and FP32 where it is needed. It reduces activation memory and can use lower-precision Tensor Core paths for matrix multiplications.

Parameters and master parameters

Parameters are the model weights used during forward and backward passes to compute activations and gradients. FP16 updates can lose small but meaningful changes, so conventional FP16 training keeps a separate FP32 master copy of each weight for the optimizer.

The optimizer updates the master parameters in FP32, then casts the updated values to FP16 for the model parameters used by forward and backward. Full FP32 training does not need master parameters because the model parameters already have the required precision.

This example assumes Adam with two FP32 moment buffers. Both approaches use 16 bytes of persistent state per parameter, so conventional mixed precision does not substantially reduce parameter, gradient, or optimizer-state memory.

Mixed precision usually saves substantial activation memory. Activations stored in FP16 or BF16 take about half the space of FP32 activations, though operations that require FP32 and other temporary buffers reduce the total saving. When activations are the bottleneck, this can permit a much larger batch, sometimes close to 2×.

Loss scaling

FP16 has limited precision and range. Small gradients can underflow to zero, so loss scaling multiplies the loss by a scale factor before backpropagation. This multiplies every gradient by the same factor. Before the optimizer step, the gradients are divided by that factor again.

Static loss scaling uses a fixed scale. Dynamic loss scaling adjusts the scale at runtime, reducing it when overflow is detected and increasing it when training is stable.

BF16 has the same exponent range as FP32, so it does not usually need loss scaling. The exponent determines whether a number can be represented at all, while the mantissa determines its precision. Loss scaling addresses underflow and existence in the representable range. It does not restore lost mantissa precision.

Automatic mixed precision

Automatic mixed precision, or AMP, selects a suitable dtype for each operation. PyTorch uses operator-specific autocast rules. Its autocast context manager or decorator runs eligible operations in lower precision and keeps numerically sensitive operations in FP32, inserting the necessary input casts. An operation that runs in FP32 is not always immediately cast back. Later autocast rules determine the dtype of its consumers.

References

• PyTorch automatic mixed precision