Part I — Basics

The modern recipe — what changed since GPT-2

Lesson 02 built the decoder-only Transformer and proved it can be trained by one parallel cross-entropy pass. But the block we built is the 2019 GPT-2 block, and at modern depth and scale it misbehaves: the residual stream's scale grows with depth and big runs spike during training, LayerNorm and GELU spend parameters that buy nothing, learned absolute positions cannot see past T_max, and a full set of K/V heads bloats the cache you carry at inference. This lesson swaps five components — each forced by one of those weaknesses — to get the Llama-style block that every frontier dense model ships in 2024–2025. The thesis throughout: each swap buys efficiency or stability at ≈ zero quality cost, nearly all at ≈ zero extra parameters and little added compute — though a few (SwiGLU's gate, QK-norm) do add a small cost.

1 · RMSNorm — keep pre-norm, drop the dead params

Lesson 02 already adopted pre-norm placement — normalizing the input to each sublayer, inside the residual branch — so the only norm change left versus the block we inherited is which normalizer sits there: LayerNorm → RMSNorm. Normalization exists to keep the residual stream's scale under control so gradients stay sane through L blocks. The inherited block uses LayerNorm: subtract the mean, divide by the standard deviation, then apply a learned scale γ and bias β:

RMSNorm asks: do we need the mean-centering and the bias at all? It drops both, keeping only the root-mean-square rescale and the learned gain:

The win is twofold. Cheaper: no mean to compute (one fewer reduction over d), no subtraction, no β to store or update — you save d parameters per norm and a memory pass. Since norms are memory-bound elementwise ops (lesson 6's roofline), fewer reductions and fewer bytes moved is a direct latency win — see the kernel-cost view in gpu_kernels/02 · memory hierarchy. Equal quality: mean-centering buys essentially nothing in a residual Transformer because the residual stream is already roughly centered; ablations show RMSNorm matches LayerNorm. The verdict is settled: RMSNorm won — nobody ships LayerNorm in a new LM.

This swap keeps the pre-norm placement from lesson 2 and only changes the normalizer; it is worth recalling why pre-norm is the placement we hold fixed. Pre-norm normalizes the input to each sublayer, inside the residual branch: x = x + sublayer(RMSNorm(x)). Post-norm (the original 2017 design) normalizes after the add: x = RMSNorm(x + sublayer(x)). The difference is everything at depth. In pre-norm there is a clean identity path from input to output — the gradient flows straight down the residual highway without passing through a norm, so it neither vanishes nor explodes across L layers; the stream's variance grows at most linearly in L and a final norm cleans it up. In post-norm the backward gradient threads through L norms in series, small Jacobian errors compound, and you must add an LR warmup just to survive the first hundreds of steps. Pre-norm removes that fragility for free (same params, same FLOPs), which is why every modern stack is pre-norm.

  post-norm (2017)                pre-norm (modern)
  x ─┬─► sublayer ─► (+) ─► Norm   x ─┬─► Norm ─► sublayer ─► (+) ─►
     └────────────────►            └───────────────────────►
  gradient crosses a Norm          clean identity path: gradient
  every layer → compounds          skips every Norm → stable at depth

2 · RoPE — rotate Q and K so attention sees relative position

The learned absolute position table has two flaws. It costs T_max · d parameters, and it is a hard wall: position 8,193 in a model trained to T_max = 8,192 has no embedding row, so the model has never seen it and cannot extrapolate. We want a positional scheme that is parameter-free and degrades gracefully past the training length. The mechanism that wins is Rotary Position Embedding (RoPE).

RoPE injects position not by adding a vector but by rotating the query and key vectors by an angle proportional to their absolute position. Split each head's d_head dimensions into pairs; pair i rotates at frequency θ_i = θ_base^−2i/d_head with base θ_base (commonly 10,000). The query at position m is rotated by angle m·θ_i; the key at position n by n·θ_i.

Here is the payoff, and it falls straight out of the trigonometry. Take one 2-D coordinate pair of the query and key, and let R(·) be the 2×2 rotation. RoPE rotates the query at position m by angle mθ_i and the key at position n by nθ_i: q_m = R(mθ_i)q and k_n = R(nθ_i)k. The score this pair contributes is then

using the fact that rotations compose as R(a)ᵀR(b) = R(b − a) (a rotation's transpose is its inverse, R(a)ᵀ = R(−a), and rotations add). So the contribution depends only on the relative offset (n − m) — never on m or n separately. The full head applies d_head/2 such independent 2-D rotations, one per pair, at frequencies θ_i, so the whole attention score q_m · k_n ends up a function of (n − m) — the relative distance — not the absolute positions m and n.

That single property is the whole reason RoPE dominates. The model learns "the token 5 back," "the token 1 back," and those relations are translation-invariant: a pattern learned at positions 100–105 transfers to 6,000–6,005 unchanged. Zero parameters — it is a fixed rotation applied inside attention, nothing to store. And it extrapolates: a position past T_max is just a larger rotation angle, which is well-defined, so the model degrades gently rather than hitting a missing-row wall. RoPE is applied to Q and K only — never to V, because V carries content, not position.

Raw RoPE only extrapolates gently, though. Pushing the context to 10–100× the trained window — via base-θ_base scaling, position interpolation, NTK-aware scaling, and YaRN, usually plus a short continued-training phase — is its own engineering step, not something you get for free here. The full treatment lives in lesson 15 · Long context; this lesson just plants the property that makes it possible.

3 · SwiGLU — a gated MLP that beats GELU per parameter

The GPT-2 MLP is d → 4d → d with GELU: proj(GELU(fc(x))), two matrices of total size 8d². SwiGLU replaces it with a gated linear unit: one branch is a Swish-activated projection, a second branch is a plain linear gate, and they are multiplied elementwise before the down-projection:

where Swish(z) = z · σ(z) (a.k.a. SiLU), ⊙ is elementwise product, and W₁, W₃ map d → d_ff while W₂ maps d_ff → d. The gate x·W₃ lets each hidden unit modulate the others multiplicatively — a richer, data-dependent nonlinearity than a single pointwise GELU — and ablations (Shazeer 2020) show it consistently improves loss per parameter.

But SwiGLU uses three weight matrices instead of two, so a naive d_ff = 4d would inflate the MLP from 8d² to 12d² — 50% more params. To keep the block's parameter budget fixed (so the comparison with GELU is fair), shrink the hidden width to d_ff ≈ (8/3)·d ≈ 2.667d. Then the three matrices total 3 · d · (8/3)d = 8d² — exactly the old MLP's count. Same parameters, same FLOPs, lower loss: a free upgrade. (In practice d_ff is rounded to a hardware-friendly multiple, so you'll see values like 11,008 rather than a clean 8/3 fraction.)

4 · GQA — share K/V across query-head groups, slash the KV cache

This swap costs almost nothing at training time and saves enormously at inference, so its motivation lives in lesson 18 — but the mechanism belongs in the architecture. At generation time the model caches the K and V vectors of every past token so it doesn't recompute them each step. In standard multi-head attention (MHA) each of the h query heads has its own K and V head, so the cache stores 2 · L · h · d_head values per token. For a large model over a long context, that KV cache — not compute — is what caps batch size and sequence length (we forward-reference the full derivation in lesson 18 · inference).

Grouped-query attention (GQA) keeps all h query heads but shares each K/V head across a group of g query heads. You go from h K/V heads down to n_kv = h / g K/V heads. The cache shrinks by exactly the group factor g. Two endpoints bracket the spectrum:

Concretely: KV-cache bytes per token = 2 · L · n_kv · d_head · bytes. Take a 7B-class model: L = 32, h = 32, d_head = 128, bf16 (2 bytes). MHA stores 2 · 32 · 32 · 128 · 2 = 524,288 bytes = 512 KB per token. With GQA at g = 4 (n_kv = 8) it's 128 KB per token — a 4× cut. Over a 32K context that is the difference between ~17.2 GB and ~4.3 GB of cache. The query side is untouched, so the model still attends with all h heads; only the stored keys/values are shared. Quality cost: small and well-characterized — which is why Llama-2 70B, Llama-3, Mistral, and essentially every recent serving-oriented model ship GQA.

5 · Bias-free linears + QK-norm / z-loss — stability for free

Two last cleanups, both about large-scale stability with negligible cost.

Drop the biases. GPT-2 puts a bias on every linear and on LayerNorm. In a Transformer these biases contribute almost nothing to quality — the residual stream and the (now bias-free) RMSNorm gain already supply the needed offsets — yet they add parameters, a little memory traffic, and at very large scale they are a documented source of training-loss spikes. So modern blocks make every projection bias=False. The saving is small in parameters but the stability and simplicity are worth it; combined with RMSNorm's dropped β, the whole block is essentially bias-free.

QK-norm. At large width and high learning rate the attention logits q·k can grow until the softmax saturates and the loss spikes. QK-norm applies an RMSNorm (or L2 normalization) to Q and K just before the dot product, capping the logit magnitude so attention can't blow up. It costs two tiny norms per attention and removes a whole class of divergence at the 10B-plus scale.

z-loss. The output softmax has a similar failure: the log-partition (the log-sum-exp normalizer Z) can drift to large values, which interacts badly with bf16 and destabilizes training. z-loss adds a small penalty λ · (log Z)² to the objective that keeps the normalizer near 1 — a cheap regularizer that PaLM and others credit with smoother large runs. Neither QK-norm nor z-loss changes what the model can represent; they are guardrails that keep the optimizer in the stable regime, paying off at the scale where a single spike can waste days of budget.

6 · The whole swap, GPT-2 vs Llama-style

None of these is a research gamble anymore; together they are the boring, settled default. The skill is not picking them — it is knowing which weakness each one repairs, so when a run spikes or a context overflows you reach for the right knob. The widget makes that mapping tangible.

7 · Toggle the recipe

Each checkbox turns one swap on or off, starting from a 7B-class shape (≈6.2B with the modern swaps on, ~7.0B vanilla) (L = 32, d = 4096, h = 32, d_head = 128, V = 128,000, T_train = 8,192). Watch the KPIs move: param count (GQA shrinks the K/V projections — the biggest mover here; RoPE removes the position table; SwiGLU is param-matched; bias-free trims a sliver), KV-cache bytes / token (GQA's group factor g divides it directly), stability at depth (pre-norm/RMSNorm dominates, with bias-free + QK-norm/z-loss adding a smaller push toward green), and max usable context (learned absolute caps you at T_train; RoPE unlocks far beyond it). The configuration that "breaks": turn everything off — you get the vanilla GPT-2 block, with the biggest cache, a hard context wall at T_train, and a stability marker in the red.

Failure modes & checklist

Checkpoint

Where this points next

You now have a defined architecture — the Llama-style block, every component justified by the vanilla weakness it repairs. But a defined model is inert: nothing yet says how to move its parameters down the loss without the run diverging or burning the budget on a schedule that spikes. RMSNorm, pre-norm, QK-norm, and z-loss only buy stability headroom; they don't choose the optimizer, the learning-rate schedule, the warmup, the gradient clipping, or the precision that actually keep a billion-parameter run on the rails. That is the training recipe — AdamW, warmup→cosine, clipping, and bf16 mixed precision — in lesson 4 · Training.

Interview prompts

• Write RMSNorm and say what it drops relative to LayerNorm and why that's safe. (§1 — RMSNorm(x)=x/√(mean(x²)+ε)·γ; it drops mean-centering and the bias β. The residual stream is already roughly centered so mean-subtraction buys ~nothing, and you save a reduction, a subtraction, and d params per norm — a memory-bound op gets cheaper.)
• Why is pre-norm stable at depth where post-norm needs warmup? (§1 — pre-norm leaves a clean identity residual path so the gradient skips every norm; variance grows ≈ linearly in L. Post-norm threads the backward gradient through L norms in series, compounding Jacobian error, so it needs warmup to survive early steps.)
• How does RoPE make attention see relative position, and why does that buy extrapolation? (§2 — rotating q and k by m·θ_i and n·θ_i makes their inner product a function of (n−m); patterns are translation-invariant. A position past T_train is just a larger, well-defined rotation angle, so quality degrades gently instead of hitting a missing-row wall — scale θ_base / NTK / YaRN for long context, lesson 15.)
• SwiGLU uses three matrices — how do you keep the parameter count equal to a GELU MLP? (§3 — set d_ff ≈ (8/3)d; then 3·d·(8/3)d = 8d², the same as GELU's two 4d matrices. Skip this and you get 12d², a 50% MLP bloat.)
• Derive the KV-cache size and explain how GQA shrinks it. (§4 — cache = 2·L·n_kv·d_head·bytes per token; MHA has n_kv=h, GQA has n_kv=h/g, so the cache divides by g. All h query heads stay; only the stored K/V heads are shared, so quality cost is tiny — e.g. 32 query / 8 KV heads → 4× smaller cache.)
• What do QK-norm and z-loss each prevent, and why are they "free"? (§5 — QK-norm normalizes Q,K before the dot product to cap attention logits so softmax can't saturate; z-loss penalizes (log Z)² to keep the output-softmax partition near 1, which behaves badly in bf16 otherwise. Neither changes representational capacity — they're guardrails against spikes at 10B-plus scale.)
• You serve a model trained at 8K out to 128K and quality collapses past ~16K. What's the likely cause and fix? (§2 + lesson 15 — raw RoPE extrapolation degrades far out; raise the base θ_base and/or apply position-interpolation / NTK / YaRN so the high-frequency pairs don't alias on unseen distances.)