Module 10 · Part III · Building an LLM

Pretraining

Pretraining Data, AdamW, Warmup and Decay, Keeping Training Stable, and Scaling Laws

By the end you can plan a pretraining run (data, optimizer, schedule, model size) and justify each choice.

Builds onModule 9·
Check yourself ▸

Before you start, can you already…

  • ☐run BPE by hand and predict how the tokenizer changes cost and behaviour. (Module 9)

If any is shaky, the linked module teaches it.

Open in Colab
5 Hours·Core·Intermediate·0/5

Pretraining is the long first phase in which a model learns from raw text by next-token prediction (section 8.4), and where almost all of its knowledge and training compute go. This module covers the decisions that shape a run: the data, the optimizer and its learning-rate schedule, keeping the run stable, and splitting a compute budget between model size and data.

10.1 Pretraining DataRun this section

A model can only learn what is in its training text, so the data is as much a part of the design as the architecture. Building a pretraining set of trillions of tokens is a pipeline of its own:

  1. Collect. Mostly web crawls (Common Crawl is the usual start), plus books, encyclopedias, papers, code, and conversational and mathematical text.
  2. Extract and filter. Web pages are mostly menus, boilerplate and spam. Text is extracted from the HTML, filtered by rules (minimum length, share of lines ending in punctuation, repeated lines) and often by a small classifier for high-quality text such as reference or educational material, plus filters for language, toxicity and personal information.
  3. Deduplicate. Duplicates waste compute, make models reproduce text verbatim, and leak test sets into training (section 13.3). Exact copies are removed by hashing; near-duplicates need the fuzzier test below.
  4. Mix. Sources are combined in chosen proportions, the data mixture, tuned by training small models on candidate mixes. More code or mathematics improves those skills, often with benefits elsewhere, and many runs finish on the highest-quality data.
  5. Synthesize. Increasingly, other models write part of the data: rewritten web text, textbook-style explanations, worked problems with verified answers. It is cheap and targets gaps, but overuse narrows the data's variety.

Finding Near-Duplicates

Near-duplicates share most of their short word sequences. Cut each document into overlapping shingles of kk consecutive words and compare the sets with the Jaccard similarity, shared shingles over all shingles:

J(A,B)=∣A∩B∣∣A∪B∣J(A, B) = \frac{|A \cap B|}{|A \cup B|}

JJ runs from 0 (nothing shared) to 1 (identical). Comparing every pair is impossible at web scale, so pipelines use MinHash: hash every shingle with many hash functions and keep each function's smallest value. Two documents agree on one such minimum with probability exactly JJ, so a signature of, say, 128 minimums estimates JJ without comparing full sets, and matching signatures are quick to find.

Figure 10.1
Near-duplicates share most of their shingles

Compare the original document with an edited copy or an unrelated one. Exact Jaccard similarity of the shingle sets, against its MinHash estimate.

Each MinHash slot uses a different hash function and keeps the smallest value; the share of slots on which two documents agree estimates their Jaccard similarity.
Predict, then try

The second document is the first with one phrase reworded. Before looking, estimate the Jaccard similarity of their 3-word shingles: closer to 0.2, 0.5 or 0.9?

In words. An original, a lightly edited copy and an unrelated text, with each pair's shared shingles, exact Jaccard similarity and MinHash estimate. The estimate settles near the exact value as the signature grows, and the edited copy scores far above the unrelated text.

Browser lab

MinHash estimates Jaccard similarity, more precisely with longer signatures

60 pairs of 200-element sets with random overlaps. For each, compare the MinHash estimate (the share of k hash functions on which the two sets’ minimums agree) with the exact Jaccard similarity. Each slot agrees with probability J, so by section 3.2’s rules the estimate’s standard error should be about √(J(1−J)/k), falling like 1/√k.

Simple multiplicative hashes stand in for the random permutations of the theory; real pipelines also band the signatures so that likely duplicates can be found without comparing every pair.

10.2 AdamWRun this section

Plain gradient descent (section 2.4) uses one learning rate for every weight. Adam (Kingma & Ba, 2015) gives each weight its own effective step size, from two running averages of its gradient :

mt=β1mt−1+(1−β1) gt,vt=β2vt−1+(1−β2) gt2\mathbf{m}_t = \beta_1 \mathbf{m}_{t-1} + (1 - \beta_1)\,\mathbf{g}_t, \quad \mathbf{v}_t = \beta_2 \mathbf{v}_{t-1} + (1 - \beta_2)\,\mathbf{g}_t^2

mt\mathbf{m}_t tracks the gradient's recent average (direction) and vt\mathbf{v}_t its average square, entry by entry (size). Language models typically use β1=0.9\beta_1 = 0.9, β2=0.95\beta_2 = 0.95. The update divides one by the square root of the other:

θ⟵θ−η(m^tv^t+ϵ+λθ)\theta \longleftarrow \theta - \eta \left(\frac{\hat{\mathbf{m}}_t}{\sqrt{\hat{\mathbf{v}}_t} + \epsilon} + \lambda \theta\right)

The hats correct for the averages starting at zero: m^t=mt/(1−β1t)\hat{\mathbf{m}}_t = \mathbf{m}_t / (1 - \beta_1^t), likewise for v\mathbf{v}. Dividing by v^t\sqrt{\hat{\mathbf{v}}_t} makes the step size roughly independent of the gradient's scale, so rarely updated weights, such as rare-token embeddings, still move. λθ\lambda \theta is weight decay (section 4.4), typically λ=0.1\lambda = 0.1; the "W" in AdamW (Loshchilov & Hutter, 2019) means it is applied directly, as here, rather than mixed into the gradient, where Adam's division would distort it.

The price is memory: two extra numbers per weight. With bf16 weights and gradients plus fp32 master weights and averages, the rule of thumb is 16 bytes per parameter, before activations.

10.3 Learning-Rate Warmup and DecayRun this section

Training starts with a warmup: the rate rises linearly from near zero over the first few thousand steps, because early on the weights are random and Adam's averages unreliable, and full-size steps can throw the model somewhere it never recovers from (Post-LN models need this most, section 8.5). Then the rate decays, usually along a cosine to about a tenth of its peak, so late training settles with small steps.

Figure 10.3
Warmup, then cosine decay

The learning rate rises linearly from zero, peaks, then follows half a cosine down to a floor. Drag the probe to read the rate at any step.

Illustrative run of 100,000 steps with peak rate 3 × 10⁻⁴, a common choice for models of a few billion parameters.
Predict, then try

Set the warmup to zero. What is the learning rate at the very first step, and why is that the riskiest moment of the whole run?

In words. The learning rate climbs linearly during warmup, peaks, then falls along half a cosine to a floor of a tenth of the peak. Longer warmup delays the peak; a higher floor keeps late steps larger.

10.4 Keeping Training Stable

A large run lasts weeks, so one divergence (the loss shooting up and never recovering) is expensive. Three standard defences, alongside warmup:

  • Gradient clipping. A gradient longer than a threshold, usually 1.01.0, is scaled down to it, so one bad batch cannot throw the weights far (Pascanu et al., 2013).
  • Initialization. Weights start small, typically normal with standard deviation around 0.020.02: too large and activations grow layer by layer, too small and the signal fades (the variance argument of section 6.3). GPT-2 also scales each output projection's initial weights by 1/N1/\sqrt{N} for NN residual layers, since NN random additions to the residual stream would otherwise grow its variance NN-fold.
  • Mixed precision. Forward and backward passes run in bf16 (fp32's range, section 4.3), while master weights and optimizer averages stay fp32 so tiny updates are not rounded away (Micikevicius et al., 2018).

10.5 Scaling Laws: How Big, and on How Much TextRun this section

Kaplan et al. (2020) found that loss falls as a smooth power law in parameters NN, training tokens DD and compute CC, over many orders of magnitude. (In this section only, NN counts parameters.) A run's compute is about:

C≈6NDC \approx 6 N D

Each token costs about 2N2N operations forward (a multiply and an add per weight) and 4N4N backward (twice the forward, section 2.5). Hoffmann et al. (2022), the "Chinchilla" paper, fitted the loss in NN and DD separately:

L(N,D)=E+ANa+BDb\mathcal{L}(N, D) = E + \frac{A}{N^{a}} + \frac{B}{D^{b}}

EE is the loss no model can beat, the randomness of text itself; the other terms penalize too small a model and too little data. The paper's fitted values are E=1.69E = 1.69, A=406.4A = 406.4, B=410.7B = 410.7, a=0.34a = 0.34, b=0.28b = 0.28 (it names the exponents and β\beta). Its direct fits of the best model at each budget give a simple rule: grow NN and DD together, at about 20 tokens per parameter. Minimizing the formula with those constants disagrees, favouring a smaller model on more data; a replication (Besiroglu et al., 2024) traced this to the fit, and refitted constants (E=1.82E = 1.82, A=482.01A = 482.01, B=2085.43B = 2085.43, a=0.3478a = 0.3478, b=0.3658b = 0.3658) agree with the 20-to-1 rule. Earlier models had been far too large for their data: Chinchilla, 70B parameters on 1.4 trillion tokens, beat the 280B Gopher on 300 billion tokens with about the same compute.

Figure 10.5
The compute-optimal model

Contours are equal predicted loss over model size N and training tokens D. The dashed line is every run with the same compute C ≈ 6ND; the best model on it is where it just touches the lowest contour.

Loss from the fitted formula L(N, D) = E + A/Nᵃ + B/Dᵇ. The 2024 refit re-estimated the constants from the paper’s own data and recovers the paper’s other finding, about 20 tokens per parameter; the original constants put the optimum at a smaller model and several times more tokens per parameter.
Predict, then try

Slide along the dashed line of equal compute from the smallest model to the largest. Predict the shape of the loss along the way before you look.

In words. Lines of equal predicted loss over model size and training tokens, with a dashed line of equal-compute runs. Along it the loss falls then rises; the lowest point is the compute-optimal model. With the refitted constants, at Chinchilla's budget it lands near 70 billion parameters and 1.3 trillion tokens, while Gopher sits on the same budget at a higher loss.

The rule minimizes cost and ignores the cost of using the model, paid on every generated token and growing with NN. So current models train far past the Chinchilla point: LLaMA 3 8B saw 15 trillion tokens, nearly 2,000 per parameter. Smaller and longer-trained costs more to train but less to run.

Exercises

Check

Name the five stages of building a pretraining dataset described in section 10.1, in order.

Apply

Estimate the training compute of an 8-billion-parameter model trained on 15 trillion tokens, in FLOPs. Give the mantissa in units of 102310^{23}.

Apply

How many gigabytes of training state does AdamW need for a 7-billion-parameter model at 16 bytes per parameter, before activations?

Apply

Under the 20-tokens-per-parameter rule, C=6NDC = 6ND with D=20ND = 20N. What is the compute-optimal model size, in billions of parameters, for C=5.76×1023C = 5.76 \times 10^{23} FLOPs?

Apply

Two documents have shingle sets of sizes 40 and 50 and share 30 shingles. What is their Jaccard similarity?

Extend

Give two separate reasons why training on many copies of the same document is harmful, beyond wasted compute.

Core Academic References

  • Kingma & Ba (2015): Adam: A Method for Stochastic Optimization. ICLR 2015. arXiv:1412.6980
  • Loshchilov & Hutter (2019): Decoupled Weight Decay Regularization (AdamW). ICLR 2019. arXiv:1711.05101
  • Pascanu, Mikolov & Bengio (2013): On the difficulty of training recurrent neural networks (gradient clipping). ICML 2013. arXiv:1211.5063
  • Micikevicius et al. (2018): Mixed Precision Training. ICLR 2018. arXiv:1710.03740
  • Kaplan et al. (2020): Scaling Laws for Neural Language Models. arXiv:2001.08361
  • Hoffmann et al. (2022): Training Compute-Optimal Large Language Models (Chinchilla). NeurIPS 2022. arXiv:2203.15556
  • Besiroglu et al. (2024): Chinchilla Scaling: A Replication Attempt. arXiv:2404.10102
  • Broder (1997): On the Resemblance and Containment of Documents (MinHash). Compression and Complexity of Sequences 1997.
  • Lee et al. (2022): Deduplicating Training Data Makes Language Models Better. ACL 2022. arXiv:2107.06499
  • Penedo et al. (2024): The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale. arXiv:2406.17557