Pretraining at scale

Imagine writing an encyclopedia by reading everything ever printed. Two problems appear at once.

First, most of what is printed is junk: flyers, receipts, the same cookie notice on a million websites, spam. You need a sorting line that throws out the junk and the photocopies before anyone reads a page.

Second, no single reader can do the reading. You hire a thousand readers, and now you have a management problem: how do they split the work, share what they learned, and keep going when one of them gets sick?

Pretraining is exactly those two problems. The first is data curation. The second is distributed training.

flowchart LR
  W[Web crawl<br/>billions of pages] --> C[Curation<br/>language, quality,<br/>deduplication]
  C --> M[Mixture<br/>weights per source]
  M --> T[Tokenize<br/>trillions of tokens]
  T --> P[Parallel training<br/>data, tensor, pipeline]
  P --> MP[Mixed precision<br/>bf16 math, fp32 master]
  MP --> S[Stability<br/>clip, watch spikes,<br/>checkpoint]
  S --> B[Base model]

Everyday picture Think of a recycling plant. Trucks tip mixed rubbish onto a conveyor belt. Magnets pull out the steel, blowers lift out the paper, people pick out what the machines miss, and only a small fraction reaches the bale at the end. Pretraining data goes down a belt like this, and just as in a recycling plant, most of what goes in never comes out.

The raw material is usually Common Crawl, a nonprofit's public archive of the web: regular snapshots, each of billions of pages. A page arrives as HTML full of menus, adverts and scripts, so the first machine on the belt is text extraction, which keeps the main body text. After that come the filters this section builds: language identification, quality filters, deduplication and mixing. FineWeb, an open dataset built this way from 96 Common Crawl snapshots, ends with about 15 trillion tokens.

flowchart LR
  H[HTML page] --> X[Extract<br/>main text]
  X --> L{Language ID<br/>is it English?}
  L -- no --> D1[drop, or route to<br/>that language]
  L -- yes --> Q{Quality rules<br/>and classifier}
  Q -- fail --> D2[drop]
  Q -- pass --> E{Exact duplicate?<br/>hash of the text}
  E -- yes --> D3[drop]
  E -- no --> N{Near duplicate?<br/>MinHash + LSH}
  N -- yes --> D4[drop]
  N -- no --> K[Keep:<br/>goes to the mixture]

1a. Language identification

Everyday picture Overhear two words of a phone call, "le" and "et", and you already guess French. Every language has a handful of little words that turn up in almost every sentence.

Tiny worked example "the cat is in the garden and it was happy" has 10 words; 7 of them are on the English list of little words (the, is, in, the, and, it, was) and none are on the French, German or Spanish lists, so the guess is English with a share of 0.7. "SKU-4431 X99 blk/wht 12pk" matches no list at all, so it is marked unknown and dropped.

Production pipelines use a trained classifier over character sequences (fastText's language ID model covers 176 languages) and keep a page only if the classifier is confident. The idea is the same: count the evidence for each language and pick the strongest.

In code: detect_language counts each language's stop words (STOP_WORDS) and returns the language with the largest share, or "unknown" below 10%.

1b. Heuristic quality filters

Everyday picture A librarian sorting donations does not read every book. A book with no pages, a pamphlet that is all hashtags, a sheet of keywords: each is rejected at a glance by a simple rule.

Tiny worked example The Gopher paper (Rae et al., 2021) published a set of such rules. Here is what they say about a navigation bar, "Home | About us | Contact | Privacy policy | Log in":

The last rule is the clever one. Real prose cannot avoid words like "the" and "of"; keyword-stuffed pages and lists of product names avoid them completely. Other rules in the set catch pages that are mostly bullet points or mostly lines ending in "...".

In code: quality_failures applies each rule and returns the names of the ones a document breaks; an empty list means keep.

1c. Quality classifiers: scoring pages against a reference

Rules catch obvious junk. To prefer good text among the survivors, pipelines train a classifier: examples of the text you want (encyclopedia articles, books, pages that trusted sites link to) against random crawl, and then keep pages that score like the reference. GPT-3 filtered Common Crawl with a simple linear classifier of this kind; FineWeb-Edu asked a large model to rate pages for educational value and trained a small classifier on those ratings. We build the simplest version, naive Bayes: every word casts a vote, and the votes add up.

Tiny worked example Our reference examples contain the word "river" 3 times and the spam examples 0 times; "click" appears 0 times in the reference and 4 times in spam. So "river" votes for good and "click" votes for bad. A whole page's score is the sum of its words' votes.

import math
N_good, N_bad, V = 77, 45, 80
def vote(count_good, count_bad):
    # log P(w | good) / P(w | bad), with add-one smoothing
    p_good = (count_good + 1) / (N_good + V)
    p_bad = (count_bad + 1) / (N_bad + V)
    return math.log(p_good / p_bad)
# "river": 3 times in good, 0 in bad
round(vote(3, 0), 3)  # → 1.158
# "click": 0 times in good, 4 in bad
round(vote(0, 4), 3)  # → -1.837

Why "naive"? It treats every word as independent evidence, which is false (words come in phrases) but works well enough to sort billions of pages cheaply. The practical danger is that a classifier keeps whatever resembles its reference set, including its blind spots: pick only encyclopedia text as "good" and you quietly filter out dialects, forums and whole topics.

In code: QualityClassifier.trained_on_examples counts words in GOOD_EXAMPLES and BAD_EXAMPLES; QualityClassifier.word_vote is the formula for one word and QualityClassifier.score adds the votes.

1d. Exact duplicates: one fingerprint per page

Everyday picture A cloakroom attendant does not compare every coat with every other coat. Each coat gets a numbered ticket, and two coats with the same ticket are the same coat.

The web is full of copies: mirrored sites, syndicated news, the same terms of service on a million shops. For exact copies the trick is a hash: a function that turns any text into a short fingerprint, always the same for the same text and almost never the same for different texts. Lowercase the text and squash its spaces first, so that trivially reformatted copies get the same fingerprint, then keep the first page with each fingerprint. One pass, one lookup per page, no comparisons.

In code: normalize_for_exact lowercases and collapses whitespace, and exact_dedup keeps the first document with each SHA-1 fingerprint.

1e. Near duplicates: shingles and Jaccard similarity

A scraper site that copies an article and adds "Read more on our site" defeats the exact hash: one changed character gives a completely different fingerprint. We need a measure of how much two pages overlap.

Everyday picture Cut each page into overlapping strips of a few words, like roof shingles, and put each page's strips in a bag. Two pages are near duplicates when their bags hold mostly the same strips.

Tiny worked example With 2-word shingles:

Three shingles are shared (the cat, cat sat, sat on) and seven are distinct across both, so the overlap is 3/7 = 0.43.

def shingles(text, k=2):
    ws = text.split()
    return {" ".join(ws[i:i + k]) for i in range(len(ws) - k + 1)}
A = shingles("the cat sat on the mat")
B = shingles("the cat sat on a mat")
# |A ∩ B| and |A ∪ B|
len(A & B), len(A | B)  # → (3, 7)
# J(A, B)
round(len(A & B) / len(A | B), 2)  # → 0.43

In code: shingles cuts a text into k-word windows (5 by default) and jaccard divides the intersection by the union.

1f. MinHash: estimating Jaccard from a few numbers

Jaccard needs both full shingle sets side by side. With billions of pages, comparing every pair is impossible (a billion pages make about 5 × 10¹⁷ pairs). MinHash compresses each page into a short list of numbers, its signature, such that comparing two signatures estimates their Jaccard.

Everyday picture Shuffle a deck containing every shingle in the world and deal from the top. Stop at the first card that belongs to page A, and separately at the first card that belongs to page B. Those two "first cards" are the same card exactly when the first card from A's-or-B's pile happens to be one they share. The more they share, the likelier that is.

Tiny worked example A = {a, b, c} and B = {b, c, d}, so J = 2/4 = 0.5. Shuffle the four letters in all 24 possible orders. In 12 of them the first letter from A ∪ B is b or c (a shared letter), and then A's first and B's first are the same letter. 12 / 24 = 0.5, exactly J.

import itertools, math
A, B = {"a", "b", "c"}, {"b", "c", "d"}
orders = list(itertools.permutations("abcd"))
def first(order, s):
    return next(x for x in order if x in s)
agree = sum(first(o, A) == first(o, B) for o in orders)
agree, len(orders)  # → (12, 24)
# P[min h(A) = min h(B)] equals J(A, B)
agree / len(orders), len(A & B) / len(A | B)  # → (0.5, 0.5)
# typical error of the estimate with k = 128 hashes
round(math.sqrt(0.5 * 0.5 / 128), 3)  # → 0.044

In code: MinHasher draws k hash functions of the form (a·x + b) mod p and MinHasher.signature keeps the minimum under each; estimate_jaccard counts agreeing slots.

1g. Locality-sensitive hashing: finding candidate pairs without comparing all of them

Signatures make each comparison cheap, but a billion pages still make too many pairs. Locality-sensitive hashing (LSH) avoids most comparisons: cut each signature into bands of numbers, and file every page into one bucket per band, keyed by that band's numbers. Only pages that share a bucket in at least one band are ever compared.

flowchart LR
  D[Page] --> S[5-word shingles]
  S --> H[k hash functions<br/>keep each minimum]
  H --> G[Signature<br/>k numbers]
  G --> B1[band 1: r numbers] --> K1[bucket]
  G --> B2[band 2] --> K2[bucket]
  G --> BB[band b] --> KB[bucket]
  K1 & K2 & KB --> C[Candidate pairs:<br/>pages sharing any bucket]
  C --> V[Check estimated Jaccard<br/>against the threshold]

def p_candidate(s, b=20, r=5):
    # 1 − (1 − s^r)^b
    return 1 - (1 - s ** r) ** b
round(p_candidate(0.8), 4)  # → 0.9996
round(p_candidate(0.3), 4)  # → 0.0475
# where the curve is steepest: about (1/b)^(1/r)
round((1 / 20) ** (1 / 5), 2)  # → 0.55

In code: lsh_candidate_probability is the formula, and near_duplicate_pairs runs the whole mechanism: signatures, band buckets, candidate pairs, and a final check against the threshold.

1h. Why duplicates hurt

Everyday picture A student who reads the same paragraph a hundred times can recite it, but has not learned a hundred paragraphs' worth. A model that sees a page thousands of times does the same: it spends capacity memorizing that page, and learns to recite it when prompted.

Tiny worked example A bigram model predicts each word from the word before it, by counting pairs. Train one on five short lines plus the boilerplate "click here to subscribe for free updates", and ask for the probability that it continues "click" into the whole boilerplate line:

Kept once, the model has many ways to continue "click". Repeated 100 times, the line becomes the only road, and the model regurgitates it 93% of the time. Lee et al. (2021) found a single 61-word sentence repeated more than 60,000 times in the C4 dataset, and that models trained on deduplicated data emit memorized training text about ten times less often. Duplicates also waste compute and leak into test sets: a benchmark question copied across the web ends up in the training data, and the model's score on it measures recall, not skill.

In code: verbatim_probability trains the bigram counts on OTHER_LINES plus the given number of copies of BOILERPLATE and multiplies the next-word probabilities along the line.

1i. The whole belt on a small crawl

CRAWL_SAMPLE holds seven pages. Running curate over it:

Two of seven survive. That is not unusual: large public pipelines keep only a small fraction of the raw crawl they start from.

In code: curate runs language ID, the quality rules, the exact hash and a MinHash comparison in that order, and returns one verdict per page.

1j. Data mixtures: how much of each source

Everyday picture A diet is not "all the food in the shop". You choose proportions: mostly staples, some vegetables, a little of the rich stuff. Pretraining data is mixed the same way from sources that differ in size and value: web text, code, books, encyclopedias, scientific papers, maths.

Each source gets a mixture weight: its share of the tokens the model will see. Weights do not follow size. A small, valuable source (an encyclopedia) is often up-weighted, which means the model sees it more than once; a huge, noisy one (the web) is sampled less than one full pass.

weights = {"web": 0.80, "wiki": 0.05, "code": 0.15}
available = {"web": 900e9, "wiki": 20e9, "code": 150e9}
D = 1000e9
# epochs_i = w_i D / N_i
{k: round(weights[k] * D / available[k], 2) for k in weights}  # → {'web': 0.89, 'wiki': 2.5, 'code': 1.0}

Repeating a small source a few times is fine; many more passes and the model starts memorizing it (the duplicate problem again, on purpose). Mixture weights are usually chosen by training small models on candidate mixtures and comparing them, and many recipes change the mixture near the end of training, up-weighting the cleanest sources.

In code: epochs_per_source applies the formula to every source.

1k. Token budgets: how much data for how big a model

Everyday picture A bigger brain can learn more, but only if you give it more to read. Given a fixed amount of study time (compute), there is a best split between "bigger brain" and "more reading".

All budgets are counted in tokens, the pieces a tokenizer cuts text into (see primer.ml.tokenization); in English a token is roughly three quarters of a word. Two rules of thumb set the scale.

N = 7e9
# D_opt ≈ 20 N
D = 20 * N
D  # → 140000000000.0
# C ≈ 6 N D
C = 6 * N * D
C  # → 5.88e+21
# GPU-hours at 400 teraFLOP/s each
round(C / 400e12 / 3600)  # → 4083

In practice, models that will serve billions of requests are trained far past this point, because a smaller model trained longer is cheaper to run forever after. Llama 2 7B saw 2 trillion tokens (about 286 per parameter); Llama 3 was trained on about 15 trillion tokens across the family (15.6 trillion for the 405B). That is why data, not compute, is now often the binding limit, and why the next topic exists.

In code: chinchilla_tokens and training_flops are the two rules of thumb.

1l. Synthetic data, and model collapse

When good human text runs short, models generate more: rewritten web pages, worked maths solutions checked by a program, question-and-answer pairs drawn out of textbooks. Used carefully this works well: it can turn a messy page into a clear one, and a checked answer is clean signal. Used carelessly it has a known failure.

Everyday picture Photocopy a photo, then photocopy the copy, and keep going. Each copy is nearly right, but fine detail disappears first, and after enough rounds you have a grey smudge. A model trained on its own outputs loses the rare, surprising parts of the data (the tails) first.

Tiny worked example Draw 20 numbers from a bell curve with spread 1. Fit a bell curve to those 20 numbers, draw 20 new numbers from the fit, fit again, and repeat. Each fit is a slightly narrower guess on average, and the narrowing compounds.

flowchart LR
  R[Real data<br/>spread 1.0] --> F1[Fit model 1]
  F1 --> S1[Sample from model 1]
  S1 --> F2[Fit model 2]
  F2 --> S2[Sample from model 2]
  S2 --> FN[... model n]
  FN --> X[Spread shrinks:<br/>tails are forgotten first]

import math
n, generations = 20, 200
# (n − 1)/n per generation, compounded
shrink = ((n - 1) / n) ** generations
f"{shrink:.1e}"  # → '3.5e-05'
# the spread is the square root of the variance
round(math.sqrt(shrink), 4)  # → 0.0059

In code: recursive_gaussian_fit runs the fit-sample-refit loop and returns every generation's fitted spread.

Everyday picture To bake one cake you need the recipe, but to learn to bake you also need your notes on every attempt: what went wrong, how much to change, how confident you are in each change. Training a model is the same. The weights are the recipe; training also keeps a gradient for every weight (what to change) and the optimizer's running notes (how it has been changing), and those notes are bigger than the recipe.

Tiny worked example Take one parameter trained with Adam (see primer.ml.optimizers) in mixed precision, the standard recipe (section 4 explains the number formats):

params = 7e9
bytes_per_param = 2 + 2 + 4 + 4 + 4
bytes_per_param  # → 16
# M_state in GB
bytes_per_param * params / 1e9  # → 112.0
# inference needs only the bf16 weights
2 * params / 1e9  # → 14.0

Activations: the memory that grows with the batch

The backward pass needs the intermediate results of the forward pass (the activations) to compute gradients, so they are kept until it runs. Their size grows with the number of tokens in flight, not with the parameter count. Korthikanti et al. (2022) counted them for one transformer layer, with 16-bit activations:

L, s, b, h, a = 32, 4096, 1, 4096, 32
# A = L · s b h (34 + 5 a s / h)
round(L * s * b * h * (34 + 5 * a * s / h) / 1e9, 1)  # → 104.2
# without stored attention scores
round(L * s * b * h * 34 / 1e9, 1)  # → 18.3
# activation checkpointing keeps only each layer's 2-byte input
round(L * 2 * s * b * h / 1e9, 2)  # → 1.07

The last line is activation checkpointing (also called gradient checkpointing, Chen et al., 2016): keep only each layer's input, and during the backward pass re-run that layer's forward pass to rebuild what it needs. It trades about one extra forward pass (roughly a third more compute) for activation memory that no longer grows with depth.

In code: training_memory itemizes the 16 bytes per parameter, and activation_bytes is the activation formula, with store_scores=False for FlashAttention-style kernels and checkpointed=True for activation checkpointing.

Everyday picture A restaurant can serve more diners in three ways. Open identical kitchens that each cook whole meals (data parallelism). Split one dish across cooks working side by side on the same step, one chopping the left half of the onions and one the right (tensor parallelism). Or build an assembly line, with each cook doing one stage and passing the plate on (pipeline parallelism). Real training runs use all three at once.

3a. Data parallelism: identical copies, averaged gradients

Every GPU holds a full copy of the model and processes a different slice of the batch. After the backward pass, the GPUs average their gradients so all copies take the same step and stay identical. This works because the gradient of an average loss is the average of the gradients.

Tiny worked example A one-weight model predicts y = w·x, with loss the mean of (w·x − y)². Batch: x = 1, 2, 3, 4 with y = 2, 4, 6, 8, and w = 1. On one GPU the gradient is (2/4)·Σ x(wx − y) = (2/4)·(−1 − 4 − 9 − 16) = −15. Split across two GPUs: GPU 1 gets x = 1, 2 and computes −5; GPU 2 gets x = 3, 4 and computes −25. Their average is (−5 − 25)/2 = −15, the same.

def grad(xs, ys, w=1.0):
    # d/dw of mean((w x − y)²)
    return 2 / len(xs) * sum(x * (w * x - y) for x, y in zip(xs, ys))
grad([1, 2, 3, 4], [2, 4, 6, 8])  # → -15.0
g1, g2 = grad([1, 2], [2, 4]), grad([3, 4], [6, 8])
g1, g2  # → (-5.0, -25.0)
(g1 + g2) / 2  # → -15.0

The averaging step is an all-reduce: every GPU contributes a vector and every GPU receives the sum. Sending every gradient to one GPU would jam its network link, so the standard algorithm is the ring all-reduce.

flowchart LR
  G0[GPU 0<br/>chunks A B C D] -->|one chunk per step| G1[GPU 1<br/>chunks A B C D]
  G1 -->|one chunk per step| G2[GPU 2<br/>chunks A B C D]
  G2 -->|one chunk per step| G3[GPU 3<br/>chunks A B C D]
  G3 -->|one chunk per step| G0

def sent(G, S):
    # 2 (G − 1) / G · S
    return 2 * (G - 1) / G * S
sent(4, 1.0)  # → 1.5
sent(8, 14e9) / 1e9  # → 24.5
round(sent(1000, 14e9) / 1e9, 2)  # → 27.97

Data parallelism alone does not solve the memory problem: every GPU still holds all 16 bytes per parameter.

In code: linear_regression_gradient computes a batch's gradient, so the averaging identity can be checked on shards; ring_all_reduce simulates both phases and counts what each GPU sends; all_reduce_traffic is the formula.

3b. Sharding the state: ZeRO and FSDP

Everyday picture A reading group with one expensive textbook does not buy a copy each. They tear it into chapters, each keeps one, and whoever needs a chapter borrows it for the evening and hands it back.

In plain data parallelism, every GPU stores an identical copy of the optimizer state, the gradients and the weights: pure waste. ZeRO (the Zero Redundancy Optimizer) keeps data parallelism's split of the batch but gives each GPU only a 1/G shard of that state, in three stages. PyTorch's FSDP (fully sharded data parallel) implements the third.

flowchart TB
  subgraph S0["Stage 0: plain data parallel"]
    A0["every GPU: weights + gradients + optimizer state"]
  end
  subgraph S1["Stage 1: shard the optimizer state"]
    A1["every GPU: weights + gradients<br/>its 1/G of the optimizer state"]
  end
  subgraph S2["Stage 2: also shard the gradients"]
    A2["every GPU: weights<br/>its 1/G of gradients and optimizer state"]
  end
  subgraph S3["Stage 3 = FSDP: shard everything"]
    A3["every GPU: its 1/G of everything<br/>borrow each layer's weights just in time"]
  end
  S0 --> S1 --> S2 --> S3

psi, G = 7.5e9, 64
stages = [16 * psi, 4 * psi + 12 * psi / G, 2 * psi + 14 * psi / G, 16 * psi / G]
[round(m / 1e9, 1) for m in stages]  # → [120.0, 31.4, 16.6, 1.9]

The price is communication. Stages 1 and 2 cost the same traffic as a plain all-reduce; stage 3 adds an all-gather of the weights in the forward pass and again in the backward pass, about 1.5 times the traffic. Frameworks hide it by fetching the next layer's weights while the current layer computes.

In code: zero_memory_per_gpu gives the state per GPU for any stage and GPU count.

3c. Tensor parallelism: splitting one matrix multiply

Everyday picture Two people fill in one large multiplication table: one does the left half of the columns, the other the right half. Neither needs the other's work until the end, when the halves are placed side by side.

The biggest operations in a transformer are matrix multiplies, X·W. They can be split across GPUs in two ways, and both give exactly the unsplit answer.

Tiny worked example X = [1, 2] and

W = [[1, 2, 3, 4], [5, 6, 7, 8]], so X·W = [11, 14, 17, 20].

• By columns: GPU 1 holds W's first two columns and computes [11, 14]; GPU 2 holds the last two and computes [17, 20]. Place side by side: [11, 14, 17, 20].
• By rows: GPU 1 holds W's first row and X's first entry: 1 × [1, 2, 3, 4] = [1, 2, 3, 4]. GPU 2 holds the second: 2 × [5, 6, 7, 8] = [10, 12, 14, 16]. Add: [11, 14, 17, 20].

X = [1, 2]
W = [[1, 2, 3, 4], [5, 6, 7, 8]]
def matmul(x, w):
    return [sum(x[i] * w[i][j] for i in range(len(x))) for j in range(len(w[0]))]
matmul(X, W)  # → [11, 14, 17, 20]
# by columns: each GPU computes half the output columns
matmul(X, [row[:2] for row in W]) + matmul(X, [row[2:] for row in W])  # → [11, 14, 17, 20]
# by rows: each GPU computes a partial sum of every output
p1, p2 = matmul(X[:1], W[:1]), matmul(X[1:], W[1:])
[a + b for a, b in zip(p1, p2)]  # → [11, 14, 17, 20]

Megatron-LM combines the two for a transformer's feed-forward layer, which is Y = activation(X·W₁)·W₂: split W₁ by columns and W₂ by rows.

flowchart LR
  X[X, full copy<br/>on both GPUs] --> A1["GPU 1: X · W1 left columns"]
  X --> A2["GPU 2: X · W1 right columns"]
  A1 --> R1[ReLU, locally] --> B1["· W2 top rows<br/>partial sum"]
  A2 --> R2[ReLU, locally] --> B2["· W2 bottom rows<br/>partial sum"]
  B1 & B2 --> AR[All-reduce:<br/>add the partial sums] --> Y[Y, full copy<br/>on both GPUs]

In code: column_parallel_matmul and row_parallel_matmul are the two splits, and tensor_parallel_mlp is the Megatron-LM feed-forward layer with its single all-reduce.

3d. Pipeline parallelism, and the bubble

Everyday picture A car assembly line with four stations. The first car takes four steps to roll off the end, and while it travels, stations further down stand idle. Only when many cars are on the line at once is every station busy.

Pipeline parallelism gives each GPU a consecutive block of layers (a stage). To keep the stages busy, the batch is cut into micro-batches that follow each other down the line. The idle time at the start and end is the pipeline bubble.

flowchart LR
  MB[Micro-batches<br/>1, 2, 3, ...] --> S1[GPU 1<br/>layers 1-8]
  S1 -->|activations| S2[GPU 2<br/>layers 9-16]
  S2 -->|activations| S3[GPU 3<br/>layers 17-24]
  S3 -->|activations| S4[GPU 4<br/>layers 25-32]
  S4 -.->|gradients flow back| S1

def bubble(p, m):
    # (p − 1) / (m + p − 1)
    return (p - 1) / (m + p - 1)
[round(bubble(4, m), 3) for m in (1, 8, 32)]  # → [0.75, 0.273, 0.086]

In code: pipeline_schedule builds the grid (positive numbers forward, negative backward, 0 idle) and bubble_fraction is the formula.

3e. All three at once

flowchart TB
  subgraph DP["Data parallel: replicas see different data, all-reduce gradients"]
    subgraph R1["Replica 1"]
      direction LR
      P1["Pipeline stage 1<br/>one server: 8 GPUs, tensor parallel"] --> P2["Pipeline stage 2<br/>one server: 8 GPUs, tensor parallel"]
    end
    subgraph R2["Replica 2"]
      direction LR
      Q1["Pipeline stage 1<br/>8 GPUs, tensor parallel"] --> Q2["Pipeline stage 2<br/>8 GPUs, tensor parallel"]
    end
  end

Everyday picture A carpenter measures a room with a tape measure and a table leg with calipers. Using calipers for everything would be slow; using the tape measure for everything would give wobbly tables. Mixed precision does each job with the coarsest number format that is good enough: the heavy matrix multiplies in 16 (or 8) bits, and the few places where tiny differences matter in 32.

The payoff is large. Halving the bits halves the memory and the data moved, and GPU tensor cores run 16-bit maths many times faster than 32-bit (see primer.ml.hardware).

4a. What a floating-point number is

Everyday picture Scientific notation, in binary. "6.02 × 10²³" has a sign, a few significant digits and an exponent. A float is the same three parts in bits: the exponent sets the range (how large or small a number can be), and the fraction bits set the precision (how many significant digits it keeps).

Tiny worked example In bf16, 1/3 is stored as sign 0, exponent field 125 and fraction field 0101011 in binary (43). The value is 2^(125 − 127) × (1 + 43/128) = 0.25 × 1.3359375 = 0.333984375. Seven fraction bits keep only about three significant decimal digits, so bf16 cannot tell 1/3 from 0.33398.

sign, E, bias, F, m = 0, 125, 127, 43, 7
# (−1)^sign × 2^(E − bias) × (1 + F / 2^m)
(-1) ** sign * 2.0 ** (E - bias) * (1 + F / 2 ** m)  # → 0.333984375
# the gap after 1 (the "epsilon") in bf16 and fp32
2.0 ** -7, 2.0 ** -23  # → (0.0078125, 1.1920928955078125e-07)

The formats that matter for training:

Below the smallest normal number a format has a few subnormal values that trade precision for extra range (fp16 reaches down to 6.0 × 10⁻⁸), and below half of the smallest subnormal a number becomes 0: underflow. Above the largest value is overflow, which becomes infinity.

In code: FloatFormat describes a format by its bit counts (with FP32, BF16, FP16, FP8_E5M2 and FP8_E4M3 defined), and quantize rounds any number to the nearest value a format can hold, reproducing overflow, subnormals and underflow.

4b. Underflow, and loss scaling

Everyday picture A kitchen scale that reads to the nearest gram shows 0 for a pinch of saffron. Weigh the saffron together with a known 1 kg jar, subtract the jar afterwards, and the pinch shows up.

Gradients are often tiny. Many are smaller than fp16's smallest subnormal (about 6 × 10⁻⁸), so a backward pass in fp16 silently turns them to 0 and those weights stop learning. Loss scaling is the jar: multiply the loss by a large number S before the backward pass, so every gradient is S times larger (the chain rule passes the factor through unchanged); then divide by S in fp32, before the update.

import numpy as np
grad, S = 1e-8, 65536
# without scaling: fp16 flushes it to zero
float(np.float16(grad))  # → 0.0
# with scaling: fp16 holds S · grad, then divide in higher precision
f"{float(np.float16(S * grad)) / S:.4e}"  # → '9.9972e-09'

A fixed scale is fragile: too small and gradients underflow, too large and they overflow to infinity. Dynamic loss scaling adapts it: if any gradient is infinite or not-a-number, skip the step and halve S; after a long run of clean steps (2,000 is common), double S.

flowchart TB
  M[fp32 master weights] -->|cast| W16[16-bit weights]
  W16 --> F[Forward pass<br/>16-bit matmuls]
  F --> L[Loss, in fp32]
  L -->|times S| B[Backward pass<br/>16-bit gradients]
  B --> CK{Any inf or NaN?}
  CK -- yes --> SK[Skip the step<br/>halve S]
  CK -- no --> U[Divide by S in fp32<br/>clip, then Adam update]
  U --> M
  SK --> M

In code: scaled_gradient_roundtrip scales, stores and unscales one gradient, and DynamicLossScaler.update skips the step and halves the scale on overflow, or doubles it after a long enough run of clean steps.

4c. Why the master copy stays in fp32

Everyday picture Pour a teaspoon of water into a full bathtub and measure with a bucket: the level has not changed, as far as the bucket can tell. Do it a thousand times and you have added four litres, but every single measurement still reads "no change".

Tiny worked example A weight of 1.0 receives an update of +0.0001 per step. Next to 1.0, bf16 can only step in increments of 0.0078, so 1.0001 rounds straight back to 1.0. After 1,000 updates the bf16 weight is still exactly 1.0; an fp32 weight has moved to 1.1, as it should.

That is why the optimizer keeps an fp32 master copy of the weights and applies updates to it, even though the matrix multiplies use 16-bit copies. The 16-bit copy is re-made from the master after every step.

In Python:

import numpy as np
w16, w32 = np.float16(1.0), np.float32(1.0)
for _ in range(1000):
    w16 = np.float16(w16 + np.float16(1e-4))
    w32 = np.float32(w32 + np.float32(1e-4))
# every update lost in 16 bits, all kept in 32
float(w16), round(float(w32), 4)  # → (1.0, 1.1)

In code: accumulate_updates adds the same update many times, rounding the weight to a chosen format after each step.

4d. fp8: the next halving

Everyday picture A ruler with only eight marks is useless for measuring a hair, unless you first slide it under a magnifying glass set to the right zoom. fp8 is that short ruler, and a per-tensor scale is the magnifying glass.

Recent GPUs multiply 8-bit floats at twice the 16-bit rate. With so few bits, one format cannot cover everything, so training uses two: E4M3 (more precision, range to 448) for weights and activations, and E5M2 (more range, to 57,344) for gradients. Because the range is so short, every tensor (or every small block of a tensor) gets its own scale factor, chosen from its recent largest value, the same idea as loss scaling applied everywhere. DeepSeek-V3 was trained largely in fp8 this way, keeping sensitive parts (normalizations, the optimizer, the master weights) in higher precision.

Everyday picture A ship on a months-long voyage does not assume the sea stays calm. It trims the sails when gusts come (clipping), leaves port slowly (warmup), keeps a log of its position (checkpoints), and has a drill for when something goes wrong (rolling back).

A large run does see storms. The loss curve occasionally jumps upward, a loss spike, sometimes recovering by itself and sometimes diverging for good. Causes include a batch of bad data, a learning rate a little too high for the model's current state, attention logits growing without bound, and numerical overflow. And separately from the maths, hardware fails: in the Llama 3 405B run, 419 unexpected interruptions occurred over 54 days, about one every three hours.

flowchart LR
  T[Train step] --> W{Loss well above<br/>recent median?}
  W -- no --> CP{Checkpoint<br/>interval reached?}
  CP -- yes --> SV[Save weights,<br/>optimizer, data position]
  CP -- no --> T
  SV --> T
  W -- "yes, and it persists" --> RB[Roll back to a<br/>checkpoint before the spike]
  RB --> SKIP[Skip the batches<br/>around the spike]
  SKIP --> T

5a. Spotting a spike

Tiny worked example Losses 3.0, 2.9, 2.8, 2.8, 2.7, then 8.5. The median of the previous four is 2.8, and 8.5 is more than twice that, so step 5 is flagged. The next step, 4.0, is compared with the median of 2.8, 2.8, 2.7, 8.5, which is still 2.8: one spike does not raise the bar, which is why the rule uses the median and not the mean.

import statistics
losses = [3.0, 2.9, 2.8, 2.8, 2.7, 8.5, 4.0, 2.7]
w, f = 4, 2.0
# L_t > f · median(L_{t−w}, …, L_{t−1})
[t for t in range(w, len(losses)) if losses[t] > f * statistics.median(losses[t - w:t])]  # → [5]

In code: detect_spikes applies the median rule to a whole loss history.

5b. Gradient clipping, when the gradient is spread across GPUs

Clipping by global norm (built in primer.ml.optimizers) caps the length of the whole gradient vector: if it is longer than c, shrink every entry by the same factor. At scale there is a twist. With ZeRO or tensor parallelism no GPU holds the whole gradient, so none can measure its length alone. Each GPU sums the squares of its own shard, one all-reduce adds those sums (a single number per GPU, so it is nearly free), and every GPU takes the square root and applies the same factor.

import math
shards = [[1.0, 2.0], [2.0, 4.0]]
# each GPU's local sum of squares
local = [sum(x * x for x in s) for s in shards]
local  # → [5.0, 20.0]
# all-reduce the sums, then the square root
norm = math.sqrt(sum(local))
norm  # → 5.0
c = 1.0
factor = min(1.0, c / norm)
[[x * factor for x in s] for s in shards]  # → [[0.2, 0.4], [0.4, 0.8]]

In code: sharded_global_norm computes the norm from per-GPU sums of squares, and train_through_a_bad_batch trains through a corrupted batch with or without primer.ml.optimizers.clip_by_global_norm.

5c. Warmup

Everyday picture Nobody floors the accelerator in a car they have never driven; they ease on until they know how it responds.

At the start of training the weights are random and Adam's running averages have seen only a handful of gradients, so its step sizes are unreliable. A full learning rate at step 1 can throw the model into a region it never recovers from. Warmup ramps the learning rate linearly from 0 to its peak over the first few hundred to few thousand steps; primer.ml.optimizers derives the warmup-then-cosine schedule with worked numbers. At scale it matters more, not less: bigger models and bigger batches tolerate smaller peak learning rates, and a spike early in a months-long run wastes the most.

5d. Checkpoints: how often to save

Everyday picture Saving a long document every few seconds wastes time on saving; saving once an hour risks losing an hour of work to a crash. Somewhere in between is the least total waste.

Tiny worked example Suppose writing a checkpoint pauses training for 1 minute, and the cluster suffers a failure every 180 minutes on average. Saving every 19 minutes spends 1/19 of the time saving, and each failure loses on average half an interval, 9.5 minutes, every 180 minutes. Both costs come to about 5.3%, and their total, 10.5%, is the smallest possible.

import math
C, M = 1.0, 180.0
# T* = √(2 C M)
T = math.sqrt(2 * C * M)
round(T, 2)  # → 18.97
# waste(T) = C/T + T/(2M)
round(C / T + T / (2 * M), 3)  # → 0.105
# a 10-second save
C = 1 / 6
round(math.sqrt(2 * C * M), 1), round(C / math.sqrt(2 * C * M) + math.sqrt(2 * C * M) / (2 * M), 3)  # → (7.7, 0.043)

In code: wasted_fraction is the waste formula and optimal_checkpoint_interval is Young's square-root rule.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.