Mixed Precision

Recap: why and how

Modern accelerators run BF16/FP16 multiplies 2-8x faster than FP32. Mixed precision keeps the model weights in FP32 (master copy) but does the heavy compute (matmuls) in lower precision. The result is faster training with the same final accuracy in 99% of cases.

BF16: same dynamic range as FP32, less precision; no GradScaler needed; preferred default in 2026. FP16: more precision, smaller dynamic range; needs GradScaler to handle small gradients without underflow.

Tip: If your hardware supports BF16 (H100, A100, M3+), default to BF16. Skip the GradScaler complexity. Only fall back to FP16 if you're on hardware that doesn't have native BF16 (consumer GPUs older than RTX 3000 series).

What goes wrong when it goes wrong

NaN losses. The most common cause: a single underflow somewhere in the FP16 forward pass that propagates through the loss. Usually fixed by switching to BF16, or by adding a GradScaler if you're stuck with FP16.

Subtle accuracy degradation. Some operations (e.g. log-softmax, layer norm) are sensitive to lower precision. PyTorch autocast keeps these in FP32 automatically, but custom layers may not. If your model accuracy drops noticeably under autocast, profile which ops are running in low precision.

The numbers

BF16 training: 1.5-3x faster wall-clock vs FP32 on Tensor Cores, ~50% memory savings on activations. FP8 (H100 and newer) can give another 30-50% on top of that for transformer training, but the tooling is still maturing in 2026.

Principle: Mixed precision is essentially free on modern hardware. If you're not using it, you're paying 2x time and memory for no good reason.

Code

BF16 mixed precision (no scaler needed)·python

import torch
from torch.amp import autocast

device = torch.device("cuda")
model.to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)

for xb, yb in train_loader:
    xb, yb = xb.to(device), yb.to(device)
    opt.zero_grad()
    with autocast(device_type="cuda", dtype=torch.bfloat16):
        logits = model(xb)
        loss = loss_fn(logits, yb)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    opt.step()

Recap: why and how

What goes wrong when it goes wrong

The numbers

Code

External links

Exercise

Progress

Comments 0