Mixed precision training

By default, PyTorch trains in float32 — 32-bit floating point. Modern GPUs (Ampere and later) have dedicated hardware for 16-bit floats that runs 2–8x faster than fp32, using half the memory. Mixed precision training is how you get that speedup without rewriting your model.

The pitch

Both 16-bit formats use half the memory of fp32, but they make different trade-offs. fp16 keeps 10 mantissa bits (more precise for moderate values) but only 5 exponent bits — so anything above ~65504 overflows. bf16 keeps fp32’s 8 exponent bits (same dynamic range) but only 7 mantissa bits — more rounding noise, never overflow.

In practice, bf16 just works in almost every case. fp16 needs extra machinery (GradScaler) to prevent gradient underflow.

The math, in NumPy

Let’s see the overflow problem with fp16 directly.

import numpy as np

# fp32 can hold this comfortably
x = np.array([65000.0, 1e-7, 1e-30], dtype=np.float32)
print("fp32:", x)

# fp16 has a max around 65504 and underflows below ~6e-8
y = x.astype(np.float16)
print("fp16:", y)
# Notice: 65000 squared overflows
print("fp16 squared (overflow!):", (y * y))

# bf16 keeps fp32's dynamic range
# NumPy doesn't have bf16 natively, but we can simulate by checking the format
# bf16: 1 sign + 8 exponent + 7 mantissa (vs fp16's 1+5+10)
# Same exponent bits as fp32 (8) -> same range, ~7 digits of precision
print("\nbf16 has fp32's exponent range but only ~7-bit mantissa precision.")
print("So large products that overflow fp16 are fine in bf16,")
print("but bf16 has more roundoff noise than fp16 for moderate values.")

fp32: [6.5e+04 1.0e-07 1.0e-30]
fp16: [6.5e+04 1.2e-07 0.0e+00]
fp16 squared (overflow!): [inf  0.  0.]

bf16 has fp32's exponent range but only ~7-bit mantissa precision.
So large products that overflow fp16 are fine in bf16,
but bf16 has more roundoff noise than fp16 for moderate values.

That’s the core difference. fp16’s range is so narrow that gradients during training routinely underflow to zero (vanishing) or overflow to inf. bf16 has the same range as fp32 — just less precision — so it trains stably without help.

The autocast pattern

In PyTorch, mixed precision is enabled with two pieces:

import torch
from torch.amp import autocast, GradScaler

# bf16 — no GradScaler needed
for x, y in loader:
    x, y = x.cuda(), y.cuda()
    with autocast(device_type="cuda", dtype=torch.bfloat16):
        pred = model(x)
        loss = loss_fn(pred, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

autocast automatically casts operations to bf16 where it’s safe (most matmuls, convolutions) and keeps them in fp32 where precision matters (reductions, softmax in some cases, loss accumulation). You don’t have to think about which ops to cast.

For fp16, you also need a GradScaler:

scaler = GradScaler()

for x, y in loader:
    x, y = x.cuda(), y.cuda()
    with autocast(device_type="cuda", dtype=torch.float16):
        pred = model(x)
        loss = loss_fn(pred, y)
    scaler.scale(loss).backward()       # scale loss UP to prevent gradient underflow
    scaler.step(optimizer)              # unscales gradients, then steps
    scaler.update()                     # adjust scale factor for next iter
    optimizer.zero_grad()

The GradScaler multiplies the loss by a big number (e.g. 65536) before backward. Gradients are then big enough to stay in fp16’s representable range. Before the optimizer step, it divides them back down. If any gradient still overflows to inf/nan, it skips the step and reduces the scale factor.

When mixed precision matters

• Large matmuls dominate (transformers, big CNNs) → 2–4x speedup, big memory savings. Always worth it.
• Memory-bound model (OOM in fp32) → often lets you 2x batch size, speeding up training even more.
• Tiny models / RNN cells → maybe no speedup. Cast overhead can eat the wins.

For LLM training it’s not optional — it’s the difference between fitting on your GPU and not.

Common pitfalls

In one breath

• Modern GPUs run 16-bit floats 2–8× faster than fp32 at half the memory; mixed precision gets that without rewriting the model.
• fp16 has 5 exponent + 10 mantissa bits — more precise but a narrow range, so gradients overflow/underflow and it needs a GradScaler.
• bf16 keeps fp32’s 8 exponent bits (same range) with only 7 mantissa bits — less precise but stable, no GradScaler. Just use bf16.
• autocast casts safe ops (matmuls, convs) to 16-bit and keeps precision-sensitive ones (reductions, loss) in fp32 automatically.
• Optimizer state (Adam’s m, v) stays fp32 by design — it’s the real memory tax; and fp8 on H100/Blackwell is the next step beyond bf16.

Quick check