Transformer Training — Presentation 08

Mixed Precision

Real training computes in 16 and even 8 bits. Exponent against mantissa, BF16 against FP16, FP32 master weights and loss scaling, FP8 with per-tensor scaling, and gradients quantised to a few bits, each emulated on the model and measured.

FP32 BF16 FP16 Loss scaling Master weights FP8 E4M3/E5M2 Quantisation
Format → Round → Scale → Accumulate → Update
00

Topics We’ll Cover

Live companion and sources

This deck follows chapter 9, mixed precision, where all 5,376 gradients of the trained model are rounded in your browser. Measured numbers: results.md §9a–9b. The formats bit by bit: Numerics Explained (the site’s rounding code is vendored from it and matches it on 202,090 values).

01

Why Train in Fewer Bits

Every number in the 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 the matrix units run several times faster on them.

Memory

Weights, gradients and activations at 2 bytes instead of 4 (deck 09).

Traffic

Two bytes per parameter per worker per step is what crosses the network when workers exchange BF16 gradients (deck 10, deck 12).

The cost

Precision and range: small numbers fall off the bottom of the format and become zero; large ones overflow.

02

Formats: Exponent Against Mantissa

A floating-point number has a sign, an exponent (its range) and a mantissa (its precision: how many values per factor of two). With \(e\) exponent and \(m\) mantissa bits and bias \(2^{e-1}-1\), 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}\).

FormatExponent / mantissa bitsSmallest subnormalLargestRelative precision
FP328 / 232−1493.4×10382−24
BF168 / 72−1333.4×10382−8
FP165 / 102−2465,5042−11
FP8 E5M25 / 22−1657,3442−3
FP8 E4M34 / 32−94482−4

BF16 (bfloat16) is 2 bytes like FP16 but keeps FP32’s 8-bit exponent: the same range as FP32 and about 2–3 significant decimal digits (Kalamkar et al. 2019). FP16 has more precision but a narrow range.

03

Rounding a Number, Live

Type a value; it is rounded to nearest (ties to even) in each format, with subnormals and overflow. FP32 uses the browser’s own Math.fround; the others the spacing formula on the previous slide.

FormatRoundedRelative error
04

BF16 and FP16 on Real Gradients

All 5,376 gradients of the trained model, rounded to each format (results.json gradRounding, the site’s chapter 9 animation):

Format (no scaling)Gradients flushed to zeroMean relative error
FP320 of 5,3762.2e-08
BF160 of 5,3761.4e-03
FP1669 of 5,3766.1e-03
FP8 E5M2405 of 5,3764.8e-02
FP8 E4M31,601 of 5,3761.2e-01

BF16’s range keeps every gradient; FP16 loses the smallest; FP8 loses many. That trade is why most large runs today compute in BF16: it needs no loss scaling. Multiplying the loss by \(2^k\) before the backward pass shifts every gradient up by \(k\) binades, into FP16’s range:

0 500 1,000 1,500 2,000 2,500 0 4 8 12 16 20 24 loss scale 2^k: k FP16 gradients lost (of 5,376) flushed to zerooverflowed
The trained model’s gradients in FP16 at loss scales 20 to 224 (results.json gradRounding): too little scale flushes small gradients to zero, too much overflows large ones. Dynamic scaling looks for the gap.
05

Mixed Precision: Master Weights and Loss Scaling

Micikevicius et al. (2018) made 16-bit training work with three ideas:

FP32 master copy

Compute in 16 bits, but the optimiser updates an FP32 copy of the weights. An update is often smaller than the spacing between 16-bit values near the weight, so adding it to a 16-bit weight would do nothing.

FP32 accumulation

The matrix products accumulate their sums in FP32, so long dot products do not lose their small terms.

Loss scaling

\(\tilde L = 2^k L \Rightarrow \tilde g = 2^k g\); the optimiser divides by \(2^k\) in FP32. Dynamic scaling halves the scale and skips the step on an overflow, and doubles it after a long run of good steps (PyTorch’s GradScaler).

06

Six Recipes, Measured

The default run with every weight’s compute copy, every activation a matrix product reads or writes, and every gradient rounded to the format (float64 stands in for FP32 in the elementwise operations and the master weights); three seeds (results.md §9a):

RecipeFinal validation lossSeeds 1, 2, 3Gradient values flushed to zero
float64 (reference)2.277 ± 0.0902.299, 2.354, 2.1780.00%
FP322.316 ± 0.0852.298, 2.408, 2.2420.00%
BF16, FP32 master weights2.329 ± 0.0592.269, 2.334, 2.3860.00%
BF16, no master copy2.296 ± 0.0272.325, 2.290, 2.2720.00%
FP16, no loss scaling2.331 ± 0.0272.340, 2.352, 2.3014.56%
FP16, dynamic loss scaling2.288 ± 0.1312.266, 2.429, 2.1690.13%
FP8 (E4M3 / E5M2), no scaling2.957 ± 0.0112.945, 2.959, 2.96715.82%
FP8 (E4M3 / E5M2), per-tensor scaling2.574 ± 0.3382.383, 2.965, 2.3741.96%

FP32, BF16 with master weights and FP16 with dynamic loss scaling all end within the seeds’ spread of the float64 reference. So, at this size, do the two recipes that should hurt: FP16 without scaling (although 4.6% of its gradient values were flushed to zero) and BF16 with no master copy (although 43% of its weight updates were rounded away). A 5,376-parameter model trained for 2,000 steps is forgiving; at scale, with longer runs and smaller late learning rates, the same losses accumulate, which is why master copies and loss scaling are standard.

07

FP8 and Per-Tensor Scaling

The newest accelerators also multiply in 8-bit floats. Micikevicius et al. (2022) recommend E4M3 (largest 448) for weights and activations and E5M2 (more range, less precision) for gradients, with a scaling factor per tensor so each tensor’s largest value lands near the top of the format:

$$s = 2^{\lfloor \log_2(\max_{\text{fmt}} / \operatorname{amax}(x)) \rfloor}, \qquad \hat x = \operatorname{round}(s x)/s$$

No scaling

Never leaves the plateau: 2.957 ± 0.011, with 16% of gradient values flushed to zero.

Per-tensor scaling

Two of 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.

In the emulation only what the matrix products read is FP8; what they write stays BF16, as in those recipes. The site’s notebook 09 trains a 60-million-parameter model in each format on a GPU: GPU run pending, so no GPU number appears here.

08

Quantised Gradients

Beyond floating point, values can be quantised to a few bits of integer: pick a full scale \(F\), divide it into \(n = 2^{b-1}-1\) steps each side of zero, round each value to the nearest step, and clip anything beyond the full scale.

$$\hat x = \Delta \cdot \operatorname{clip}\big(\operatorname{round}(x/\Delta), -n, n\big), \qquad \Delta = F/n$$
BitsLevelsFinal validation lossSeeds 1, 2, 3
82552.286 ± 0.0692.321, 2.331, 2.207
4152.413 ± 0.2092.653, 2.269, 2.315
232.860 ± 0.1822.650, 2.966, 2.964

Every weight gradient quantised per tensor (full scale = its largest value) before clipping and the optimiser step; three seeds (results.md §9b). Adam is forgiving here too, since it normalises each weight’s step (Alistarh et al. 2016 analyse quantised SGD). Quantised gradients are what compressed distributed training sends between workers (deck 12), and analogue or optical hardware imposes the same quantisation on whatever passes through it.

09

What to Take Away