Mixed Precision

Modern accelerators run half-precision arithmetic two to eight times faster than single precision and use half the memory per tensor. Mixed-precision training exploits this without sacrificing convergence: most operations execute in float16 or bfloat16, while a small set of numerically sensitive ones (loss, softmax, layer norm, weight updates) stay in float32. The result is a training run that fits a larger model in the same memory and finishes substantially faster, with model quality nearly indistinguishable from full-precision training.

This section covers PyTorch's Automatic Mixed Precision (AMP) machinery, the differences between the two half-precision formats, when the gradient scaler is required, and the half-dozen pitfalls that bite practitioners moving from FP32 to AMP for the first time. Chapter 9 covers the related but distinct topic of low-precision inference via post-training quantization.

Why Mixed Precision Works

The motivation is hardware. NVIDIA's Tensor Cores (Volta and later), AMD's Matrix Cores (RDNA3 and later), Apple Silicon's Neural Engine, and Google's TPUs all expose dedicated matrix-multiply units that run at half precision dramatically faster than the general-purpose float32 path. On an A100, a float16 matmul achieves roughly 312 TFLOPS versus 19.5 TFLOPS for float32; the difference is large enough that mixed-precision training is the default for any serious workload.

The catch is numeric range. Float16 has 5 exponent bits, giving a representable range of roughly $\pm 6 \times 10^{4}$ and a smallest positive normal of $\approx 6 \times 10^{-5}$. Many gradient values, especially in deep networks late in training, fall below that smallest value and silently round to zero. Bfloat16 has 8 exponent bits (the same as float32), giving the same range as float32 but only 7 bits of mantissa precision. Mixed-precision training keeps the weights, gradients, and optimizer state in float32 (the master copy) while the forward and backward passes run in the half-precision format the hardware prefers.

torch.autocast: The Forward-Pass Wrapper

PyTorch's torch.autocast is a context manager that automatically chooses the right precision for each operation inside its scope. It is a thin layer around the operator dispatch system: matmul and conv run in the chosen half precision, while softmax, log, exp, and the loss reduction stay in float32. The user does not need to manually cast tensors; autocast inserts the conversions transparently.

import torch

device = "cuda"
model = NeuralNetwork(50, 10).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

for features, labels in train_loader:
    features = features.to(device, non_blocking=True)
    labels   = labels.to(device, non_blocking=True)

    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        logits = model(features)
        loss = torch.nn.functional.cross_entropy(logits, labels)

    loss.backward()
    optimizer.step()

The device_type argument selects which dispatch table to use; the most common values are "cuda", "cpu", and "xpu" (Intel). The dtype argument selects the half-precision format: torch.float16 or torch.bfloat16 on CUDA; torch.bfloat16 on CPU (with the BF16 ISA extension); MPS supports both. The autocast scope should wrap the forward pass and the loss computation; the backward pass and the optimizer step happen outside it so that gradients land in float32 storage.

bfloat16 vs float16

For new code, the choice is almost always bfloat16. Its float32-equivalent range means gradients do not underflow, the gradient scaler (covered below) is not needed, and training rarely requires any recipe changes from the float32 baseline. The only downsides are slightly less precision than float16 (7 vs 10 mantissa bits) and lack of support on pre-Ampere NVIDIA GPUs (V100 and older).

Float16, by contrast, has a wider hardware install base (every GPU from Pascal onward) and slightly better precision per bit, but its narrow exponent range demands gradient scaling. The exact recipe is to multiply the loss by a large constant before backward() (so gradients are pulled up into the representable range), then unscale them before clipping or the optimizer step. PyTorch's torch.amp.GradScaler automates this, dynamically adjusting the scale factor to be as large as possible without causing overflow.

import torch

scaler = torch.amp.GradScaler("cuda")

for features, labels in train_loader:
    features = features.to(device, non_blocking=True)
    labels   = labels.to(device, non_blocking=True)

    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        logits = model(features)
        loss = torch.nn.functional.cross_entropy(logits, labels)

    # Scale the loss to lift gradients into the float16 range.
    scaler.scale(loss).backward()
    # Unscale before clipping (clipping reads .grad directly).
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    # step() and update() handle non-finite gradients gracefully.
    scaler.step(optimizer)
    scaler.update()

Common AMP Pitfalls

Five mistakes account for the majority of AMP-related training failures.

Using float16 without the scaler

Gradients silently underflow to zero; training appears to start but loss does not decrease. The fix is to add GradScaler; the alternative is to switch to bfloat16, which sidesteps the issue.

Calling clip_grad_norm_ on scaled gradients

The clipping threshold is interpreted against gradients that are still scaled up by the loss scaler. Either the clip threshold has to be multiplied by the current scale (fragile), or, much better, scaler.unscale_(optimizer) must be called before clipping. The pattern is in Code Fragment E.6.2 above.

Loss computation inside autocast on tiny values

Some loss functions accumulate many tiny intermediate values whose sum is dominated by precision loss in float16. Wrap the loss in with torch.autocast(device_type="cuda", enabled=False): after the forward pass (the activations are already half-precision; only the reduction needs to escape) to force float32 accumulation.

NaN in a single batch propagating to all parameters

One overflow in a forward pass produces NaN logits, NaN loss, NaN gradients, and NaN parameters after the next step. The scaler's scaler.step already skips parameter updates when any gradient is non-finite, so float16 training tolerates the occasional overflow. With bfloat16, ensure the training loop checks torch.isfinite(loss) and skips the step when not. Anomaly detection (Section E.9) helps localize the source.

Saving and loading checkpoints across precisions

Mixed-precision training keeps a float32 master copy of weights. The state dict therefore contains float32 tensors regardless of the autocast format. This is correct: do not save half-precision weights for resumption; doing so loses precision that float32 master copies are designed to preserve. For deployment, a separate cast-to-half-precision step is fine and is what most inference servers do.

Memory Savings

Mixed precision halves the memory used by activations (the largest contributor for moderate batch sizes), but does not halve total memory: parameters, gradients, and optimizer state still live in float32. For a transformer trained with AdamW, the memory breakdown is approximately 4 bytes per parameter for weights, 4 bytes for gradients, 8 bytes for Adam's two momentum buffers, and a variable amount for activations. AMP shrinks the activation cost and may shrink the gradient cost, but the optimizer state remains the limiting factor for the largest models. The fix for that is sharded optimizer state, covered in Section E.7 on distributed training.