Training LLMs

You now have every piece of a language model and every piece of an optimizer. Training is the loop that joins them, repeated hundreds of thousands of times. This page is about that loop:

• how batches are cut from a stream of tokens;
• how a large batch is built from small ones;
• how training fits in 16-bit numbers;
• how much compute a model of a given size needs.

A note on scale: the exercises run in your browser, where NumPy has no optimized matrix library and a $512 \times 512$ matrix product takes most of a second. So they train bigram-sized models in a few hundred steps. The loop is exactly the one used for a 7-billion-parameter model; only the model inside it is smaller.

The training loop

Every step does the same six things:

1. Sample a batch of token sequences.
2. Forward: compute the logits and the loss.
3. Backward: compute the gradient of the loss with respect to every parameter.
4. Clip the gradient's global norm (section 3).
5. Step: update the parameters with AdamW (Adam with decoupled weight decay), at the learning rate the schedule gives for this step.
6. Every few hundred steps, evaluate the loss on held-out data, never used for training.

Example. A bigram model (a $3 \times 3$ table of logits, all zeros at the start) trained on the repeating text 0 1 2 0 1 2 ..., with plain gradient descent. The training loss at a few steps:

step	loss
0	1.099
10	0.235
20	0.115
40	0.055
59	0.037

At step 0 every token has probability $1/3$, so the loss is $\log 3 \approx 1.099$: the uniform baseline of section 2. Then it falls towards 0, because in this text the next token is always certain.

The training loss going down only shows that the model fits what it has seen. The validation loss, measured on the held-out data, shows whether it generalizes to text it has never seen. When the two separate (training loss still falling, validation loss flat or rising), the model is memorizing its training data.

Batches from a token stream

The training corpus is tokenized once, into one long array of ids. A batch is $B$ random windows of $T + 1$ consecutive tokens. The first $T$ tokens of a window are the input $x$; the same window shifted by one is the target $y$, so that $y_t = x_{t+1}$: at every position, the target is the next token.

Example. With the data [10, 11, 12, 13, 14], $T = 3$ and a window starting at position 1: the input is [11, 12, 13] and the target is [12, 13, 14]. The model learns to predict 12 after 11, 13 after 11 12, and 14 after 11 12 13: three training examples from one window.

One window gives $T$ training examples at once, one per position, which is why causal masking matters (section 7). This is how the get_batch function of nanoGPT works.

Gradient accumulation

LLMs (large language models) are trained with batches of millions of tokens, far more than fit in the memory of a GPU (graphics processing unit) at once. Gradient accumulation splits a batch into $M$ micro-batches, runs forward and backward on each, adds up the gradients, and takes one optimizer step at the end.

Because the loss is a mean, the gradient of the full batch is the mean of the micro-batch gradients when the micro-batches have the same number of tokens. Otherwise each one must be weighted by its share of the tokens.

Example. Two micro-batches with 3 and 1 targets have mean losses $2.0$ and $4.0$. The mean over all 4 targets is $(3 \cdot 2.0 + 1 \cdot 4.0)/4 = 2.5$. A plain average of the two means would give $3.0$, as if the small micro-batch counted as much as the big one.

Done correctly, accumulation gives the same update as one big batch, only slower. The same idea, spread over many GPUs instead of over time, is data parallelism (section 9).

Mixed precision

How a computer stores a real number

A floating-point number is stored like scientific notation, $1.234 \times 10^5$, but in base 2. It has three parts:

• a sign: plus or minus;
• an exponent (the $5$ in the example), which sets the range: how large or how small the number can be;
• a mantissa (the $1.234$), which sets the precision: how many significant digits are kept.

With more exponent bits a format can store huger and tinier numbers; with more mantissa bits it can tell closer numbers apart.

Three formats

Every number in the loop can be stored in 32 bits, the FP32 (32-bit floating point) format. But 16-bit formats halve the memory and run several times faster on modern accelerators. The catch is range and precision. The three formats in use are FP32, FP16 (16-bit floating point) and BF16 (brain floating point, 16 bits; also written bfloat16):

(Exp. and Mant. are the numbers of exponent and mantissa bits; one more bit is the sign. "Min normal" is the smallest number stored at full precision; smaller ones lose digits.)

FP16: precise, but a small range

FP16 has precision but little range: its largest number is 65504, and gradients are often smaller than $10^{-5}$; below about $6 \times 10^{-8}$ they round to exactly zero. Micikevicius et al. (2018) made FP16 training work with two techniques.

An FP32 master copy of the weights, to which the updates are applied. A small update added to a large weight in FP16 is lost to rounding.

Example. In FP16, the number right after $1.0$ is $1.000977$: nothing in between can be stored. So $1.0 + 0.0001$ rounds back to exactly $1.0$, and a weight of 1 receiving updates of $0.0001$ would never change. In FP32 the update is kept.

Loss scaling: multiply the loss by a large factor $S$ before the backward pass, so all gradients are $S$ times larger and stay representable, then divide them by $S$ (in FP32) before the update.

Example. A gradient of $10^{-8}$ becomes exactly 0 in FP16. Multiplied first by $S = 2^{16} = 65536$, it is about $0.000655$, which FP16 stores without trouble; dividing by $S$ in FP32 gives back about $10^{-8}$.

But a gradient of $1.0$ multiplied by the same $S$ is $65536$, above FP16's maximum of 65504: it becomes inf. The scale must be large enough for the small gradients and small enough for the large ones.

With dynamic loss scaling, $S$ starts high. It is halved whenever a gradient overflows to inf (and that step is skipped), and doubled again after a long run of good steps. PyTorch's GradScaler uses exactly this, with defaults $S = 2^{16}$, factors 2 and 0.5, and 2000 good steps before growing.

BF16: the range of FP32, less precision

BF16 keeps FP32's 8 exponent bits and drops mantissa bits. It has FP32's range, so gradients do not underflow and loss scaling is unnecessary (Kalamkar et al., 2019), at the cost of only 2 to 3 significant decimal digits: the number right after $1.0$ is $1.0078$. With BF16 one can compute in 16 bits and keep the master weights and optimizer states in FP32, with no loss scaling at all.

How much compute: $6ND$

The forward pass costs about 2 FLOPs (floating-point operations) per parameter per token, one multiply and one add, and the backward pass about twice that. So training a model with $N$ parameters on $D$ tokens costs

where $C$ is the total compute, $N$ the number of parameters and $D$ the number of training tokens (Kaplan et al., 2020, section 2.1). The 6 is 2 for the forward pass plus 4 for the backward pass.

Example. A 1-billion-parameter model trained on 20 billion tokens needs about $6 \cdot 10^9 \cdot 2 \cdot 10^{10} = 1.2 \times 10^{20}$ FLOPs.

MFU (model FLOPs utilization; Chowdhery et al., 2022) measures how well the hardware is used: the FLOPs the model needs per second, that is $6N$ times the tokens processed per second, divided by the peak FLOPs per second of the hardware. PaLM (Pathways Language Model) reports an MFU of 46.2%. The rest is lost to memory transfers, communication between chips and operations that do not use the hardware fully.

Example. Imagine 8 accelerators, each with a peak of $10^{15}$ FLOPs per second, running at an MFU of 40%: together they do $8 \cdot 10^{15} \cdot 0.4 = 3.2 \times 10^{15}$ useful FLOPs per second. The $1.2 \times 10^{20}$ FLOPs of the run above then take $1.2 \times 10^{20} / 3.2 \times 10^{15} = 37{,}500$ seconds, about 10.4 hours.

Scaling laws: how big, for how long

For a fixed compute budget $C$, should you train a larger model on fewer tokens, or a smaller one on more? Kaplan et al. (2020) found that the loss falls as a smooth power law in $N$, $D$ and $C$: double the compute, and the loss drops by a predictable amount.

Hoffmann et al. (2022) redid the analysis with learning-rate schedules matched to the length of each run, and found that $N$ and $D$ should grow in equal proportion: about 20 training tokens per parameter (their Table 3 gives 20.2 billion tokens for a 1-billion-parameter model). Their 70-billion-parameter Chinchilla, trained on 1.4 trillion tokens, outperformed the much larger Gopher (280 billion parameters) at the same compute.

Combining the two rules, $C = 6ND$ and $D = 20N$, gives $C = 6N \cdot 20N = 120 N^2$, so the compute-optimal size for a budget is $N = \sqrt{C / 120}$.

Example. A budget of $10^{21}$ FLOPs gives $N = \sqrt{10^{21}/120} \approx 2.9$ billion parameters, trained on $D = 20N \approx 58$ billion tokens.

Training far beyond 20 tokens per parameter can still pay off: a smaller model costs more to train than the optimum, but less every time it is used. LLaMA (Large Language Model Meta AI), for example, notes that its 7B model keeps improving even after 1 trillion tokens, about 140 per parameter.

References

• J. Kaplan et al., Scaling Laws for Neural Language Models, arXiv, 2020. Section 2.1: $C \approx 6N$ FLOPs per training token.
• J. Hoffmann et al., Training Compute-Optimal Large Language Models, arXiv, 2022 (NeurIPS 2022). Chinchilla; Table 3: compute-optimal tokens per model size.
• P. Micikevicius et al., Mixed Precision Training, ICLR, 2018. FP32 master weights and loss scaling.
• D. Kalamkar et al., A Study of BFLOAT16 for Deep Learning Training, arXiv, 2019. Table 1: FP32, FP16 and BF16 formats.
• A. Chowdhery et al., PaLM: Scaling Language Modeling with Pathways, arXiv, 2022. Section 4.1 and appendix B: model FLOPs utilization.
• PyTorch source, torch.amp.GradScaler: the dynamic loss-scaling rule and its defaults.
• A. Karpathy, nanoGPT: get_batch and the training loop in train.py.