Two loops, one function: what moves, what is frozen, and what each costs. Capstone for Part 1.
This is the last chapter of Part 1 and it adds no new component. Everything is built: tokenizer, embedding, blocks, unembedding, loss, gradient, update. What remains is to see the two ways the machine is used, with all the pieces in place, and to feel the asymmetry between them. Then the capstone: one complete training step traced by hand through a toy transformer, and the first working version of Dispatch.
The question this chapter answers: when a frontier model is trained and when it answers you, what exactly is happening, what is different, and why does that difference shape everything in Parts 2 and 4?
An orchestra in rehearsal. The conductor stops them every few bars: "violas, you were late; brass, too loud". Every section adjusts at once, from one hearing of one passage. The score is fixed; the players change. They play the passage again, better, and the conductor stops them again. Weeks of this.
Then the concert. The conductor says nothing. The players do not change. They play the piece from the beginning to the end, one bar at a time, each bar following from the last, and if a bar goes wrong there is no going back: the next bar simply follows from the wrong one. The audience hears only the concert. They never see the rehearsal, but everything they hear was decided in it.
Rehearsal is training: the score (the data) is fixed, the players (the parameters) adjust, and the whole passage is corrected at once from one hearing. The concert is inference: the players are frozen, the music unfolds one bar at a time, and every bar depends on the ones already played. The one thing this analogy gets slightly wrong is important enough to say now: in rehearsal, the orchestra hears the correct passage at every bar, not its own mistakes. That is teacher forcing, and it is the first thing this chapter explains.
| In the picture | In the machine | The word we will use |
|---|---|---|
| Rehearsal | Forward on real text, loss at every position, backward, update | training step |
| Playing to the written score, not to your own last bar | Every position's prediction is conditioned on the true preceding tokens | teacher forcing |
| Correcting all sections from one hearing | All T positions billed and backpropagated in one pass | parallel training |
| The concert | Frozen θ, one token appended per pass | inference, autoregressive decoding |
| The upbeat before the first note | Processing the whole prompt in one parallel pass | prefill |
| Each bar following the last | One forward pass per generated token | decode |
| Not re-reading the score from the top each bar | Keep every past token's keys and values instead of recomputing them | KV cache |
| The soloist's freedom within the score | How a token is chosen from the distribution | sampling, temperature, top-k, top-p |
Take a batch: B sequences of T tokens each, drawn from the corpus, shape [B, T]. Run one forward pass. Because of the causal mask (Chapter 3), the output at position i depends only on positions ≤ i, so the model's prediction for "what comes after position i" is legitimate at every i simultaneously. The pass produces B·T distributions, and the true next token is known at every position: it is just the text shifted by one. Bill all B·T predictions, average, backpropagate once, update once.
The consequence to absorb is that a sequence of 8,000 tokens is 8,000 training examples for the price of one forward-backward pass, and the model never sees its own output during pretraining. At position i it is conditioned on the true tokens 1 … i, whatever it would have predicted for them. This is teacher forcing. It is why training parallelises so well, and it has a cost that surfaces later: a model that has only ever seen correct prefixes can drift when, at inference, its own errors become its context. Chapters 14 to 16 are partly about training on the model's own outputs to close that gap.
one training step at frontier scale (Llama 3 405B, public): batch [B, T] = up to 16M tokens, e.g. 2,048 sequences × 8,192 tokens forward 16M distributions over 128k tokens ≈ 2N × 16M = 1.3 × 10¹⁹ ops loss 16M bills, averaged backward gradient for all 405B parameters ≈ 4N × 16M = 2.6 × 10¹⁹ ops update Adam, every parameter, once ≈ 1 million such steps ≈ 3.8 × 10²⁵ ops total
Now the concert. A prompt arrives: the system prompt, the conversation so far, the new message, tokenized to P ids. Two phases follow, and they are different enough that serving systems treat them as different workloads (Chapter 19).
Prefill. All P prompt tokens go through the model in one parallel pass, exactly like a training forward pass without the loss. This produces the distribution at the last position, from which the first output token is sampled, and, as a by-product, the keys and values of every prompt token at every layer.
Decode. The sampled token is appended. Now the model needs the distribution at position P+1. Naively that is a fresh forward pass over P+1 tokens. But the first P tokens have not changed, θ has not changed, and the causal mask means their keys and values do not depend on anything after them. So their K and V at every layer are exactly what they were during prefill. Keep them. Each decode step then processes one token: compute its q, k, v, attend over the stored keys and values plus its own, run the MLP, unembed, sample, append its k and v to the store. That store is the KV cache.
The cache turns decoding from quadratic in the reply length to linear, and it is universal: no serving system generates without one. It also creates the memory problem that dominates serving. Every token in every live conversation holds 2 × L × (KV heads) × d numbers in the cache, and a busy server holds thousands of conversations.
Why is decode slow even though it does little arithmetic? One decode step for one conversation is one token through the model: 2N ≈ 0.8 × 10¹² operations for the 405B, which a GPU does in about a millisecond. But to do them, every one of the 405 billion weights has to be read from memory once, because each token touches every matrix.
weights to read per decode step 810 GB (405B × 2 bytes) GPU memory bandwidth (top, 2025) ≈ 5–8 TB/s per GPU; the weights span ≥ 5 GPUs time to stream the weights once ≈ 810 / (5 × 6,000) ≈ 27 ms → at most ≈ 37 tokens/s for ONE conversation arithmetic time for that token ≈ 0.8 × 10¹² / (5 × 10¹⁵) ≈ 0.2 ms → the GPUs are 99% idle
Decode is memory-bound: the weights are read far more slowly than they are used. The fix is to serve many conversations at once, reading the weights once per step for all of them; that is batching, and it is why a provider's throughput and your personal latency are different numbers. Prefill, by contrast, pushes thousands of tokens through the same weights in one read, and is compute-bound. Chapter 19 is built on this contrast.
The model returns a distribution; something has to pick. Chapter 1 introduced greedy and sampling and the temperature; here is the full set of standard knobs, all of which act on the distribution after the model has produced it.
T: divide the logits by T before softmax. Below 1 sharpens toward the top token; above 1 flattens toward uniform. T → 0 is greedy.k most probable tokens, zero the rest, renormalise. Removes the long tail of nonsense that has 0.01% each but adds up.p, then renormalise. Adapts to the shape: a sharp distribution keeps one or two tokens, a flat one keeps many.p times the top token's. Another adaptive truncation, popular in open-model serving.max_tokens: when to end. The model ends its own turn by sampling an end-of-turn special token; the caller can end it earlier with a token budget. That is the stop_reason of Chapter 1.None of these change the model. Two consequences. First, "the model is creative / repetitive / random" is often a statement about the sampler settings, and changing them is free. Second, the newest frontier APIs increasingly fix these knobs on the provider side public: the lab has tuned sampling as part of post-training and does not expose it. The distribution is still there; you just no longer choose how to draw from it.
| Training step | Inference (one conversation) | |
|---|---|---|
| Parameters θ | change every step | frozen |
| Input | real text, all positions at once | prompt, then the model's own tokens one at a time |
| What each position sees | the true preceding tokens (teacher forcing) | the prompt plus what was sampled so far |
| Passes | one forward, one backward, per batch | one forward per generated token, no backward |
| Kept between passes | nothing (activations are freed after backward) | the KV cache |
| Memory per parameter | ≈ 16 bytes (weights, gradients, Adam state) | 2 bytes (or 1, quantised; Chapter 27) |
| Cost per token | ≈ 6N operations | ≈ 2N operations, plus reading all weights once per step |
| Randomness | which batch, data order, dropout if any | the sampler |
| Bottleneck | compute, and communication between GPUs | memory bandwidth, and cache size |
| What changes the model's behaviour | the data and the loss | only the context |
The last row is the one most people get wrong. Nothing you type in a conversation changes θ. The assistant that "remembers" your name from earlier in the chat is attending, through the KV cache, to tokens that are still in its context. Close the conversation and it is gone. When a provider says a model has been updated, that means someone re-entered the left column: new data, new bills, new θ. Part 3 is about the specific left-column loops that turn a pretrained model into an assistant. Fine-tuning, in any of its forms, is nothing more than re-entering the left column with different data and, usually, a lower learning rate.
Everything from Chapters 1 to 5 in one place, small enough to check. The toy model has V = 6 tokens, width C = 4, one block with a single head of width 2 and an MLP of hidden width 4. The sequence is "the cat sat on" and the targets are "cat sat on mat". The script code/ch07/paper_trace.py prints every intermediate; the walkthrough below is its output with the chapter that explains each step.
$ python code/ch07/paper_trace.py| Step | What the script prints | Where it was explained |
|---|---|---|
| 1 · embed | x0 = W_E[ids], a [4, 4] matrix: four tokens, four coordinates each | Ch 2 §2.3, a lookup |
| 2 · attend | attention weights, lower-triangular, rows summing to 1: [1, 0, 0, 0], [.564, .436, 0, 0], … | Ch 3: scores, √d, mask, softmax |
| 3 · add | x1 = x0 + attention: the stream after its first note | Ch 4 §4.2, the residual |
| 4 · MLP | the ReLU gate matrix, e.g. [0 1 1 0] for token 1: two of four detectors fired | Ch 2 §2.6 |
| 5 · add again | x2 = x1 + mlp | Ch 4 |
| 6 · unembed, softmax | logits [4, 6], then p with each row summing to 1 | Ch 1 §1.2–1.3 |
| 7 · bill | loss = mean −ln p[target] = 2.0911; uniform would be ln 6 = 1.79, so the random model is slightly worse than guessing | Ch 1 §1.4 |
| 8 · gradient on logits | (p − onehot)/T: every row has one negative entry (the target) and small positives elsewhere | Ch 5 §5.4 |
| 9 · backward through everything | gradient norms for all eight matrices; the two residual adds pass the gradient straight through | Ch 5 §5.3–5.4 |
| 10 · check | numeric dL/dW_U[1,2] = 0.04930 analytic = 0.04930 | Ch 5, the ε nudge |
| 11 · update, re-run | loss 2.0911 → 1.7221 | Ch 5 §5.2 |
attention weights (rows = queries)
[[1. 0. 0. 0. ]
[0.564 0.436 0. 0. ]
[0.485 0.206 0.31 0. ]
[0.519 0.153 0.223 0.104]]
…
loss = mean −ln p[target] = 2.0911
gradient on the logits (p − onehot)/T
[[ 0.019 -0.234 0.026 0.004 0.075 0.111]
[ 0.017 0.016 -0.216 0.01 0.093 0.081]
[ 0.012 0.007 0.068 -0.229 0.075 0.067]
[ 0.015 0.01 0.038 0.01 -0.17 0.097]]
gradient norms per matrix: {'W_E': 1.964, 'W_Q': 0.856, 'W_K': 0.716, 'W_V': 0.345, 'W_O': 0.31, 'W_in': 0.212, 'W_out': 0.842, 'W_U': 0.934}
numeric dL/dW_U[1,2] = 0.04930 analytic = 0.04930
after one step (lr=0.5): loss 2.0911 → 1.7221
The exercise is to read the script top to bottom with a pen, reproducing at least the attention weights for the second row ([.564, .436, 0, 0]) and the gradient on the logits for the first row from the printed p. If you can do that, you can trace a frontier model's step in your head; the only differences are the shapes and the number of blocks.
The builder spine gets its first real program: a chat loop. It is fifteen lines, and every line is one of the concepts of this chapter made concrete: the conversation history is the context, the API streams tokens as the decode loop produces them, the usage numbers separate prefill from decode, and nothing you type reaches θ.
# code/ch07/dispatch_chat.py
import anthropic
client = anthropic.Anthropic()
SYSTEM = """You are Dispatch, the on-call assistant for Postbox, a notification delivery
service. Answer briefly and concretely. If you do not know something about Postbox, say so."""
history = [] # the conversation: the only thing that changes between turns
total_in = total_out = 0
while True:
user = input("\nyou > ").strip()
if user.lower() in {"quit", "exit"}: break
history.append({"role": "user", "content": user})
with client.messages.stream(model="claude-opus-5", max_tokens=1024, system=SYSTEM, messages=history) as stream:
print("dispatch > ", end="", flush=True)
for text in stream.text_stream: # tokens arrive one decode step at a time
print(text, end="", flush=True)
response = stream.get_final_message()
history.append({"role": "assistant", "content": response.content})
total_in += response.usage.input_tokens; total_out += response.usage.output_tokens
print(f"\n [prefill {response.usage.input_tokens} tokens · decode {response.usage.output_tokens} tokens"
f" · session so far: {total_in} in, {total_out} out]")(example session — your text will differ; the token accounting is the point) you > Pager: push-queue depth > 50k. Where do I start? dispatch > Three checks, in order: (1) is depth growing or draining over the last 15 min; (2) push-worker error rate and restarts; (3) inbound event rate for a spike. Tell me which you can see. [prefill 78 tokens · decode 61 tokens · session so far: 78 in, 61 out] you > Depth is growing. Workers look healthy. dispatch > Then consumers are fine and producers are ahead of them. Check inbound rate first; if it spiked, that's the cause and you need to scale workers or shed load. [prefill 158 tokens · decode 43 tokens · session so far: 236 in, 104 out] you > What did I say the pager was about? dispatch > A push-queue depth alert above 50k. [prefill 218 tokens · decode 12 tokens · session so far: 454 in, 116 out]
Read the prefill numbers. Turn two's prefill is turn one's prompt plus turn one's reply plus the new message: the context grows by everything said. Turn three "remembers" the pager because the pager is still in the context, at prefill cost, every turn. A provider can avoid re-running the prefill on the unchanged prefix by keeping its KV cache between calls (prompt caching, Chapter 20); the model itself remembers nothing. Dispatch will gain tools in Chapter 20, retrieval in Chapter 21, and an agent loop in Chapter 22, and every one of those additions is more context, never a change to θ.
Train without teacher forcing: feed the model its own samples as context. Positions can no longer be computed in parallel, because position i+1's input is not known until position i has been sampled. Training becomes sequential like decoding, thousands of times slower. Some post-training methods do exactly this on purpose, for small amounts of data, to correct the drift teacher forcing leaves behind (Chapter 16); nobody does it for pretraining.
Decode without the KV cache. Each token costs a forward pass over the entire context. A 1,000-token reply after a 10,000-token prompt costs about 10.5 million token-passes instead of 11 thousand. Every serving system would be a thousand times slower.
Serve one conversation per GPU. Decode reads 810 GB of weights per token to do 0.2 ms of arithmetic; the GPUs are idle 99% of the time. Batching a hundred conversations per step reads the weights once for all of them; throughput rises a hundredfold for nearly the same time per step. Chapter 19 makes this the central fact of serving.
Let the chat client change θ after each turn. Every user's conversation would nudge the shared model; a bad conversation would degrade it for everyone; the model would forget its training in favour of whatever was said most recently; and reproducibility would be gone. The wall between the two loops is a feature. Personalisation is done in the context, or by training a separate copy under controlled conditions.
Say back the whole machine. Text is split by a learned merge list into tokens, and a token is an integer. The integer selects a row of the embedding table, a vector of C numbers in a space where direction means something. That vector is the residual stream. It passes through L blocks. In each block, a normalised copy is read by H attention heads, each of which projects every token into a query, a key, and a value, scores each query against every earlier key, softmaxes the scaled scores into weights, and blends the values; the heads are concatenated, mixed by W_O, and added to the stream. Then a normalised copy is read by the MLP, which runs thousands of learned detectors and adds their write-backs to the stream. Positions enter as rotations of the queries and keys. After the last block, a final norm and the unembedding turn the stream into V logits; softmax turns those into a distribution over the next token. That distribution is the model's only output.
In training, a batch of real text goes through once, every position's distribution is billed by −ln p(true next token), the bills are averaged, and the chain rule walks backwards through the same graph delivering each matrix a gradient of its own shape: p − onehot at the logits, outer products at the matmuls, gates at the ReLUs, straight through at the residual adds. Adam nudges every parameter against its gradient, a million times, on fifteen trillion tokens. In inference, the parameters are frozen; the prompt is prefilled in one pass, then one token per pass is sampled and appended, with the keys and values of the past kept in a cache so nothing is recomputed. Decode is bound by reading the weights, not by arithmetic. Beacon is this machine with 405 billion numbers, and Dispatch is a loop that sends it context and prints what it samples.
p for the second row of the paper trace ([0.068, 0.065, 0.135, 0.038, 0.371, 0.323]) and its target index 2, compute the gradient on that row's logits, divide by T = 4, and confirm it matches the script. Then compute the loss contribution of that position and state whether it is above or below the mean.dispatch_chat.py for four turns. Record the prefill count each turn and fit it to "previous prefill + previous decode + new message". Then add a max_tokens=20 limit and observe the stop_reason. Finally, modify the loop to drop the oldest turn once the history exceeds 2,000 tokens and note what Dispatch forgets.