Chapter 9 of 16

Mixed precision

FP32 master weights, bf16 and fp16 with loss scaling, FP8, quantised gradients.

Slides for this chapter: Transformer Training 08, Mixed Precision

Every number in this site's trainer is a 64-bit double. Real training uses 16-bit and even 8-bit numbers, because they halve or quarter the memory and the traffic and run several times faster on the matrix units. The cost is precision and range: small numbers fall off the bottom of the format and become zero. Below, all 5,376 gradients of the trained model, computed in your browser, are rounded to each format, and then to FP16 again after multiplying the loss by a growing loss scale.

Loading the animation…

Concept

Formats. A floating-point number has a sign, an exponent (its range: how large and how small it can be) and a mantissa (its precision: how many values per factor of two). FP32 has 8 exponent and 23 mantissa bits. FP16 keeps 5 and 10: plenty of precision but a narrow range, smallest value $2^{-24} \approx 6 \times 10^{-8}$ , largest 65,504.

BF16 ("brain floating point", bfloat16) is 2 bytes like FP16, but keeps FP32's 8-bit exponent and only 7 mantissa bits: the same range as FP32 (in the animation, BF16 sends 0 of the trained model's 5,376 nonzero gradients to zero, FP16 69, FP8 E4M3 1,601 ), and about 2–3 significant decimal digits (Kalamkar et al. 2019). That trade is why most large runs today compute in BF16: it needs no loss scaling. Two bytes per value is also what crosses the network when workers exchange gradients or weight updates in BF16: two bytes per parameter, per worker, per step, before any compression (chapters 11 and 13). The Numerics site shows these formats bit by bit; this chapter's rounding code is vendored from it.

Mixed precision. Micikevicius et al. (2018) made 16-bit training work with three ideas: compute in FP16 but keep an FP32 master copy of the weights that the optimiser updates; accumulate the matrix products in FP32; and loss scaling. The master copy matters because an update is often smaller than the spacing between FP16 or BF16 values near the weight, so adding it to a 16-bit weight does nothing. Loss scaling matters because many gradients are smaller than FP16's smallest value: during the FP16 runs without scaling below, 4.6% of all the gradient values computed were flushed to zero, and a loss scale of $2^{16}$ brings the trained model's gradients in the animation from 69 lost to 0. Multiplying the loss by $2^k$ before the backward pass shifts every gradient up by $k$ binades, into range; the optimiser divides by $2^k$ again in FP32. Dynamic loss scaling starts high, halves the scale and skips the step whenever a gradient overflows to infinity, and doubles it again after a long run of good steps (PyTorch's GradScaler).

Measured on the Decoder's model. The chart trains the default run in each recipe, emulated by rounding every weight's compute copy, every activation a matrix product reads or writes, and every gradient to the format (scripts/training/precision.py), three seeds each. FP32, BF16 with master weights and FP16 with dynamic loss scaling all end within the seeds' spread of the float64 reference (2.277 ± 0.090): FP32 2.316, BF16 2.329, FP16 with loss scaling 2.288. So, at this size, do the two recipes that should hurt: FP16 without loss scaling ends at 2.331 ± 0.027 although 4.6% of its gradient values were flushed to zero, and BF16 with no master copy at 2.296 ± 0.027 although 43% of all its weight updates were rounded away (smaller than half the BF16 spacing at the weight). A 5,376-parameter model trained for 2,000 steps is forgiving: the lost gradients and updates are the smallest ones, and Adam rescales the rest. At scale, with longer runs and smaller learning rates at the end, the same losses accumulate, which is why the master copy and loss scaling are standard.

Emulated on the CPU by rounding in float64 (scripts/training/precision.py); the GPU run of the same comparison is pending (notebook 09). results.md §9a.

FP8. The newest accelerators also multiply in 8-bit floats. Micikevicius et al. (2022) recommend E4M3 (4 exponent, 3 mantissa bits, largest 448) for weights and activations, and E5M2 (more range, less precision) for gradients, with a scaling factor per tensor so that each tensor's largest value lands near the top of the format. In the emulation, only what the matrix products read is FP8; what they write stays BF16, as in those recipes. Without scaling, the FP8 run never leaves chapter 1's plateau (2.957 ± 0.011; 16% of its gradient values were flushed to zero). In these runs, with per-tensor scaling, two of the three seeds train (2.383, 2.965, 2.374): three mantissa bits are coarse, and at this scale whether the run escapes the plateau is a coin toss.

Quantised values, full scale and clipping. Beyond floating point, values can be quantised to a few bits of integer: pick a full scale (the largest magnitude to represent), divide it into $2^{b - 1}-1$ steps each side of zero, and round each value to the nearest step; anything beyond the full scale is clipped to it. With the full scale set to the tensor's largest value nothing clips, but small values round to zero; a smaller full scale keeps more resolution for the small values and clips the large ones. Gradients quantised this way are what compressed distributed training sends between workers (chapter 13), and analogue or optical hardware imposes the same quantisation on whatever passes through it. Measured on the Decoder's model, with every weight gradient quantised per tensor (full scale = its largest value) before the optimiser step: 8 bits 2.286 ± 0.069, 4 bits 2.413 ± 0.209, 2 bits (three levels: −1, 0, 1 times the scale) 2.860 ± 0.182. Adam is forgiving here too, since it normalises each weight's step by the size of its past gradients (Alistarh et al. 2016 analyse quantised SGD).

Run it yourself on a GPU. Notebook 09 (open in Colab) trains a 60-million-parameter model in FP32, BF16, FP16 with and without a GradScaler, and an FP8 emulation, measuring loss, tokens per second and memory (about 40 minutes on an A100). The notebook is one of the planned GPU replications listed under Runs to come; every number on this page comes from the in-browser model.

Maths

A format with $e$ exponent and $m$ mantissa bits, bias $2^{e-1}-1$ , represents normal values $2^{E-\text{bias}}(1 + M/2^m)$ for $1 \le E \le 2^e - 2$ and subnormals $2^{1-\text{bias}}\,M/2^m$ ; the spacing near $|x|$ is $2^{\max(\lfloor\log_2|x|\rfloor,\,1-\text{bias}) - m}$ , so the relative rounding error is at most $2^{-m-1}$ for normal values. Values below half the smallest subnormal round to 0.

Format	Exponent / mantissa bits	Smallest subnormal	Largest	Relative precision
FP32	8 / 23	$2^{-149}$	$3.4 \times 10^{38}$	$2^{-24}$
BF16	8 / 7	$2^{-133}$	$3.4 \times 10^{38}$	$2^{-8}$
FP16	5 / 10	$2^{-24}$	65,504	$2^{-11}$
FP8 E5M2	5 / 2	$2^{-16}$	57,344	$2^{-3}$
FP8 E4M3	4 / 3	$2^{-9}$	448	$2^{-4}$

Loss scaling: $\tilde L = 2^k L \Rightarrow \tilde g = 2^k g$ ; the optimiser uses $\operatorname{round}(\tilde g)/2^k$ . Per-tensor scaling for FP8: $s = 2^{\lfloor \log_2(\max_{\text{fmt}} / \operatorname{amax}(x)) \rfloor}$ , $\hat x = \operatorname{round}(s x)/s$ . Integer quantisation with full scale $F$ and $b$ bits: $\hat x = \Delta \cdot \operatorname{clip}\big(\operatorname{round}(x/\Delta), -n, n\big)$ with $n = 2^{b-1}-1$ , $\Delta = F/n$ .

Code

Rounding a double to a format, src/lib/training/rounding.ts (vendored from the Numerics site; bit-identical to the Python reference on every FP8 value, every midpoint and 400 random values):

// the spacing of the format at |x|: 2^(max(e, emin) − m)
const q = pow2(Math.max(exponentOf(a), f.emin) - f.m);
const t = a / q;
let n = Math.floor(t);
const frac = t - n;
if (frac > 0.5 || (frac === 0.5 && n % 2 === 1)) n += 1;
if (n === 0) return sign; // underflow: rounds to (signed) zero
const v = n * q;
if (v > f.max) {
  if (sat) return sign + f.maxCode;
  return sign + (f.kind === "ieee" ? f.infCode : f.nanCode);
}
return sign + magnitudeCode(v, f);

Comments

Be the first to leave a comment on this section.