Building an LLM · Draft

9 — What changed since GPT-2

Chapter 8 ended with a GPT-2 that follows instructions, in the sense that it answers in the right format and stops. Its answers were often wrong, and tuning can’t supply knowledge the pretrained model lacks; the limit is the model underneath. GPT-2 came out in 2019. The open models that people run today, Meta’s Llama 3 and Alibaba’s Qwen3 among them, are better mostly because they’re bigger and were trained on far more text. But their blocks are also built a little differently.

The differences are smaller than the gap in capability suggests. A Llama or Qwen block still normalizes, attends, adds, normalizes, runs a feed-forward network, and adds. What changed is the details inside those steps, plus a few tricks that make generation cheaper. This chapter goes through them one at a time, each written as a diff against the chapter 4 code, checked against the paper that introduced it, and run where running it shows something.

Everything ran on an Apple M3 Pro with 18 GB of memory, using Python 3.12 and PyTorch 2.15.0, on the GPU through MPS unless the CPU is named. Seeds are given with each run.

Normalizing without the mean: RMSNorm

The chapter 4 block normalizes every vector twice, once before attention and once before the feed-forward network. Layer norm does two things to a vector: it subtracts the mean, so the numbers are centered on zero, and it divides by the standard deviation, so their spread is 1. Then a learned scale and a learned shift are applied.

RMSNorm keeps only the second half. It divides each vector by its root mean square, the square root of the average of the squared entries, and multiplies by a learned scale. No mean is subtracted and there is no shift.

RMSNorm(x)i=gi⋅xi1d∑j=1dxj2+ϵ\text{RMSNorm}(x)_i = g_i \cdot \frac{x_i}{\sqrt{\frac{1}{d}\sum_{j=1}^{d} x_j^2 + \epsilon}}

On a four-number vector the difference is easy to see. Take x = [2, 4, 6, 8]. Its mean is 5 and its variance is 5, so layer norm gives

layer norm: [-1.342, -0.447,  0.447,  1.342]

The squares average to (4 + 16 + 36 + 64) / 4 = 30, so the root mean square is √30 ≈ 5.477, and RMSNorm divides by that:

RMSNorm:    [ 0.365,  0.730,  1.095,  1.461]

The layer-norm output is centered; the RMSNorm output is not. Add 10 to every entry and the contrast sharpens. Layer norm returns exactly the same four numbers, because subtracting the mean removes any constant shift. RMSNorm returns [0.791, 0.923, 1.055, 1.187]: the vector is still rescaled so its root mean square is 1, but the shift is kept in it.

Zhang and Sennrich (2019) proposed this on the hypothesis that the re-centering in layer norm isn’t what makes it useful, and that the rescaling is. They reported comparable quality and running time reduced by 7% to 64% depending on the model. The saving comes from skipping the mean and the shift: one fewer pass over the vector and one fewer parameter vector.

As a diff against the chapter 4 LayerNorm:

-class LayerNorm(nn.Module):
+class RMSNorm(nn.Module):
     def __init__(self, d, eps=1e-5):
         super().__init__()
         self.eps = eps
         self.weight = nn.Parameter(torch.ones(d))
-        self.bias = nn.Parameter(torch.zeros(d))
 
     def forward(self, x):
-        mean = x.mean(dim=-1, keepdim=True)
-        var = x.var(dim=-1, keepdim=True, unbiased=False)
-        return self.weight * (x - mean) / torch.sqrt(var + self.eps) + self.bias
+        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return self.weight * x / rms

x.pow(2) squares every entry, and .mean(dim=-1, keepdim=True) averages along the last dimension (the 768 numbers of one token’s vector) while keeping that dimension as size 1, so the division broadcasts back over the vector.

Position as a rotation: RoPE

GPT-2 learns one vector per position, 1,024 of them, and adds the right one to each token’s embedding before the first block (chapter 2). That works, with two weaknesses. Position 1,025 has no row in the table, so the model simply can’t take a longer input. And the model sees absolute positions: “word at 500, word at 503” and “word at 10, word at 13” look unrelated to it, though for most of language what matters is that the two words are three apart.

Rotary position embeddings, from Su et al. (2021), drop the table. Nothing is added to the embeddings. Instead, inside every attention layer, each query and each key is rotated by an angle that grows with its position, just before the dot product.

The idea is clearest in two dimensions. A query q = [1, 0] and a key k = [0.6, 0.8] have dot product 0.6. Now rotate each by 0.5 radians per position, so a vector at position m is turned by 0.5m. Put the query at position 2 and the key at position 5:

q at 2:  [ 0.540,  0.841]
k at 5:  [-0.959, -0.282]
dot:     -0.7556

Move both seven places along, to 10 and 13, and the two vectors point somewhere else entirely, but the dot product doesn’t change:

q at 10:   [0.284, -0.959]
k at 13:   [0.414,  0.910]
dot:       -0.7556
q at 100, k at 103:  -0.7556

Rotating both vectors by the same extra angle doesn’t change the angle between them, and a dot product depends only on lengths and that angle. So the attention score between a query and a key depends on their contents and on how far apart they are, never on where the pair sits in the text. Swap the order (query at 5, key at 2) and the score is 0.8404, so direction is kept as well as distance. At equal positions the rotation cancels and the plain dot product, 0.6, comes back.

Real heads have 64 dimensions, not 2. RoPE splits them into 32 pairs and rotates each pair as above, each at its own speed. Pair i turns by θ_i radians per position, with

θi=10000−2i/d,i=0,1,…,d/2−1\theta_i = 10000^{-2i/d}, \quad i = 0, 1, \dots, d/2 - 1

For a 64-dimensional head the first pair turns 1 radian per position and completes a circle every 6.3 positions; the last turns 0.000133 radians per position and would need about 47,000 positions to go round once. Fast pairs can tell neighbors apart; slow pairs still change measurably across thousands of tokens. The same check as before, with random 64-dimensional vectors (seed 0):

query at 3,    key at 0:     -8.2811
query at 503,  key at 500:   -8.2811
query at 1503, key at 1500:  -8.2811
query at 0,    key at 3:    -10.7468
no rotation:                -11.4345

In code it is two small functions and three changed lines in attention. rope_angles builds a table of cos and sin values once; rope applies it. x[..., 0::2] takes the even-numbered dimensions and x[..., 1::2] the odd ones, so each pair is (x1, x2); the rotation is the usual 2-by-2 one, and torch.stack(..., dim=-1).flatten(-2) interleaves the results back into the original order.

def rope_angles(head_dim, n_pos, base=10000.0):
    i = torch.arange(head_dim // 2)
    theta = base ** (-2 * i / head_dim)                 # one frequency per pair of dimensions
    ang = torch.arange(n_pos)[:, None] * theta[None, :]  # (positions, head_dim/2)
    return ang.cos(), ang.sin()

def rope(x, cos, sin):
    """Rotate each (even, odd) pair of x's last dimension by its position's angle."""
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)
 class GPT(nn.Module):
     def __init__(self, cfg):
         ...
         self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.emb_dim)
-        self.pos_emb = nn.Embedding(cfg.context_length, cfg.emb_dim)
+        cos, sin = rope_angles(cfg.emb_dim // cfg.n_heads, cfg.context_length)
+        self.register_buffer("cos", cos, persistent=False)
+        self.register_buffer("sin", sin, persistent=False)

 class CausalSelfAttention(nn.Module):
-    def forward(self, x):
+    def forward(self, x, cos, sin):
         ...
         v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        q, k = rope(q, cos[:T], sin[:T]), rope(k, cos[:T], sin[:T])
         scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)

Two pieces of plumbing go with it. GPT.forward no longer adds pos_emb, and because nn.Sequential passes only one argument from layer to layer, the blocks become an nn.ModuleList that the forward pass loops over, handing cos and sin down to each attention layer. register_buffer stores a tensor on the module so it moves with .to(device), without making it a parameter the optimizer would train; persistent=False keeps it out of saved checkpoints, since it can always be recomputed. The values are never rotated, only queries and keys, because position is supposed to affect who attends to whom, not what gets passed along.

The base of 10,000 is the paper’s choice. Later models raise it to slow every pair down and stretch the usable range: Llama 3 uses 500,000, and Qwen3-0.6B’s published config uses 1,000,000. RoPE doesn’t make long context free. A model trained on 4,096-token sequences has never seen the angle combinations of much larger distances, and nothing guarantees it handles them; models reach long contexts by being trained further on long text, usually with a larger base. What RoPE removes is the hard wall of a learned table, and the need for the model to learn relative distance on its own.

A gated feed-forward network: SwiGLU

The feed-forward network in each block expands a 768-number vector to 3,072, applies GELU to every number, and projects back. Shazeer (2020) tested variants where the expansion is done twice, by two separate matrices, and one copy is passed through an activation and used to scale the other, number by number. The version later models chose uses the Swish activation, also called SiLU, x · sigmoid(x):

FFNSwiGLU(x)=(SiLU(xWgate)⊙xWup) Wdown\text{FFN}_{\text{SwiGLU}}(x) = \big(\text{SiLU}(x W_{\text{gate}}) \odot x W_{\text{up}}\big)\, W_{\text{down}}

The ⊙ is element-wise multiplication. SiLU looks like ReLU with the corner smoothed off and a small dip below zero:

x:     -4      -3      -2      -1     0     1      2      3      4
ReLU:   0       0       0       0     0     1      2      3      4
SiLU:  -0.072  -0.142  -0.238  -0.269  0     0.731  1.762  2.858  3.928

The “gated” part is what’s new. In the GELU network, each of the 3,072 hidden numbers is a fixed function of one projection of the input. In SwiGLU, each hidden number is a product of two projections, one of which decides how much of the other gets through, and both are learned. A single hidden unit can therefore pass a signal through, block it, or flip its sign, depending on the input.

Three matrices instead of two would mean 50% more parameters, so the paper shrinks the hidden size by a factor of 2/3 to keep the count the same. For a 768-wide model, 2/3 × 3,072 = 2,048:

GELU MLP   768 → 3072 → 768, two matrices:    4,718,592 weights
SwiGLU     768 → 2048 (×2) → 768, three:      4,718,592 weights

Exactly equal at this width. LLaMA states the same choice, “a dimension of 2/3 4d instead of 4d”. Larger models often round the hidden size to a convenient multiple or pick it separately; Llama 3 8B uses 14,336 for a 4,096-wide model, 3.5× rather than 2.67×.

-class FeedForward(nn.Module):
-    def __init__(self, cfg):
-        super().__init__()
-        self.fc = nn.Linear(cfg.emb_dim, 4 * cfg.emb_dim)
-        self.act = GELU()
-        self.proj = nn.Linear(4 * cfg.emb_dim, cfg.emb_dim)
-
-    def forward(self, x):
-        return self.proj(self.act(self.fc(x)))
+class SwiGLU(nn.Module):
+    def __init__(self, d, hidden):
+        super().__init__()
+        self.gate = nn.Linear(d, hidden, bias=False)
+        self.up = nn.Linear(d, hidden, bias=False)
+        self.down = nn.Linear(hidden, d, bias=False)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))

F.silu is PyTorch’s built-in x · sigmoid(x). Why the gate helps is not well explained. The paper’s own conclusion offers no explanation and, half-joking, credits “divine benevolence”. The evidence is empirical: its SwiGLU and GEGLU variants reached the lowest perplexity of the variants tested, and the result has held up across model families since.

Not recomputing the past: the KV cache

Every generation loop in this track so far has done something wasteful. To pick token 101, generate runs the whole model over tokens 1 to 100. To pick token 102, it runs over tokens 1 to 101, recomputing everything for the first 100 from scratch, and gets the same numbers as last time.

The waste is easy to locate. Because of the causal mask, the computation at position 5 depends only on tokens 1 to 5. Adding token 6 changes nothing at positions 1 to 5, in any layer. The only new work is position 6 itself: its query, key and value in each layer, and its attention over the keys and values of everything before it. Those earlier keys and values are exactly what the previous step computed and threw away.

A KV cache keeps them. Each block stores the keys and values it has produced so far, and each new step sends in a single token, computes that token’s q, k and v, appends k and v to the stored ones, and attends over the lot. Queries don’t need caching: an old query is only ever used to compute an old position’s output, which is already done.

It is not a post-GPT-2 invention. OpenAI’s own GPT-2 code passes a past tensor of stacked keys and values into every attention layer and returns a present one. It belongs in this chapter because the chapter 4 model left it out for clarity, and because GQA, next, exists mainly to shrink it.

Two changes to the chapter 4 attention are enough. The forward pass takes the previous cache and returns the new one, and the causal mask is sliced from the right place, because the new queries sit at positions start to start + T while the keys run from 0:

-    def forward(self, x):
+    def forward(self, x, cache=None):
         B, T, C = x.shape
         q, k, v = self.qkv(x).split(C, dim=-1)
         q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
         k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
         v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        if cache is not None:                  # append the new keys and values to the old ones
+            k = torch.cat([cache[0], k], dim=2)
+            v = torch.cat([cache[1], v], dim=2)
+        start = k.shape[2] - T                  # how many positions came before x
         scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)
-        scores = scores.masked_fill(self.mask[:T, :T], float("-inf"))
+        scores = scores.masked_fill(self.mask[start:start + T, :start + T], float("-inf"))
         weights = torch.softmax(scores, dim=-1)
         out = self.attn_drop(weights) @ v
         out = out.transpose(1, 2).reshape(B, T, C)
-        return self.proj(out)
+        return self.proj(out), (k, v)

torch.cat(..., dim=2) joins tensors along the sequence dimension; the cache for one block has shape (batch, heads, positions so far, 64). The model’s forward pass needs one more change: position embeddings must start at start, not 0, since the single new token is at position 101, not position 0.

def forward_cached(model, idx, caches):
    """Run only the new tokens in idx; caches holds one (k, v) per block, or None."""
    start = 0 if caches[0] is None else caches[0][0].shape[2]
    pos = torch.arange(start, start + idx.shape[1], device=idx.device)
    x = model.tok_emb(idx) + model.pos_emb(pos)
    new = []
    for block, cache in zip(model.blocks, caches):
        a, kv = attn_step(block.attn, block.norm1(x), cache)
        x = x + a
        x = x + block.ffn(block.norm2(x))
        new.append(kv)
    return model.head(model.final_norm(x)), new

@torch.no_grad()
def generate_cache(model, idx, n):
    caches = [None] * len(model.blocks)
    logits, caches = forward_cached(model, idx, caches)        # the whole prompt once
    for _ in range(n):
        nxt = logits[:, -1].argmax(-1, keepdim=True)
        idx = torch.cat([idx, nxt], dim=1)
        logits, caches = forward_cached(model, nxt, caches)    # then one token at a time
    return idx

attn_step is the attention forward pass above, written as a standalone function that takes the attention layer as its first argument, so it runs on the chapter 6 model’s layers without changing their class. @torch.no_grad() turns off gradient tracking, which generation never needs.

Both versions generate greedily from the same prompt, “The history of the bicycle begins”, with chapter 6’s GPT-2 weights. The first check is that they produce the same thing: for every length below, the cached and uncached runs give identical token sequences. The cache changes the cost, not the result. Timings are single runs after a short warm-up, made while holding the lock that keeps this track’s timed runs from overlapping:

New tokensCPU, no cacheCPU, cacheSpeed-upGPU, no cacheGPU, cache
10026.8 tokens/s138.1 tokens/s5.2×62.5 tokens/s117.1 tokens/s
50010.2 tokens/s130.8 tokens/s12.8×24.9 tokens/s108.8 tokens/s
9945.1 tokens/s110.6 tokens/s21.7×11.4 tokens/s40.3 tokens/s

994 is what fits in the 1,024-token context after the 6-token prompt. On the CPU the shape is exactly what the argument predicts. Without a cache, step n reprocesses n tokens, so the average speed falls steadily as the text grows: from 26.8 tokens per second over the first 100 to 5.1 over a full context, 196 seconds in total. With a cache each step processes one token, plus attention over the stored keys, which grows only slowly, so the speed stays above 100 tokens per second the whole way, and the full context takes 9 seconds.

The GPU numbers agree at 100 and 500 tokens and are odd at 994, where the cached run slows to 40 tokens per second. These timings were noisier than the CPU ones across repeated runs, and this chapter didn’t track the cause down. One likely contributor is that every step hands the GPU tensors of a new shape, since the cache grows by one position each time, and work that depends on the shape can’t be reused from the previous step. Production inference code avoids that by allocating the cache at its full size once and writing into it.

The generated text, for the record, is chapter 6’s GPT-2 being itself: “The history of the bicycle begins with the invention of the bicycle. The bicycle was first introduced in the United States in 1848 by the inventor, John W. H. Hutton.” and then the same sentence again.

Sharing keys and values: grouped-query attention

The cache trades compute for memory, and the memory adds up. Every layer stores a key and a value vector for every head and every position. For one token:

bytes per token=2×layers×KV heads×head dim×bytes per number\text{bytes per token} = 2 \times \text{layers} \times \text{KV heads} \times \text{head dim} \times \text{bytes per number}

For GPT-2 124M in 32-bit floats, that is 2 × 12 × 12 × 64 × 4 = 73,728 bytes per token, 72 MiB for a full 1,024-token context. Small. Llama 3 8B has 32 layers and 128-dimensional heads, stores numbers in 16-bit bfloat16, and supports contexts up to 128K tokens. With one key and value per head for each of its 32 query heads, its cache would take 512 KiB per token, 4 GiB at 8,192 tokens and 64 GiB at 131,072, for a single sequence, on top of 16 GB of weights.

Shazeer (2019) observed that generating one token at a time is limited mostly by moving these large key and value tensors from memory, not by arithmetic, and proposed multi-query attention: keep all the query heads but give them a single shared key head and value head. The cache shrinks by the number of heads, at some cost in quality. Ainslie et al. (2023) proposed the middle ground, grouped-query attention: split the query heads into groups and give each group one key and value head. With 32 query heads and 8 groups, every 4 query heads share a key and a value.

Llama 3 8B cache, bf16      per token     8,192 tokens    131,072 tokens
32 KV heads (plain MHA)     512 KiB       4.0 GiB         64 GiB
 8 KV heads (GQA, actual)   128 KiB       1.0 GiB         16 GiB
 1 KV head  (MQA)            16 KiB       0.125 GiB        2 GiB

The Llama 3 paper gives 8 key/value heads for all three of its sizes (8B, 70B, 405B) and names the reason: inference speed and a smaller key-value cache during decoding. In Llama 2 only the 34B and 70B models used it.

The code change is small. Chapter 3’s fused qkv layer has to be split back into three, because the key and value projections now get narrower than the query one: they produce n_kv_heads × head_dim numbers instead of emb_dim. Before the dot product, each key and value head is repeated to line up with the query heads in its group:

-        self.qkv = nn.Linear(emb_dim, 3 * emb_dim, bias=qkv_bias)
+        self.n_kv_heads = n_kv_heads
+        self.W_q = nn.Linear(emb_dim, emb_dim, bias=qkv_bias)
+        self.W_k = nn.Linear(emb_dim, n_kv_heads * self.head_dim, bias=qkv_bias)
+        self.W_v = nn.Linear(emb_dim, n_kv_heads * self.head_dim, bias=qkv_bias)
 ...
-        q, k, v = self.qkv(x).split(C, dim=-1)
-        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        q = self.W_q(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.W_k(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.W_v(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        group = self.n_heads // self.n_kv_heads
+        k = k.repeat_interleave(group, dim=1)   # kv head 0 serves query heads 0..group-1
+        v = v.repeat_interleave(group, dim=1)

repeat_interleave(2, dim=1) turns heads [a, b] into [a, a, b, b]. The cache stores k and v before the repeat, which is where the saving comes from. Production kernels avoid materializing the repeated copy at all; this version keeps the arithmetic visible.

Many feed-forward networks, few used: mixture of experts

The last change is the largest in spirit. In a dense model every token passes through every weight. A mixture-of-experts layer replaces the one feed-forward network in a block with several, called experts, and adds a small router: a linear layer that scores each expert for the current token. Only the top few experts run, and their outputs are added together, weighted by the router’s softmaxed scores.

With 4 experts and top-2 routing, a token whose router scores are [2.0, -1.0, 0.5, 1.5] goes to experts 0 and 3. Softmax over those two scores, e^2.0 and e^1.5, gives weights 0.62 and 0.38, so the layer’s output is 0.62 × expert0(x) + 0.38 × expert3(x). Experts 1 and 2 aren’t computed at all for this token. The next token may go elsewhere.

class MoE(nn.Module):
    def __init__(self, d, hidden, n_experts=8, top_k=2):
        super().__init__()
        self.router = nn.Linear(d, n_experts, bias=False)
        self.experts = nn.ModuleList(SwiGLU(d, hidden) for _ in range(n_experts))
        self.top_k = top_k

    def forward(self, x):                       # x: (tokens, d)
        scores, chosen = self.router(x).topk(self.top_k, dim=-1)
        weights = torch.softmax(scores, dim=-1)
        out = torch.zeros_like(x)
        for e, expert in enumerate(self.experts):
            rows, slot = (chosen == e).nonzero(as_tuple=True)   # tokens routed to expert e
            if len(rows):
                out[rows] += weights[rows, slot, None] * expert(x[rows])
        return out

topk returns the k largest scores and their indices. nonzero(as_tuple=True) lists which tokens picked expert e and in which of their two slots, so each expert runs once on just its tokens.

The point is to separate how much a model knows from how much it costs per token. Mixtral 8x7B has 8 experts per layer and routes each token to 2: 47B parameters in total, 13B used for any one token. Qwen3-235B-A22B has 128 experts per layer and activates 8, for 235B total and 22B active. The total still has to sit in memory, so MoE saves compute, not RAM. Training is also harder: a router left alone tends to send most tokens to a few experts, so MoE models add a load-balancing loss that rewards spreading tokens out. The idea is older than GPT-2: Shazeer et al. (2017) introduced the sparsely-gated layer, and the Switch Transformer simplified it to routing each token to a single expert.

At this track’s scale an MoE has nothing to show; on 40,738 tokens of training text there is too little to specialize on. It’s included here because several of the largest open models, Mixtral and Qwen3’s flagship among them, are built this way.

Putting the changes together

Each change is a switch on the same small model, so they can be tried on the same data. The test bed is the setup from chapter 5: The Time Machine, the first 90% of its tokens for training (40,738) and the last 10% for validation (4,527), and a 4-block model with 256-wide vectors and 128-token windows, about 28.9 million parameters, most of them in the two vocabulary-sized ends. Each variant trains for 600 steps of 16 random windows with AdamW at a learning rate of 0.001 and no dropout, and is scored by its best validation loss, checked every 50 steps. Every variant runs with seeds 1, 2 and 3, and all six variants share the same batches for a given seed. The heads are 8 per block, so GQA uses 2 key/value heads, 4 queries per group. “Modern” turns on all four changes and drops the biases, as Llama does.

VariantParametersBest validation loss, seeds 1 / 2 / 3Mean
GPT-2 style (chapter 4)28,923,9045.871 / 5.793 / 5.8675.843
+ RMSNorm only28,921,6005.857 / 5.747 / 5.8305.811
+ RoPE only28,891,1365.647 / 5.717 / 5.6705.678
+ SwiGLU only28,916,7365.764 / 5.715 / 5.7305.736
+ GQA only (2 KV heads)28,529,1525.859 / 5.924 / 5.8305.871
All four, no biases28,484,3525.657 / 5.637 / 5.6305.641

Each run took 83 to 93 seconds on the GPU. The seed-to-seed spread within a row is about 0.05 to 0.1, so read the table coarsely. RMSNorm and GQA land within noise of the baseline, which is what they’re for: RMSNorm is meant to be cheaper with no loss of quality, and GQA to shrink the cache with little loss of quality, and at this size neither saving is measurable. SwiGLU is a little better. RoPE is clearly better: 0.17 lower on average, and lower than the baseline in every seed. Everything together is the best row, and essentially the RoPE row.

The curves show where RoPE’s gain comes from.

5.56.06.57.07.50100200300400500600training step (batch 16 × 128 tokens)GPT-2 styleGPT-2 style + RoPE onlyall four changes
Validation loss every 50 steps, three seeds per model (one line each), on The Time Machine. Values above 7.7 are cut off at the top. M3 Pro GPU (MPS).

The GPT-2-style model sits near 6.7 for its first 50 to 150 steps, depending on the seed, before it starts to improve. With RoPE, validation loss is already near 6.0 by step 100 in every seed. A plausible reading, consistent with the “RoPE only” row matching the full modern model, is that the GPT-2-style model first has to learn from scratch what its 128 position vectors mean, while a rotary model gets relative position built into every attention score from step one. On a book this small, a head start is most of the contest: every variant reaches its best within 150 to 400 steps and then overfits, the modern ones faster, since they learn faster. This says that the changes help a small model learn sooner. It says nothing on its own about which architecture is better after trillions of tokens; that evidence is in the papers, at scales no laptop reaches.

What a current model looks like

Llama and Qwen models are the chapter 4 design with these changes made. Pre-normalization, residual connections, causal multi-head attention and a feed-forward network per block are all still there. The figures below come from the Llama 3 paper (Table 3), the Qwen3 technical report (Table 1), and, for numbers the report doesn’t give, the config.json published with Qwen3-0.6B.

GPT-2 124MQwen3-0.6BLlama 3 8B
Layers122832
Vector width7681,0244,096
Query heads / KV heads12 / 1216 / 832 / 8
Head size64128128
Feed-forwardGELU, 3,072SwiGLU, 3,072SwiGLU, 14,336
NormalizationLayerNormRMSNormRMSNorm
Positionslearned table, 1,024RoPE, base 1,000,000RoPE, base 500,000
Vocabulary50,257151,669128K
Tied input/output embeddingyesyesno

A few details beyond the four changes. Qwen3’s heads are wider than the model divided by the number of heads: 16 heads of 128 make 2,048, twice the 1,024 width, and the output projection maps it back. The head count and the width no longer have to be tied together. Qwen3 also removed the biases on q, k and v that Qwen2 had, and added QK-Norm, an RMSNorm applied to each query and key before the dot product, for training stability. Meta’s reference Llama 3 code has no biases in any linear layer and keeps a separate output matrix rather than reusing the embedding. Vocabularies grew to between two and a half and three times the size; Llama 3 took 100K tokens from tiktoken and added 28K to better support languages other than English. And Qwen3’s largest models are mixture-of-experts.

The biggest difference isn’t in the table. GPT-2 124M was trained on about 40 GB of web text; Llama 3 on about 15 trillion tokens, and its paper says plainly that its gains over Llama 2 came “primarily” from better data and more training, not architecture. Then come the stages from chapter 8: instruction tuning on far more and better examples, and preference tuning. The skeleton built in chapters 2 to 4 is still the skeleton.

Where this leaves things

This chapter closes the track. It started with a sentence that had to become numbers, and each chapter added the piece the last one lacked. Text became token IDs through byte-pair encoding. IDs became vectors through an embedding table, with a second table for position. Attention let each vector gather from the ones before it, under a mask that hides the future. The transformer block wrapped attention and a feed-forward network in layer norms and residual connections, and twelve of them stacked gave GPT-2’s shape. Training on next-token prediction made a small version of it write recognizable English from one book. OpenAI’s weights filled the real-sized version, and it wrote fluent web text. A two-way head and a minute of training turned it into a spam filter, and a thousand examples in a fixed format turned it into something that answers when asked.

What changed since GPT-2 is smaller than the gap in capability suggests. Normalization lost its mean. Position moved from a table into a rotation inside attention. The feed-forward network gained a gate. Keys and values are shared across heads and cached between steps, and in the largest models only some of the feed-forward networks run for each token. Every one of these is a few lines of difference from the chapter 4 code. The rest of the distance between GPT-2 and a current model is scale, data, and the tuning that follows, and the code for those is the same loop from chapter 5, run for much longer on much more.