Core ML

Mixed Precision Training

Run the forward and backward pass in a lower-precision format for speed and memory savings, while keeping a full-precision master copy of the weights for accuracy.

The forward and backward pass run in fast fp16, but the weight update applies to a full-precision fp32 master copy, so small updates never get lost to rounding

Why Does This Exist?

Every weight, activation, and gradient in a neural network is normally stored as a 32-bit float. Modern GPUs can run arithmetic on 16-bit floats roughly two to eight times faster and at half the memory, because the hardware has dedicated circuits built specifically for lower-precision math at high throughput. That's a large, close-to-free speedup sitting on the table — except it isn't quite free. Cut every number in the network down to 16 bits carelessly, and small gradients that mattered in 32-bit precision round straight down to zero in 16-bit, and training that would have converged quietly stalls or diverges instead. Mixed precision training is the specific set of techniques that captures most of the speed and memory win without paying that accuracy cost.

Think of It Like This

A rough draft and a master copy

A writer drafts quickly in a rough, low-effort format — no time spent on precise wording, just getting sentences down fast — and that draft is genuinely useful for moving quickly through a chapter. But the writer keeps a separate, carefully maintained master copy of the manuscript, and every real revision gets applied there, in full detail, not to the rough draft.

Small, careful edits made directly to the rough draft — a word changed here, a comma there — are exactly the kind of edit that's easy to lose or garble in a low-effort format. Keeping the precise master copy as the place edits actually land is what protects those small changes, even while the fast rough draft is what's used for the bulk of the day-to-day writing.

How It Actually Works

fp16 versus fp32, and where fp16 falls short

A 32-bit float (fp32) devotes 23 bits to its fraction and can represent numbers across an enormous range, roughly $10^{-38}$ to $10^{38}$ . A 16-bit float (fp16) devotes only 10 bits to its fraction and its range shrinks to roughly $6\times10^{-5}$ to $65{,}504$ . Neural network gradients, especially deep in a large model, routinely fall below that smallest representable fp16 magnitude — cast directly to fp16, they silently become exactly zero, and the parameter they belonged to simply stops updating.

Loss scaling fixes the range problem

Before casting gradients down to fp16, mixed precision training multiplies the loss by a large scaling factor — commonly a power of two, often adjusted automatically during training — so every gradient computed from that scaled loss is proportionally larger too, pushed up into fp16's representable range. After the backward pass, the gradients are divided back down by the same factor before the optimizer step, which recovers the original magnitude exactly, since multiplying and then dividing by the same constant is lossless arithmetic (modulo the rounding fp16 itself introduces along the way).

The fp32 master copy is what protects the actual update

The forward and backward pass run using fp16 copies of the weights and activations, which is where the speed and memory savings come from. But the optimizer's weight update is applied to a separate master copy of the weights kept in full fp32 precision the entire time. After computing the (loss-scaled, then unscaled) gradient in fp16, that gradient gets cast up to fp32 and applied to the fp32 master weights — only afterward does a fresh fp16 copy get cast back down for the next forward pass. This is the detail from the analogy above: the fast format is where the bulk of the compute happens, but the accurate format is where every real update actually lands.

bfloat16 sidesteps the range problem differently

bfloat16 keeps fp32's full exponent range (so it doesn't underflow the way fp16 does) but shrinks the fraction down to 7 bits instead of 10 — trading some precision for avoiding the loss-scaling machinery fp16 needs entirely. Hardware built for large-model training increasingly supports bfloat16 natively, and many large-scale training recipes use it specifically to skip loss scaling as a moving part.

Show Me the Code

A gradient too small to survive a direct cast to fp16, recovered correctly once loss scaling is applied first.

1import numpy as np2
3grad_fp32 = np.float32(1e-8)4direct_cast = np.float16(grad_fp32)5print(f"direct cast to fp16: {direct_cast}  (zero? {direct_cast == 0})")6
7scale = 1024.08scaled_then_cast = np.float16(grad_fp32 * scale)9recovered = np.float32(scaled_then_cast) / scale10print(f"scaled cast to fp16: {scaled_then_cast}  (zero? {scaled_then_cast == 0})")11print(f"recovered gradient:  {recovered:.4e}")12# -> direct cast to fp16: 0.0  (zero? True)13# -> scaled cast to fp16: 1.0251998901367188e-05  (zero? False)14# -> recovered gradient:  1.0012e-08

The direct cast rounds this gradient down to exactly zero — a real update that would have moved this parameter is lost. Scaling by 1024 first keeps the value in fp16's representable range, and dividing back out afterward recovers a value within rounding error of the true gradient.

Watch Out For

Assuming bfloat16 needs the same loss scaling fp16 does

bfloat16's exponent range matches fp32, so the underflow problem loss scaling exists to fix mostly doesn't apply to it — adding loss scaling anyway is a leftover habit from fp16, not a necessity, and it adds a moving part without a matching benefit.

Forgetting the master weights need to stay in fp32

Applying the optimizer update directly to fp16 weights instead of a full-precision master copy reintroduces the exact underflow problem loss scaling was protecting against, just one step later — a small fp32 update, once added into fp16 weights, can round away to nothing. The master copy has to persist in fp32 for the whole run, not just during the backward pass.

The Quick Version

fp16 arithmetic runs faster and uses less memory than fp32, but its narrower range lets small gradients underflow to exactly zero.
Loss scaling multiplies the loss before the backward pass and divides the resulting gradients back down afterward, keeping them inside fp16's range without changing their final value.
A full-precision fp32 master copy of the weights is where every real update is applied — the fp16 copies are only used for the fast forward and backward pass.
bfloat16 keeps fp32's exponent range at lower fraction precision, which is why it often skips loss scaling entirely.
The speedup is close to free on hardware built for it, but only once the master-copy and (if needed) loss-scaling machinery is in place.