Mixed Precision and Loss Scaling

Half precision halves the memory and doubles the throughput, and silently rounds small gradients to zero. Watch it happen, then fix it.

Overview

Why bother

Modern accelerators run half-precision matrix multiplies several times faster than single-precision ones, and half-precision tensors take half the memory. That means bigger batches, bigger models, and a real speed-up on the same hardware.

The catch is the format. fp32 has 8 exponent bits and 23 mantissa bits. fp16 has 5 and 10, giving a smallest normal value around 10⁻⁵ and a maximum of 65,504.

Both ends of that range are a problem, and the small end is the one that bites.

Mixed Precision and Loss Scaling

This module needs JavaScript: the numbers are computed in the page rather than recorded.

Worth knowing

fp16 has 5 exponent bits and 10 mantissa bits. Its smallest representable value is about 6×10−8.
Gradients are routinely smaller than that. They round to exactly zero, and the parameter never moves.
Loss scaling multiplies the loss before the backward pass, shifting every gradient up into the representable range.
The optimiser unscales before stepping, so the update is unchanged. The scaling exists only to survive the trip through fp16.

Mixed Precision and Loss Scaling

Training in half the bits, the failure that causes, and the one-line trick that removes it.

Watch the gradients disappear

The histogram is gradient magnitudes spread over several orders of magnitude, which is what a real network produces. The two marked lines are fp16's floor and ceiling.

With the loss scale at 2⁰ — that is, no scaling — the readout reports how many of the 3,000 gradients round to zero. Those parameters receive no update at all. Not a small update; none.

This is the failure that made early fp16 training unreliable. The model trains, the loss goes down, and a fraction of the network is silently frozen.

Loss scaling

The fix is almost embarrassingly simple. Multiply the loss by a large constant before the backward pass. By the chain rule every gradient is multiplied by the same constant, so the whole distribution shifts up into the representable range. Then divide by that constant before the optimiser steps.

Drag the loss scale upward and watch the underflow count fall to zero.

The update is mathematically unchanged. The scaling exists purely so the numbers survive the trip through fp16 — it is a change of units, not a change of algorithm.

Push it too far and the readout starts reporting overflow: gradients that have become infinite. That is a failure too, and it is where the automatic version comes from.

Dynamic loss scaling

Nobody tunes the constant by hand. The standard algorithm starts high, multiplies by two whenever a stretch of steps completes without overflow, and halves it and skips the step whenever an infinity or NaN appears.

Skipping is the important part: an overflowed gradient is not clipped or repaired, it is discarded, because there is no way to recover what it should have been. A few skipped steps early in training cost nothing.

This is what torch.cuda.amp.GradScaler does, and it is why mixed precision is now a decorator rather than a project.

What stays in fp32

"Mixed" is the operative word. In a standard setup:

Weights are kept in an fp32 master copy. Updates are frequently far smaller than the weights themselves, and adding a tiny number to a large one in fp16 rounds to no change at all.

Matrix multiplies and convolutions run in fp16, where the speed is.

Reductions — sums, means, softmax, normalisation statistics — accumulate in fp32, because adding many small numbers is exactly where limited mantissa bits go wrong.

Loss is computed in fp32.

bfloat16

The newer alternative has 8 exponent bits and 7 mantissa bits — the same range as fp32, with less precision.

That trade removes this entire page. Range was the problem; precision, for gradients, mostly was not. bfloat16 needs no loss scaling, which is why it is the default on hardware that supports it, and why large-model training largely stopped talking about GradScaler.

Where it goes wrong

fp16 without loss scaling. Silent underflow, and a partly frozen model.

A fixed scale. Use the dynamic version.

Reductions in fp16. Softmax and layer norm in half precision produce mysterious NaNs.

Assuming a speed-up. It comes from the tensor cores. Layers that are memory-bound rather than compute-bound gain little.

Check yourself

0 of 3

Answer without scrolling back up.

  1. What does loss scaling change about the parameter update?

  2. What does dynamic loss scaling do when a gradient overflows?

  3. Why does bfloat16 not need loss scaling?

Cheat sheet

Mixed Precision and Loss Scaling

Modern accelerators run half-precision matrix multiplies several times faster than single-precision ones, and half-precision tensors take half the memory. That means bigger batches, bigger models, and a real speed-up on the same hardware.

DEEP LEARNING · vizlearn.in/deep_learning/mixed_precision_training.html

About the author

Ashish Jangra builds and maintains VizLearn. Every module here is written and the visualisation behind it hand-built, so the numbers in a readout come from the same code that draws the picture. Corrections are genuinely welcome and get priority over everything else — if a page states something wrong, or an animation misrepresents what the algorithm does, get in touch.