Modern GPUs run 16-bit matmuls on tensor cores roughly twice as fast as fp32, and a 16-bit tensor is half the bytes. So training in fp16/bf16 nearly halves activation memory and doubles throughput — for free, it seems. Except a 16-bit float doesn't only lose precision. fp16 loses range, and that's the failure that makes naive 16-bit training NaN within a few steps. I built a live tool that lays out the formats bit-by-bit and rescues underflowing gradients with a slider; here's the whole story.
The problem: fp16's exponent is tiny
Every float is ±1.mantissa × 2^exponent. The exponent bits set the range; the mantissa bits set the precision. fp16 spends its 16 bits as 5 exponent + 10 mantissa, so its smallest subnormal is 2⁻²⁴ ≈ 6e-8. Real gradients routinely fall below that and underflow to exactly zero — a silent dead gradient, and learning stalls. (The top end matters too: above 65504 fp16 overflows to Inf.)
Loss scaling borrows exponent headroom
The fix is almost too simple. Multiply the loss by a big constant S before .backward(). By the chain rule that scales every gradient by S, lifting the tiny ones back into fp16's representable window. Then unscale — divide the grads by S, in fp32, before the optimizer step — so the update is mathematically identical. You only borrowed range, temporarily.
S = 2**14
(loss * S).backward() # grads are now S× bigger -> fit in fp16
for p, m in zip(model.parameters(), master):
m.grad = p.grad.float() / S # UNSCALE in fp32 before stepping
opt.step()
# math is identical: (loss*S).backward()/S == loss.backward()
Two more pieces make it stable
First, an fp32 master copy of the weights. The tiny w -= lr·g update rounds to nothing if you add it straight into an fp16 weight, so the authoritative weights live in fp32 and the step happens there; you copy fp32 → fp16 only for the fast forward/backward.
Second, dynamic loss scaling, because one fixed S is fragile — too small still underflows, too big overflows. So auto-tune it: push S up until a gradient overflows to Inf/NaN, skip that step, back S off, and grow it again after a clean streak. It converges on the largest safe S on its own.
grads = [p.grad.float() / S for p in model.parameters()]
if any(not torch.isfinite(g).all() for g in grads):
S = max(1, S / 2) # overflow -> back off, SKIP step
else:
opt.step(); good += 1
if good % 2000 == 0: S *= 2 # been safe a while -> grow S
bf16 sidesteps the whole thing
bfloat16 splits its 16 bits differently: 8 exponent + 7 mantissa. That's fp32's full exponent, so bf16 has the same range and gradients simply don't underflow — you can drop the loss scaler entirely. The cost is a coarse 7-bit mantissa, so it rounds harder, but that's usually fine. This is the default for most modern LLM training.
with autocast(dtype=torch.bfloat16): # same range as fp32
loss = criterion(model(xb), yb)
loss.backward() # NO loss scaling required
opt.step()
In practice you write three lines
You never hand-roll the master copy and the scaler. autocast runs each op in the right precision — matmuls and convs in 16-bit, softmax/layernorm/loss kept in fp32 — and GradScaler handles dynamic scaling, unscaling, and skip-on-overflow.
scaler = GradScaler()
with autocast(dtype=torch.float16):
loss = criterion(model(xb), yb)
scaler.scale(loss).backward() # loss * S, then backward
scaler.step(opt) # unscale + skip if Inf/NaN
scaler.update() # adjust S up/down
Mixed precision is one of three orthogonal efficiency pillars that stack: gradient checkpointing trades compute for memory, accumulation trades time for batch size, and mixed precision changes the number format itself for ~2× speed and ~½ memory. Drag the probe down to 3e-8 and watch fp16 flip to underflow while bf16 holds, then slide the loss scale and watch the dead gradients climb back into range:
https://dev48v.infy.uk/dl/day50-mixed-precision-training.html
Top comments (0)