Mixed-Precision Training

Lesson 2 of 3 in Mixed Precision and Stability.

If 16-bit math is faster and half the size, why not train purely in 16-bit? Because two quiet failures eat the run. Gradients underflow — in FP16, much of the gradient distribution flushes to zero, and the Loss stops improving for no visible reason. And updates vanish — a weight update is learning rate × gradient, often several orders of magnitude smaller than the weight itself, and adding a tiny number to a big one in low precision rounds the sum right back to the big number. Learning silently stalls while the hardware runs at full speed.

The fix, from Micikevicius et al. (2017), is Mixed precision training: do the expensive math in 16-bit, but keep the authoritative numbers in FP32. The recipe has three ingredients — 16-bit forward and backward passes, an FP32 master copy of the weights that receives the updates, and (for FP16 only) loss scaling. Walk one step of it:

One training step in mixed precision

  1. FP32 master weights

    The authoritative parameters. They stay in FP32 for the whole run; everything 16-bit is a disposable working copy.

  2. Cast a 16-bit working copy

    Weights are rounded down to FP16 or BF16 — the copy the matrix units will actually read.

  3. Forward pass in 16-bit

    The matrix multiplies run on the fast 16-bit hardware path. Sensitive ops — softmax, normalization, the loss itself — typically still compute in FP32.

  4. Compute loss, then scale it (FP16 path)

    Multiply the loss by a factor S. By the chain rule every gradient is multiplied by S too — shifting the small ones up above FP16’s underflow line. BF16 runs usually skip this.

  5. Backward pass → 16-bit gradients

    Gradients flow backward in 16-bit, scaled by S on the FP16 path.

  6. Unscale gradients, scan for Inf/NaN

    Divide by S to restore true magnitudes, then check whether anything overflowed during the backward pass.

  7. Overflow detected?

    An Inf or NaN in the gradients means S was too aggressive for this batch — the update would poison the weights.

  8. Skip the update, lower S

    Throw this step’s gradients away and reduce the scale. Losing one batch is cheap; corrupting the weights is not.

  9. Optimizer updates the FP32 master weights

    The tiny update is added in full FP32 precision, so it cannot be rounded away. Periodically, dynamic scaling nudges S back up to reclaim gradient headroom.

  10. Next batch

    Cast a fresh 16-bit copy from the just-updated master weights and repeat.

Each ingredient answers one of the two quiet failures. The FP32 master weights answer vanishing updates: the addition weight + update happens in a format with ~7 digits of precision, so an update a millionth the size of the weight still lands. The 16-bit copy is regenerated from the master every step — rounding never accumulates in the numbers that matter.

Loss scaling answers gradient underflow, and only FP16 needs it. Multiplying the loss by S multiplies every gradient by S (the chain rule is linear), sliding the whole gradient histogram up into FP16’s representable window; dividing by S after the backward pass restores the true magnitudes before the optimizer sees them. The scale factor is itself managed dynamically — grow S while things are clean, halve it and skip the step when an overflow appears. BF16 has FP32’s range, so there is nothing to slide — which is a large part of why it displaced FP16 as the training default. The FP32 master copy stays even with BF16, though: seven mantissa bits round away small updates even more eagerly than FP16’s ten.

You will rarely wire any of this by hand — frameworks ship it as automatic mixed precision, choosing per-op precision, managing the master copy, and running the loss-scaler state machine. But when a run misbehaves, the log lines you will be reading (grad scaler: skipping step, reducing scale) only make sense if you know this loop.

The memory ledger of one parameter

Mixed precision looks like it should halve training memory. It does not — count the ledger. With the standard mixed-precision Adam recipe, each Parameter carries: a 16-bit working weight (2 bytes), a 16-bit gradient (2 bytes), and then the FP32 trio — master weight (4), momentum (4), variance (4). That is 16 bytes per parameter of persistent training state, the accounting Rajbhandari et al. (2019) use to motivate ZeRO: a 7B-parameter model carries ~112 GB of state before a single activation is stored, and a 70B model over a terabyte.

Two consequences. First, the win from 16-bit compute is mostly speed and activation memory, not optimizer memory — the FP32 states dominate the ledger, which is exactly why the previous module’s sharding machinery (ZeRO, FSDP) exists. Second, this same ledger is what a checkpoint must serialize if a run is to resume exactly where it stopped — weights alone are not enough, because Adam’s momentum and variance are part of the run’s state. Hold that thought for the next lesson, where checkpoints become the main character.

Interactive checkpoint quiz (1 questions) — open this page in a browser to take it.