Mixed Precision Training

Run the fast math in 16-bit while every real weight change lands in a 32-bit master copy. Small gradients survive because loss scaling lifts them into range first.

A 1e-8 gradient vanishes in fp16 yet survives as 1.02e-05 after scaling by 1024, then lands on fp32 master weights

Why Does This Exist?

Storing every weight, activation and gradient in 32-bit floats is safe but slow. Modern chips run 16-bit math at roughly twice the speed and half the memory because they carry dedicated low-precision units. That speedup sits on the table, except cutting straight to 16-bit breaks training. Gradients deep in a large model routinely fall near $10^{-8}$ , while 16-bit normals start near $6.10 \times 10^{-5}$ . Cast directly, those small gradients become exactly zero and their parameters freeze.

Mixed precision keeps most of the speed without paying that accuracy cost. This page covers the two guards: scale the loss so gradients sit in range, and keep a full-precision master copy where updates actually land. The memory companion that trades batch size for the same budget lives in gradient accumulation. The multi-chip companion lives in distributed training. The inference-time cousin that shrinks finished weights lives in model quantization.

One basic makes the rest click. A float has an exponent for range and a fraction for detail. Cutting bits shrinks one of the two, so fp16 loses range while bfloat16 keeps range and loses detail instead.

Think of It Like This

A rough draft plus a master copy

A writer drafts fast in a rough notebook with thick markers, moving quickly through chapters without fussing over small edits. Alongside it sits a clean master manuscript where every real change lands in fine print.

Tiny edits made only in the rough book, like a comma moved one slot, smudge into nothing because the marker is too thick. Apply those same edits to the master copy and they survive exactly, even while the rough draft does the bulk of the daily writing.

Where it stops: paper never rounds a letter into a different letter, while low-precision math does. The rough draft here can turn a small number into exactly zero, which paper never does. That is why scaling and the master copy are load-bearing, not optional polish.

How It Actually Works

Call fp32 the full format with about $10^{-38}$ to $10^{38}$ range, fp16 the fast format with about $6.10 \times 10^{-5}$ to $65{,}504$ normal range, and bfloat16 the alternate fast format that keeps fp32 range with fewer fraction bits. The forward and backward passes run in the fast format for speed. The optimizer step does not.

Loss scaling fixes range for fp16. Before the backward pass, multiply the loss by a scale $S$ , often 1024 or adapted automatically. Every gradient from that scaled loss is $S$ times larger too, so values that sat below fp16 range move into it. After the backward pass, divide gradients by the same $S$ before the optimizer step. Multiply then divide by one constant recovers the original magnitude up to fp16 rounding along the way.

The fp32 master copy protects the update. Cast fp32 master weights down to fp16 for the fast passes, compute scaled gradients in fp16, unscale, cast up to fp32, then apply the update to the fp32 master. A fresh fp16 copy goes down for the next step. Small updates never squeeze through fp16 storage because storage lives in fp32 the whole run.

Worked example: a gradient that vanishes, then survives

Take true gradient $g = 1 \times 10^{-8}$ . Direct cast to fp16 gives $0.0$ , because $10^{-8}$ sits below what fp16 can hold. Now scale by $S = 1024$ :

1 \times 10^{-8} \times 1024 = 1.024 \times 10^{-5} \approx 1.02 \times 10^{-5}

This scaled value is representable, so the cast keeps about $1.02 \times 10^{-5}$ instead of zero. Divide back out:

1.02 \times 10^{-5} / 1024 \approx 1.00 \times 10^{-8}

The recovered value matches the true gradient within rounding error. That is the diagram headline: one path showing $10^{-8}$ collapsing to zero, one path showing $1.02 \times 10^{-5}$ surviving and returning as $10^{-8}$ onto master weights. Bfloat16 skips most of this machinery because it keeps fp32 range natively and trades fraction detail instead.

Code

The point is the check: rebuild the vanish-then-survive arithmetic from the worked example.

1grad: float = 1e-82scale: float = 1024.03
4direct: float = 0.0 if grad < 6.0e-08 else grad5print(direct)  # -> 0.06
7scaled: float = grad * scale8print(f"{scaled:.2e}")  # -> 1.02e-059
10recovered: float = scaled / scale11print(f"{recovered:.2e}")  # -> 1.00e-08

The printout matches the hand computation: direct path 0.0, scaled value about 1.02e-05, recovered value about 1.00e-08. Change scale to 1 and the middle print drops back toward zero.

Watch Out For

Adding loss scaling on top of bfloat16

Bfloat16 keeps fp32 exponent range, so the underflow that scaling fixes mostly never happens there. Copying an fp16 recipe verbatim adds a moving scale with no matching benefit.

Symptom: scale values swing or overflow checks fire while plain runs stay stable. Fix: default bfloat16 runs to no scaling, and reserve scaled loss for fp16 paths where small gradients truly vanish.

Updating fp16 weights directly

Applying the optimizer step to fp16 weights instead of the fp32 master reintroduces the same rounding one step later. A correct small fp32 update added into fp16 storage can round away to nothing.

Symptom: loss looks healthy early then stalls with many frozen layers late in training. Fix: keep the master in fp32 for the whole run, cast down only for compute, and confirm the update target in code is the master tensor.

The Quick Version

Fp16 math runs faster at half the memory, but normals start near 6.10e-05, so a 1e-8 gradient casts to zero.
Scaling the loss by 1024 lifts that gradient to about 1.02e-05, which survives the cast.
Dividing back by 1024 returns about 1.00e-08, matching the true gradient within rounding.
Every real update applies to fp32 master weights, while fp16 copies handle only fast forward and backward passes.
Bfloat16 keeps fp32 range with fewer fraction bits, so it often skips loss scaling entirely.
Direct fp16 updates or needless scaling on bfloat16 reintroduce the exact stall the method removes.