Mixed precision training: Exercises — AI-ML deep dive

Before we start, three quick reminders so every symbol below is earned:

Definition The words we use constantly

Normal vs subnormal. A floating-point number is stored as $\pm m \times 2^{e}$ . When the exponent $e$ hits its smallest allowed value, the format can squeeze out a few extra tiny numbers by letting the leading mantissa bit be $0$ — these are subnormals. For FP16 the smallest normal positive number is $\approx 6.1 \times 1 0^{- 5}$ ; below that lie subnormals reaching down to $\approx 6 \times 1 0^{- 8}$ (the smallest subnormal). Crucially, most GPU/Tensor-Core paths run in flush-to-zero (FTZ) mode: subnormals are treated as $0$ . So for training we take the effective underflow cliff to be the smallest normal, $\approx 6.1 \times 1 0^{- 5}$ — anything below that flushes to $0$ . This is the number every underflow argument below uses.
Precision $ϵ$ (machine epsilon, the ULP convention). $ϵ$ is the gap between $1.0$ and the next representable number — one ULP ("unit in the last place") at magnitude $1$ . For FP16, $ϵ = 2^{- 10} \approx 9.77 \times 1 0^{- 4} \approx 1 0^{- 3}$ . Rounding to nearest loses any addition smaller than half a ULP, i.e. below $\frac{1}{2} ϵ ∣ x ∣$ ; we use the full-ULP $ϵ$ as the round-number threshold. Near a value $x$ , neighbours are spaced by $ϵ \cdot 2^{⌊ l o g_{2} ∣ x ∣ ⌋}$ .
Loss scale $S$ — one positive number you multiply the loss by before backprop. By the chain rule it multiplies every gradient by $S$ , sliding the whole gradient histogram to the right (bigger), away from the underflow cliff.

Sign symmetry. Every bound below is stated for magnitudes $∣ g ∣$ . Floating point has a dedicated sign bit, so the set of representable magnitudes is identical for positive and negative numbers. A gradient of $- 4 \times 1 0^{- 8}$ underflows to $- 0$ exactly as $+ 4 \times 1 0^{- 8}$ underflows to $+ 0$ ; a gradient of $- 9 \times 1 0^{4}$ overflows to $- \infty$ exactly as $+ 9 \times 1 0^{4}$ overflows to $+ \infty$ . So we always reason about $∣ g ∣$ and the sign just rides along.

The blueprint below is not decoration — it is the map for every underflow/overflow argument on this page. Read it as instructed.

Level 1 — Recognition

L1.1 — Name the bits

Fill in the blanks: FP16 has 1 sign bit, 5 exponent bits and 10 mantissa bits. BF16 keeps FP32's 8 exponent bits but only 7 mantissa bits.

Recall Solution

FP16 = 1 + 5 + 10 = 16 bits. BF16 = 1 + 8 + 7 = 16 bits. Both are 16 bits total, but they spend those bits differently: FP16 buys precision (10 mantissa bits), BF16 buys range (8 exponent bits, matching FP32). That is why BF16 usually needs no loss scaling — its range already covers tiny gradients. Reveal: Why does BF16 rarely need loss scaling? ::: Its 8 exponent bits give it FP32's dynamic range, so gradients almost never underflow.

L1.2 — Which copy is which?

In mixed precision, which weight copy is used for the forward pass, and which one accumulates the update?

Recall Solution

Forward/backward pass uses the FP16 copy (fast on Tensor Cores).
The optimizer update happens on the FP32 master copy, which is then cast back to FP16 for the next step. The reason is that updates are tiny and cumulative; only FP32 has enough precision to not lose them (proven in L2.1).

Level 2 — Application

L2.1 — The vanishing update

A weight sits at $w = 1.0$ . Learning rate $η = 1 0^{- 3}$ , gradient $g = 2 \times 1 0^{- 4}$ . FP16 precision is $ϵ = 2^{- 10} \approx 1 0^{- 3}$ . (a) What is the true update $η g$ ? (b) If you update the weight in FP16, what value do you get, and why? (c) If you update the FP32 master, what value do you get, and why is that offset representable there?

Recall Solution

(a) $η g = 1 0^{- 3} \times 2 \times 1 0^{- 4} = 2 \times 1 0^{- 7}$ . (b) The new weight should be $1.0 - 2 \times 1 0^{- 7} = 0.9999998$ . FP16 neighbours around $1.0$ are spaced by one ULP $= ϵ = 2^{- 10} \approx 9.8 \times 1 0^{- 4}$ . Our offset $2 \times 1 0^{- 7}$ is far below half that spacing, so round-to-nearest snaps back: the nearest representable value to $0.9999998$ is still $1.0$ . The update rounds away to zero — the weight never moves. (c) FP32 has a 23-bit mantissa, so just below $1.0$ its numbers are spaced by one ULP $= 2^{- 23} \approx 1.19 \times 1 0^{- 7}$ (the exponent bucket for values in $[0.5, 1)$ ). Our offset $2 \times 1 0^{- 7}$ is larger than that ULP ( $2 \times 1 0^{- 7} > 1.19 \times 1 0^{- 7}$ ), so $1.0 - 2 \times 1 0^{- 7}$ lands on — or rounds to — a genuine FP32 grid point distinct from $1.0$ . That is why FP32 keeps the update while FP16 loses it: the offset clears FP32's ULP but not FP16's. After 5000 such steps the master has drifted by $5000 \times 2 \times 1 0^{- 7} = 1 0^{- 3}$ , now large enough to change even the FP16 copy. This is exactly why the master weight exists.

L2.2 — Sizing the loss scale

Early-layer gradients cluster near $g = 4 \times 1 0^{- 6}$ . FP16's smallest normal value (the FTZ cliff) is $\approx 6.1 \times 1 0^{- 5}$ . (a) Without scaling, what happens to these gradients under flush-to-zero? (b) Choose the smallest power of two $S$ that lifts $∣ g ∣$ to a comfortable $\geq 1 0^{- 3}$ . Compute the scaled gradient. Confirm the same works for $g = - 4 \times 1 0^{- 6}$ .

Recall Solution

(a) $4 \times 1 0^{- 6} < 6.1 \times 1 0^{- 5}$ , so the magnitude falls below the smallest normal. On FTZ hardware the subnormal region is treated as zero, so the gradient flushes to $0$ — that neuron gets no learning signal. (b) We want $S \cdot 4 \times 1 0^{- 6} \geq 1 0^{- 3}$ , i.e. $S \geq \frac{1 0 ^{- 3}}{4 \times 1 0 ^{- 6}} = 250$ . The smallest power of two $\geq 250$ is $2^{8} = 256$ . Scaled gradient: $256 \times 4 \times 1 0^{- 6} = 1.024 \times 1 0^{- 3}$ — comfortably above both the cliff and our $1 0^{- 3}$ target, so it survives in FP16. For the negative case, $256 \times (- 4 \times 1 0^{- 6}) = - 1.024 \times 1 0^{- 3}$ : identical magnitude, opposite sign — the sign bit carries it, no special handling. After backprop we divide by $S$ to recover the true $\pm 4 \times 1 0^{- 6}$ on the FP32 master.

L2.3 — Overflow check at the top

Loss scale $S = 65536$ . The largest gradient magnitude in the batch (before scaling) is $∣ g_{m a x} ∣ = 1.2 \times 1 0^{- 3}$ . FP16's max is $65504$ . Does scaling this gradient overflow, and does the sign matter?

Recall Solution

Scaled magnitude $= 65536 \times 1.2 \times 1 0^{- 3} = 78.6432$ . That is far below $65504$ , so no overflow — and this holds for both $+ 78.6432$ and $- 78.6432$ since the representable magnitudes are sign-symmetric. Loss scaling only overflows when the largest magnitude times $S$ exceeds $65504$ ; here we have huge headroom, meaning $S$ could be pushed higher. This is the logic dynamic loss scaling uses to increase $S$ when it's been safe for many steps.

Level 3 — Analysis

L3.1 — Memory accounting

A model has $P = 100$ million parameters, trained with the Adam optimizer (2 FP32 state tensors per parameter: momentum + variance). (a) FP32-only footprint for weights + gradients + Adam state (ignore activations). (b) Mixed-precision footprint: FP16 weights + FP16 gradients + FP32 master weights + 2× FP32 Adam state. (c) Which scheme uses more memory for this non-activation part, and by how much?

Recall Solution

Bytes: FP32 = 4, FP16 = 2. $P = 1 0^{8}$ . (a) Weights $4 P$ + grads $4 P$ + Adam $2 \times 4 P$ = $16 P = 16 \times 1 0^{8}$ B $= 1600$ MB. (b) FP16 weights $2 P$ + FP16 grads $2 P$ + FP32 master $4 P$ + Adam $8 P$ = $16 P = 1600$ MB. (c) They are identical (1600 MB) for the weight/gradient/optimizer part! Mixed precision does not save memory on this bookkeeping — it even adds the FP32 master copy. The real savings come from activations (halved in FP16), which dominate deep-network memory. This is the counter-intuitive truth the parent's Example 3 hides behind its totals.

L3.2 — Where does the speedup actually live?

A training step spends 90% of its time in matrix multiplies (GEMM) and 10% in elementwise ops (which get no Tensor Core speedup). GEMM runs 4× faster in FP16. (a) By Amdahl's law, what is the overall speedup? (b) What is the theoretical ceiling if GEMM became infinitely fast?

Recall Solution

Let old total time $= 1$ : GEMM $= 0.9$ , elementwise $= 0.1$ . (a) New GEMM time $= 0.9/4 = 0.225$ . New total $= 0.225 + 0.1 = 0.325$ . Speedup $= 1/0.325 \approx 3.08 \times$ . (b) As GEMM $\to 0$ , total $\to 0.1$ , speedup $\to 1/0.1 = 10 \times$ . The un-accelerated 10% caps you at $10 \times$ no matter how fast Tensor Cores get — a classic Amdahl ceiling. See 4.2.03-Tensor-cores for what stays FP32.

L3.3 — Reading a dynamic-loss-scale log

A dynamic scaler starts at $S = 2^{16} = 65536$ , halves on overflow, and doubles after $2000$ clean steps. The log shows overflow at steps 3, 5, and 8, then no overflow for a long time. (a) What is $S$ right after step 8's overflow? (b) At step 2010 (2002 clean steps since step 8), what is $S$ ?

Recall Solution

(a) Start $2^{16}$ . Three halvings (steps 3, 5, 8): $2^{16} \to 2^{15} \to 2^{14} \to 2^{13} = 8192$ . (b) After 2000 clean steps since step 8, the scaler doubles once: $2^{13} \to 2^{14} = 16384$ . (Only one doubling — 2002 clean steps crosses the 2000 threshold exactly once.) $S = 16384$ .

Level 4 — Synthesis

L4.1 — Combine with gradient accumulation

You want an effective batch of 512 but only 128 fits per step, so you accumulate over 4 micro-batches (3.3.08-Gradient-accumulation). Loss scale $S = 1024$ . (a) Where in the pipeline do you apply $S$ — once per micro-batch, or once total? (b) After accumulating 4 scaled gradients, by what factor must you divide before the optimizer step to get the mean true gradient over 512 examples?

Recall Solution

(a) Apply $S$ to each micro-batch's loss before its backward pass — because each backward must keep its gradients out of the underflow zone. The scaling is linear, so scaled gradients simply add up: $\sum_{i} S g_{i} = S \sum_{i} g_{i}$ . (b) Accumulated value $= S \cdot \sum_{i = 1}^{4} g_{i}$ . To get the mean true gradient you divide by $S$ (unscale) and by the number of micro-batches $4$ (average): total divisor $= S \times 4 = 4096$ . If your per-micro-batch loss is already the mean over its 128 examples, dividing the sum of 4 by 4 gives the mean over 512. Order doesn't matter — just do both divisions before optimizer.step.

L4.2 — Batch norm in FP16?

2.4.06-Batch-normalization computes a mean and a variance across the batch, then normalises. Why is it standard to keep BatchNorm's statistics accumulation in FP32 even in a mixed-precision model?

Recall Solution

Variance is $E [x^{2}] - (E [x])^{2}$ — a difference of two large, nearly-equal numbers. In FP16 (precision $1 0^{- 3}$ ) this catastrophic cancellation can produce garbage or even a tiny negative variance, and the following $1/ var$ then blows up — a numerical-stability disaster. Summing many terms also loses small contributions once the running sum is large (the L2.1 phenomenon again). So the reductions (sum, mean, var) run in FP32; the cheap elementwise scale-and-shift can stay FP16. This is why frameworks whitelist/blacklist specific ops rather than casting the whole model.

Level 5 — Mastery

L5.1 — Design a scaling window

Profiling one model shows: smallest gradient magnitude $∣ g_{m i n} ∣ = 1 \times 1 0^{- 7}$ , largest $∣ g_{m a x} ∣ = 8 \times 1 0^{- 2}$ . FP16 usable (normal) range: $[6.1 \times 1 0^{- 5}, 65504]$ . You want every gradient magnitude representable after scaling by a single power-of-two $S$ . (a) Lower bound on $S$ so $∣ g_{m i n} ∣$ clears the FTZ cliff. (b) Upper bound on $S$ so $∣ g_{m a x} ∣$ stays under the ceiling. (c) Is there a valid power-of-two $S$ ? If so, give one.

Recall Solution

(a) Need $S \cdot ∣ g_{m i n} ∣ \geq 6.1 \times 1 0^{- 5}$ : $S \geq \frac{6.1 \times 1 0 ^{- 5}}{1 0 ^{- 7}} = 610$ . (b) Need $S \cdot ∣ g_{m a x} ∣ \leq 65504$ : $S \leq \frac{65504}{8 \times 1 0 ^{- 2}} = 818800$ . (c) Valid window $610 \leq S \leq 818800$ — enormous. Powers of two inside: $2^{10} = 1024$ up to $2^{19} = 524288$ . A safe choice is $S = 2^{14} = 16384$ (check: $16384 \times 1 0^{- 7} = 1.64 \times 1 0^{- 3} \geq 6.1 \times 1 0^{- 5}$ ✓; $16384 \times 8 \times 1 0^{- 2} = 1310.7 \leq 65504$ ✓). Any $S$ in that range works for both signs of every gradient; dynamic scaling would settle somewhere here on its own.

The picture below shows exactly this: the gradient histogram as a bar, the FP16 window as a lane, and how multiplying by $S$ slides the bar into the lane.

L5.2 — When the window is empty

Suppose instead $∣ g_{m i n} ∣ = 1 \times 1 0^{- 7}$ and $∣ g_{m a x} ∣ = 2 \times 1 0^{4}$ . Repeat the bound calculation. What must you do?

Recall Solution

Lower bound: $S \geq \frac{6.1 \times 1 0 ^{- 5}}{1 0 ^{- 7}} = 610$ (as before). Upper bound: $S \leq \frac{65504}{2 \times 1 0 ^{4}} = 3.2752$ . Now we need $S \geq 610$ and $S \leq 3.2752$ — no valid $S$ exists. The gradient dynamic range itself, $∣ g_{m a x} ∣/∣ g_{m i n} ∣ = 2 \times 1 0^{11}$ , is comparable to FP16's whole usable span ( $65504/6.1 \times 1 0^{- 5} \approx 1.07 \times 1 0^{9}$ ) — in fact it exceeds it, so no single rigid slide fits both ends. Fixes: (1) switch to BF16, whose FP32-matching range needs no scaling; (2) use per-tensor / per-layer loss scales so early and late layers get different $S$ ; (3) reduce the huge $∣ g_{m a x} ∣$ at its source with gradient clipping or normalization. This is precisely the scenario that motivated BF16 hardware.

L5.3 — End-to-end sanity number

Put L2.1's logic on a schedule. With master-weight FP32 accumulation, $η$ from a cosine schedule (3.3.05-Learning-rate-schedules) currently at $η = 5 \times 1 0^{- 4}$ , gradient $g = 3 \times 1 0^{- 4}$ : compute the single-step update and confirm it is below FP16 precision (so the master is doing real work).

Recall Solution

Update $= η g = 5 \times 1 0^{- 4} \times 3 \times 1 0^{- 4} = 1.5 \times 1 0^{- 7}$ . FP16 precision (one ULP at magnitude $1$ ) $\approx 1 0^{- 3}$ , and $1.5 \times 1 0^{- 7} ≪ 1 0^{- 3}$ , so a direct FP16 update near a weight of order $1$ would round away entirely. The FP32 master (ULP $\approx 1.19 \times 1 0^{- 7}$ just below $1.0$ ) just barely keeps a $1.5 \times 1 0^{- 7}$ change — this is the razor's edge the whole technique balances on. Over a cosine schedule the update shrinks further as $η$ decays, making the master copy more essential late in training, not less.

Recall Quick self-audit before you leave

Master weights are FP32 because updates are smaller than FP16 precision (one ULP ≈ 1e-3) and would round away.
Loss scaling fixes underflow of tiny-magnitude gradients below the FTZ cliff (6.1e-5), not overflow.
Bigger $S$ risks overflow of the largest-magnitude gradients (over 65504).
Underflow/overflow depend only on magnitude; the sign bit rides along symmetrically.
The biggest memory saving in mixed precision comes from activations, not weights.
BF16 needs no loss scale because it has FP32's dynamic range.