L6.4 · LESSON 4 OF 6

Mixed precision

Describe a pretraining loop and reason about data, batch size, precision and a simplified compute budget.

THE IDEA

Fewer bits, more speed, new failure modes

Storing numbers in 16 bits halves memory and speeds up arithmetic compared with 32 bits. The price is range and precision. A float16 number can be as large as 65504 and as small as about 6·10⁻⁸, with only about three significant digits. Larger values overflow to infinity, tinier ones underflow to zero, and adding a small number to a big one can change nothing.

So "mixed precision" mixes: matrix products in 16 bits, but master weights, optimizer state and long sums in 32 bits, often with loss scaling (multiply the loss, hence the gradients, by a constant before backpropagation and divide afterwards).

number to store in float16

enter a number

loss scale1

gradient 1e-8 × 1 = 1.00e-8 stored in float16 → 0.00e+0 after dividing by the scale: 0 (the gradient was lost)

Frameworks keep a float32 copy of the weights and optimizer state and use scaling, because small updates vanish in 16 bits. bfloat16 keeps float32's range with less precision, so it rarely needs scaling.

Store any number in float16, then see how scaling rescues a tiny gradient.

Predict first. Does using smaller weight storage guarantee all optimizer state uses the same precision?

Chapter context

Pretraining optimizes next-token predictions over a dataset. Its costs depend on the architecture, sequence lengths, implementation and hardware. A useful budget model states which work it approximates and which work it ignores.

TRY IT YOURSELF

Part 1. Do it hands-on

Find the float16 limits and apply loss scaling by hand, then reproduce overflow, underflow and a stalled running sum with NumPy's real float16.

number to store in float16

enter a number

loss scale1

gradient 1e-8 × 1 = 1.00e-8 stored in float16 → 0.00e+0 after dividing by the scale: 0 (the gradient was lost)

largest float16 70000 in float16 (inf, 0 or 70000) 1e-8 × 1024

To do: The largest finite half-precision (float16) number is…
To do: What does 70000 become when stored in float16? Type inf, 0 or 70000
To do: A gradient of 1e-8 underflows to 0 in float16. Loss scaling multiplies it by 1024 first. Type the scaled value (e.g. 0.00001024)

Part 2. Do it in code

to_half(x): the value stored as np.float16 and read back as a Python float.
overflows(x): True if x is finite but becomes infinite in float16.
underflows(x): True if x is non-zero but becomes 0 in float16.
scale_roundtrip(g, scale): store g·scale in float16, unscale in float32 (divide by scale) and return the float. Use it to show loss scaling saves tiny gradients.
half_sum(values) and float_sum(values): add the values one at a time keeping the running total as np.float16 / as np.float32 and return it as a float.

import numpy as np

def to_half(x):
    pass

def overflows(x):
    pass

def underflows(x):
    pass

def scale_roundtrip(g, scale):
    pass

def half_sum(values):
    pass

def float_sum(values):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Does using smaller weight storage guarantee all optimizer state uses the same precision?

ANoBYes

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

A commonly used rough dense-transformer training estimate is

C\approx6ND,

with N model parameters, D training tokens and C floating-point operations. It omits some architecture-dependent work and is not a hardware throughput guarantee. Effective token batch is approximately $B\times T\times A$ for microbatch B, length T and A accumulation steps.

Worked example

For N=1 million parameters and D=100 million training tokens, the rough compute estimate is 6×10^6×10^8=6×10^14 FLOPs. A microbatch of 4 sequences of length 8 accumulated twice contributes 64 token positions per optimizer update. Neither calculation includes padding exclusions or communication time.

Reference code

The following example uses the Python standard library and does not call a remote service.

N,D=1_000_000,100_000_000
assert 6*N*D == 600_000_000_000_000
batch,length,accumulation=4,8,2
assert batch*length*accumulation == 64

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Compute-optimal ratios are model- and study-dependent. Train/test leakage can arise from shared documents and overlap. Tokens per update and optimizer steps are different units.