Mixed precision training: Visual walkthrough — AI-ML deep dive

We are unpacking the central trick from Mixed precision training — loss scaling — and the FP32 master-weight partner that makes it safe. Prerequisites we lean on: 1.3.07-Numerical-stability (why small numbers vanish), 4.2.03-Tensor-cores (why we want FP16 at all), and 3.3.05-Learning-rate-schedules (where $η$ comes from).

Step 1 — What a floating-point number actually is

WHAT. A number in a computer is stored as three pieces: a sign (plus or minus), an exponent (which power-of-two "bucket" the number lives in), and a mantissa (the fine detail inside that bucket). Think of a ruler: the exponent picks which segment of the ruler you're standing on, the mantissa picks the tick within that segment.

WHY. Everything about mixed precision comes from one fact: FP16 has fewer bits than FP32, so both its reach (biggest/smallest number) and its resolution (smallest gap between neighbours) shrink. We have to see those two shrinkages before anything else makes sense.

PICTURE. The number line below is drawn in log-space. Blue marks the range FP32 can reach; orange marks the much narrower range of FP16. Notice orange stops on the small-number side well before the interesting gradients live.

Step 2 — Where a gradient comes from, and why some are tiny

WHAT. A gradient $g = \frac{\partial L}{\partial w}$ is "how much the loss $L$ changes when we nudge weight $w$ ." The backward pass builds it by the chain rule, multiplying one small factor per layer:

g = top layer \frac{\partial L}{\partial a _{n}} \cdot \times \frac{\partial a _{n}}{\partial a _{n - 1}} \dots our weight \frac{\partial a _{1}}{\partial w}

$L$ — the single loss number at the top of the network.
$a_{k}$ — the activation (output) of layer $k$ .
each fraction — "how layer $k$ 's output moves when its input moves."

WHY. Each factor is often less than $1$ . Multiply many sub- $1$ numbers and the product shrinks exponentially with depth. That is why early-layer gradients land around $1 0^{- 7}$ — right in the danger zone we drew in Step 1.

PICTURE. Watch the green gradient value slide left (smaller) as we move from the last layer toward the first. The orange FP16 floor is a cliff; gradients that fall off it become exactly $0$ .

Step 3 — The trick: multiply the loss, and let the chain rule spread it

WHAT. Pick a scalar $S ≫ 1$ (say $S = 1024$ ) and train on a scaled loss $L^{'} = S \cdot L$ instead of $L$ .

WHY here, why this tool? We want to move every gradient up into FP16's range. Why scale the loss and not each gradient? Because the chain rule is linear in the top factor:

\frac{\partial L ^{'}}{\partial w} = \frac{\partial ( S L )}{\partial w} = S \cdot \frac{\partial L}{\partial w} = S \cdot g

$L^{'} = S L$ — one multiply, at the very top of the graph.
$S$ pulls straight through the derivative because it's a constant — every gradient in every layer comes out multiplied by $S$ automatically, for the price of a single multiplication.

That is the elegance: one nudge at the top, and the whole backward pass lifts together.

PICTURE. Same gradients as Step 2, but each is now $\times S$ . The green dots slide right, clearing the orange floor.

Step 4 — Unscale before you touch the weights

WHAT. After the backward pass we hold $g_{scaled} = S \cdot g$ . Before updating anything we divide it back:

g_{true} = \frac{g _{scaled}}{S}

WHY. The scaling was a temporary lift to survive FP16. The optimizer must see the real gradient, or the effective learning rate would silently be $S$ times too big. We do this division in FP32, where the number now sits comfortably.

PICTURE. A round-trip: the gradient rides up by $\times S$ through the FP16 backward pass (orange arc, safe above the floor), then comes back down by $\div S$ into FP32 (blue arc) landing on its true value.

Step 5 — Why we still need FP32 master weights

WHAT. Even with a healthy gradient, the update $η g$ can be far tinier than the weight itself. Keep the authoritative weights in a FP32 master copy; only cast a FP16 copy for the fast forward/backward math.

WHY. The update rule is

w_{t + 1} = w_{t} - η g_{t}

$w_{t}$ — current weight, often near $1$ (so FP16 resolution near it is $ϵ_{16} \approx 1 0^{- 3}$ ).
$η$ (the Greek letter "eta") — the learning rate: a small fixed number that sets how big a step we take, e.g. $η = 1 0^{- 4}$ . (It comes from your learning-rate schedule.)
$g_{t}$ — the gradient at step $t$ .
$η g_{t}$ — the actual step; with $g_{t} = 1 0^{- 3}$ this is $1 0^{- 7}$ .

Adding $1 0^{- 7}$ to a number near $0.5$ in FP16 rounds right back to $0.5$ : the step is smaller than the gap between FP16 neighbours. This is a different failure from Step 2 — not the gradient underflowing, but the sum swallowing the update.

PICTURE. Left panel: in FP16 the update arrow is shorter than the tick spacing — $w$ snaps back. Right panel: in FP32 the ticks are dense enough that the same update actually moves $w$ .

Step 6 — The full loop, and the edge cases

WHAT. Chain the pieces into one training step:

WHY / edge cases. Two things can go wrong at the extremes, and both have fixes:

Overflow (too much scale). If $S$ is too big, a large gradient times $S$ exceeds FP16's ceiling ( $\approx 65504$ ) and becomes inf/nan. Fix: dynamic loss scaling — when you see inf/nan, skip this step and halve $S$ . If many steps pass cleanly (e.g. 2000), double $S$ to reclaim headroom. $S$ auto-tunes to each model. (Halving and doubling keep $S$ a power of two, preserving the exact round-trip from Step 4.)
Zero / degenerate gradient. A genuinely-zero gradient ( $g = 0$ ) stays $0$ under any scale ( $S \cdot 0 = 0$ ) — scaling can't invent signal that isn't there, and that's correct: those weights shouldn't move.
BF16 escape hatch. Batch normalization statistics and softmax sums can overflow FP16's small range. BF16 has FP32's full range (just coarser mantissa), so on hardware that supports it you often skip loss scaling entirely — the floor problem never appears.

PICTURE. The dynamic-scaling controller as a sawtooth: $S$ climbs by doubling while stable, drops by halving on each overflow, hovering just under the danger line.

The one-picture summary

One gradient's journey: born tiny (below the FP16 floor) → lifted by $\times S$ into the safe FP16 band → carried through the backward pass → dropped by $\div S$ in FP32 → added to an FP32 master weight where even a $1 0^{- 7}$ step counts. Fast math in FP16, careful bookkeeping in FP32.

Recall Retell the whole thing in plain words (Feynman check)

FP16 is a short ruler: it can't measure very tiny things (they read as zero) and its ticks are far apart (tiny changes disappear). Gradients — especially in early layers — are exactly those tiny things. So before the backward pass I multiply the loss by a big number $S$ ; because the chain rule carries constants straight through, every gradient comes out $S$ times bigger and now fits on the ruler. I pick $S$ as a power of two so the multiply and later divide are exact exponent shifts. Right before I change the weights I divide by $S$ to get the honest gradient back, and I do that in FP32, the long precise ruler. I also keep the real weights in FP32 the whole time, because a weight update can be smaller than an FP16 tick and would otherwise vanish. If $S$ ever gets so big that gradients overflow to infinity, I halve it and try again; if things stay calm, I double it to keep the most headroom. BF16 skips the lifting trick because it already has FP32's reach.

Recall

Why scale the loss instead of each layer's gradient? ::: One multiply at the top; the chain rule multiplies every downstream gradient by that same constant automatically. What are the two different FP16 failures? ::: (1) small gradients underflow to zero below the FP16 floor; (2) small weight updates get swallowed because they're below the $\approx 1 0^{- 3}$ spacing near the weight. What is the smallest normal FP16 value, and the smallest subnormal? ::: Normal $n_{16} = 2^{- 14} \approx 6 \times 1 0^{- 5}$ ; subnormal $d_{16} = 2^{- 24} \approx 6 \times 1 0^{- 8}$ . Why must the loss scale $S$ be a power of two? ::: Multiplying/dividing by $2^{k}$ only shifts the exponent, leaving the mantissa untouched — the scale/unscale round-trip is then exact. What does dynamic loss scaling do on an overflow? ::: Skip the step and halve $S$ ; double $S$ after a long stable run. Why does BF16 often need no loss scaling? ::: It shares FP32's exponent range, so gradients don't underflow — only precision is reduced.

Related: Hinglish version · 3.3.10-Gradient-checkpointing · 5.1.02-Model-parallelism