Mixed precision: 16-bit numbers (bf16)

Read these first if they are new to you

PyTorch tensors on the GPU Intermediate
Gradient descent Intermediate
Backpropagation: gradients by the chain rule Intermediate

The idea in one paragraph

By default, PyTorch stores each number in 32 bits. A 16-bit format called bf16 keeps the same range of sizes but fewer digits: about 2 to 3 significant decimal digits instead of about 7. Neural networks tolerate that rounding in most of their arithmetic. Mixed precision runs the big matrix multiplications in bf16, which needs half the memory and runs much faster on GPUs built for it. It keeps the few sensitive steps, such as updating the weights, in 32 bits.

A float stores a sign, a power of two and a fraction

A computer stores a number such as 6.5 in base-2 scientific notation: 6.5 = 1.625 × 2². The stored pieces are a sign bit (plus or minus), an exponent (the power of two, here 2) and a fraction (the digits after the point of 1.625, written in binary). A number stored this way is a floating-point number, or float. The number of exponent bits sets how large and how small a number can be. The number of fraction bits sets how many significant digits it keeps.

Format	Sign	Exponent	Fraction	Largest value	Significant decimal digits
float32	1	8	23	3.4 × 10³⁸	about 7
bf16 (bfloat16)	1	8	7	3.4 × 10³⁸	2 to 3
float16	1	5	10	65,504	3 to 4

bf16 is the first 16 bits of a float32, rounded. It keeps all 8 exponent bits, so it covers the same range of sizes, and drops 16 of the 23 fraction bits. The older 16-bit format, float16, splits its bits the other way: more fraction, much less range.

Try it. Type a number or pick one from the list, and compare the bits and the stored values. Goal: find the smallest whole number that bf16 cannot store exactly. Then find a number that float16 cannot store at all.

The same number in three formats

Number or pick

sign exponent fraction

What to notice. The bf16 row is the left half of the float32 row, give or take rounding in its last bit. Its exponent matches float32's, so very large and very small numbers survive. float16 has more fraction bits but only 5 exponent bits, so 70000 becomes infinity and 1e-30 becomes zero.

Rounding can swallow a small weight update

Floats are not evenly spread. Between 1 and 2, bf16 has 7 fraction bits, so it can store only 2⁷ = 128 values, spaced 1/128 = 0.0078125 apart. float32 has 2²³ values there, spaced about 0.00000012 apart.

Worked example: a weight that will not move

A weight stored in bf16 equals 1.0. A training step wants to add 0.001.
The exact result, 1.001, is not a bf16 value. The nearest bf16 values are 1.0 and 1.0078125.
1.001 is much closer to 1.0, so the sum rounds to 1.0. The update is lost.
Each later step that adds 0.001 is lost the same way, so this weight never changes.
In float32, the sum is 1.00100004673, within 0.00000005 of 1.001. The update is kept.

That is why mixed precision keeps a float32 copy of the weights and applies each update to it.

Your turn. Between 256 and 512, bf16 values are spaced 256/128 = 2 apart. What does 257 round to? (Answer: 256. 257 lies halfway between 256 and 258, and the rule for a tie picks the value whose last bit is 0, which is 256.)

Mixed precision picks the format operation by operation

Most of the time and memory in training goes into large matrix multiplications and the activations, the intermediate results kept for the backward pass. These tolerate bf16. A few operations do not: long sums, such as the loss over a batch, and the update of the weights. Mixed precision keeps those in float32.

PyTorch does the sorting for you. Inside a torch.autocast block, each operation runs in the format that PyTorch has listed as safe for it. The weights themselves stay float32. The snippet shows the rounding from the worked example and one multiplication under autocast. It runs on the CPU; on a GPU you write device_type="cuda" (see tensors on the GPU).

bf16.py

import torch

third = torch.tensor(1 / 3)                       # float32 by default
print(third.item())                               # float32's closest value
print(third.to(torch.bfloat16).item())            # bf16's closest value

w = torch.tensor(1.0, dtype=torch.bfloat16)
print((w + 0.001).item())                         # a small update to a bf16 weight

big = torch.tensor(70000.0)
print(big.to(torch.float16).item(), big.to(torch.bfloat16).item())

a, b = torch.randn(64, 64), torch.randn(64, 64)
with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    c = a @ b                                     # runs in bf16 inside autocast
print(a.dtype, c.dtype)

0.3333333432674408
0.333984375
1.0
inf 70144.0
torch.float32 torch.bfloat16

The inputs a and b are float32, but autocast ran the multiplication in bf16, so the result c is bf16.

Why bf16 rather than float16

Gradients are often tiny numbers. In float16, numbers below about 6 × 10⁻⁵ start to lose digits, and numbers below about 3 × 10⁻⁸ become zero, so small gradients vanish. Training in float16 therefore multiplies the loss by a large factor first and divides the gradients by it afterwards, a trick called loss scaling. bf16 has float32's range, so it needs no loss scaling. That is why bf16 is the usual choice on GPUs that support it.

This is a simplification

"Much faster" depends on the GPU: older GPUs have no fast bf16 arithmetic at all, and the gain is large only when matrix multiplications take most of the time. Autocast's lists of which operations run in 16 bits differ between devices and PyTorch versions. Inside a 16-bit matrix multiplication, the hardware usually adds up the products in 32 bits, which is why the rounding hurts less than the bit counts suggest.

Check yourself

Why can bf16 store 1 × 10³⁰ while float16 cannot?

Answer
bf16 has 8 exponent bits, like float32, so its largest value is about 3.4 × 10³⁸. float16 has 5 exponent bits, and its largest value is 65,504. Anything bigger becomes infinity.
Near 1, bf16 values are 0.0078125 apart. A bf16 weight equals 1.0 and the update is +0.005. What is the result?

Answer
1.0078125. The exact sum, 1.005, is closer to 1.0078125 (0.0028 away) than to 1.0 (0.005 away). The weight moves, but by 0.0078125, not 0.005.
A model has 50 million weights. How much memory do the weights take in float32, and in bf16?

Answer
float32 uses 4 bytes per number: 200 MB. bf16 uses 2 bytes: 100 MB. With mixed precision the weights stay in float32; the saving comes from the activations and the arithmetic.

Go further

Optional background, for readers who want more depth.

bfloat16 floating-point format Wikipedia Article Intermediate

The bit layout with worked encodings, and how bf16 relates to float32.
Automatic Mixed Precision: adding torch.autocast PyTorch recipes Tutorial Intermediate

A short recipe that times a network with and without autocast. It uses float16, so it also shows loss scaling.

Where this shows up

Units that build on this one

The training loop Intermediate