Step 8 — Rebuild with the modern design

Step 8 — Rebuild with the modern design 🏗️¶

Goal: replace your Step 1 model with the architecture people actually use today — one change at a time, measuring each one, until you can open any modern model’s config.json and read it like a sentence.

Why this step matters¶

The transformer from the 2017 paper Attention Is All You Need is history, not a target. Almost nothing in it survived unchanged.

You built it in Step 1 so you would understand it. Now build the current one, so that when you open a real model in Step 11, every line means something to you.

The shape of this chapter 🗺️¶

Five changes. Each one replaces exactly one thing you wrote in Step 1:

Sanskrit gives you a specific reason to care about number 2: word order in Sanskrit is unusually free, because grammatical role is carried by the ending rather than the position. A position scheme that encodes the gap between two words rather than their absolute slot numbers is a better match for that.

1 · RMSNorm — LayerNorm with the mean taken out 📏¶

LayerNorm centres the vector (subtract the mean), then scales it (divide by the standard deviation). RMSNorm skips the centring and just divides by the root-mean-square.

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))    # gain only, no bias

    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return x * rms * self.weight

Compare that to what LayerNorm does and the saving is visible: one less pass over the tensor, one less running statistic, one less parameter vector. It turns out the centring was not doing much work.

Also confirm your normalization runs before the block, not after. You proved why in Step 2, experiment 2.

2 · RoPE — rotate the queries and keys 🔄¶

Instead of adding a position vector to each token, RoPE rotates the query and key vectors by an angle based on their position.

def build_rope_cache(head_dim, max_len, base=10000.0):
    # Each PAIR of dimensions gets its own rotation frequency. Low dimensions
    # rotate fast (fine, local position); high dimensions rotate slowly
    # (coarse, long-range position).
    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
    freqs = torch.outer(torch.arange(max_len).float(), inv_freq)
    return freqs.cos(), freqs.sin()


def apply_rope(x, cos, sin):
    B, H, T, D = x.shape
    x1, x2 = x[..., 0::2], x[..., 1::2]               # even and odd dimensions
    c = cos[:T].view(1, 1, T, D // 2)
    s = sin[:T].view(1, 1, T, D // 2)
    out = torch.stack([x1 * c - x2 * s,               # a plain 2D rotation
                       x1 * s + x2 * c], dim=-1)
    return out.flatten(-2)

The useful consequence: when a query at position 10 meets a key at position 3, the two rotations partially cancel, and what survives depends on 10 - 3 = 7.

Two things fall out of this that matter later:

• RoPE has no parameters at all. The learned position table disappears from your model. On a small model that is a real, if modest, saving.

• You can stretch the angles. That is what makes context extension possible in Step 19 — you cannot stretch a lookup table.

3 · SwiGLU — add a gate to the feed-forward 🚪¶

SwiGLU splits the feed-forward into two paths: one carries content, the other is a gate deciding how much of the content gets through.

class SwiGLU(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        hidden = int(4 * cfg.n_embd * 2 / 3)           # ← the important line
        hidden = 64 * ((hidden + 63) // 64)            # round up; kernels prefer it
        self.gate = nn.Linear(cfg.n_embd, hidden, bias=cfg.bias)
        self.up   = nn.Linear(cfg.n_embd, hidden, bias=cfg.bias)
        self.down = nn.Linear(hidden, cfg.n_embd, bias=cfg.bias)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

4 · Grouped-query attention — share the keys and values 🔗¶

In Step 1, every query head had its own keys and values. In grouped-query attention, several query heads share one set.

The change is visible in three lines of the constructor — note the asymmetry:

self.q_proj = nn.Linear(cfg.n_embd, cfg.n_head    * self.head_dim, bias=cfg.bias)
self.k_proj = nn.Linear(cfg.n_embd, cfg.n_kv_head * self.head_dim, bias=cfg.bias)  # narrower
self.v_proj = nn.Linear(cfg.n_embd, cfg.n_kv_head * self.head_dim, bias=cfg.bias)  # narrower

and one line in the forward pass:

if self.n_kv_head != self.n_head:
    rep = self.n_head // self.n_kv_head
    k = k.repeat_interleave(rep, dim=1)               # one K/V serves a group of Qs
    v = v.repeat_interleave(rep, dim=1)

Quality barely changes. What drops is the size of the KV cache you carry during generation — and that cache is the thing that actually limits how many users you can serve at once.

While we are here, Step 1’s five lines of attention collapse into one call:

y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=...)

That is the same scale, mask, softmax, dropout, and weighted sum you wrote by hand — and on a GPU it dispatches to FlashAttention, which computes exactly the same thing without ever writing the T × T matrix to memory. Same maths, far less memory. More on that in Step 9.

5 · Remove every bias term ✂️¶

From the attention projections, from the feed-forward layers, from the normalization. bias=False everywhere. They were not earning their place, and removing them is slightly faster and slightly more stable.

This is the cheapest change in the chapter and the one most likely to show up as noise in your ablation. Include it anyway — knowing which changes don’t matter is part of the point.

6 · QK-norm, if and only if training is unstable 🌡️¶

Normalizing the queries and keys before the attention dot product stops those values from getting large, which lets you use a higher learning rate without the loss spiking.

if self.qk_norm:
    q, k = self.q_norm(q), self.k_norm(k)             # keeps the dot products small

Add it only if you need it. You will know you need it when you meet a real loss spike in Step 9 — the one you already caused on purpose in Step 2, experiment 4.

What you should see ▶️¶

python modern_gpt.py --ablate --smoke

device=cpu vocab=79 tokens=2,220

variant                 params    train      val    Δ val   secs
----------------------------------------------------------------
Step 1 baseline          0.11M   2.4311   2.4353              2s
+ RMSNorm                0.11M   2.4402   2.4445  +0.0091     2s
+ RoPE                   0.10M   2.4135   2.4218  -0.0227     2s
+ SwiGLU                 0.11M   2.4083   2.3983  -0.0235     2s
+ no bias                0.11M   2.4265   2.4317  +0.0334     2s
+ GQA (kv=2)             0.10M   2.4043   2.3970  -0.0348     2s
+ QK-norm                0.10M   2.2221   2.2183  -0.1787     2s

Now read it properly, because the honest reading is the lesson:

Run it longer and it gets worse, not better:

python modern_gpt.py --ablate --smoke --n-embd 128 --iters 600

Step 1 baseline          0.41M   0.0635   0.0599             25s
+ RMSNorm                0.41M   0.0641   0.0602  +0.0003     26s
+ RoPE                   0.41M   0.0714   0.0632  +0.0030     27s
+ SwiGLU                 0.44M   0.0666   0.0673  +0.0042     29s

A validation loss of 0.06 means the model has memorised the toy corpus outright — Step 2, experiment 5, arriving uninvited. Once every variant has memorised, they all score the same and every delta is noise.

The deliverable of this chapter is the harness, not this table. Point it at your real corpus from Step 7, run each variant on three seeds, and only believe a change whose effect is larger than the spread across seeds.

Where people usually get stuck¶

Copying a modern architecture wholesale and never learning which piece did what. You end up with a working model and no understanding, which is exactly the situation this book exists to avoid.

Comparing variants with different parameter counts. Cheap to avoid: print the count on every row and adjust the hidden size when the shape changes, as SwiGLU does.

Running one seed and believing it. Three seeds, and report the spread. If your effect is smaller than the spread, you do not have an effect.

Expecting these changes to help at toy scale. Most of them were designed for models a hundred times bigger than the one you can train this weekend. GQA in particular is a serving optimisation. You are adopting them so your model matches what the field does — the measurable gains come later.

Forgetting n_head % n_kv_head == 0. 8 heads with 2 KV heads works; 8 with 3 does not. The assert in the file will tell you, which is why it is there.

You are ready to move on when¶

You have a small table showing the loss before and after each individual change with parameter counts next to it, measured on your own Sanskrit data across more than one seed — and you can open any modern model’s config.json and explain what every field means.

A good test: someone shows you "num_key_value_heads": 8 next to "num_attention_heads": 32 and asks what it buys. You can answer in one sentence, and say when it matters.