Half the bytes: mixed precision and checkpoints

Build a Small GPT From Scratch · Part 5 of 8

Two habits separate a run that finishes from one that does not: storing numbers in fewer bytes, and saving progress often enough that an interruption costs minutes instead of hours.

Where the memory goes

Training with ordinary 32-bit floats needs four numbers for every parameter: the weight, its gradient, and the two running averages Adam keeps. That is 4 numbers of 4 bytes, or 16 bytes per parameter, before counting a single activation. A 125-million-parameter model therefore needs about 2.0 GB just for this, on a free GPU that has to share its memory with everything else.

Mixed precision

Mixed precision does most of the arithmetic in 16-bit numbers, which take half the memory and run faster on modern GPUs, while keeping a 32-bit master copy of the weights so tiny updates are not lost. There are two common 16-bit formats. float16 has more precision but a narrow range: its smallest normal number is about 6.1e-05 and its largest 65,504. bfloat16 keeps the range of 32-bit floats and gives up some precision, which is why it needs no extra care and is the usual choice where the hardware supports it.

Gradients are often very small, and in float16 a value like 1e-8 simply turns into zero. The fix is loss scaling: multiply the loss by a large factor before backpropagation so gradients are big enough to survive, then divide the factor back out before the update. Training libraries do this automatically; seeing it once makes their warnings less mysterious.

import numpy as np
g = 1e-8                                   # a typical tiny gradient
print(np.float16(g))                       # 0.0  (underflow: the signal is gone)
scale = 1024.0
kept = np.float16(g * scale)               # still representable
print(float(kept) / scale)                 # about 1e-08 once the scale is divided back out

Checkpoints

A free Colab session can be interrupted at any time, and a multi-hour run will be. A useful checkpoint stores everything needed to carry on as if nothing had happened: the model weights, the optimizer state (Adam’s running averages), the step number and, ideally, the random-number state. Saving only the weights lets you resume, but the optimizer then starts cold, and the learning-rate schedule from the previous post would restart from the wrong place.

import torch

model = torch.nn.Linear(4, 2)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
step = 120

torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step}, "ckpt.pt")

state = torch.load("ckpt.pt")              # later, in a fresh session
model.load_state_dict(state["model"]); opt.load_state_dict(state["opt"]); step = state["step"]
print("resumed at step", step)             # resumed at step 120

Takeaways

• Plain 32-bit training needs about 16 bytes per parameter before activations.
• Mixed precision does most arithmetic in 16 bits; bfloat16 keeps the range, float16 needs loss scaling.
• A tiny gradient can underflow to zero in float16; scaling the loss prevents it.
• A checkpoint must hold weights, optimizer state and step so a resumed run is the same run.

Next in the series (Part 6): How much will a training run cost? C = 6ND