Step 5 — Assembling the GPT

Released with Lecture 5 · Starts from: Step 4 (or solutions/step-04.ipynb) · Solution: notebook · read online

Goal

Assemble Step 4’s head into the complete GPT architecture: multi-head attention, MLP blocks, residual connections, LayerNorm, and the embedding/unembedding maps. At the end you have the full model class — untrained, but verified for shapes, causality, and an exact parameter count you derive on paper first. Step 6 trains it.

On paper first

1. The architecture as one formula. With $L$ blocks, write the model as the composition (pre-norm convention, as in GPT-2):

   $$X^{(0)} = \text{TokEmb}(x) + \text{PosEmb},\qquad \begin{aligned} X^{(\ell)} &{}\mathrel{+}= \text{MHA}(\text{LN}_1^{(\ell)}(X^{(\ell)}))\\ X^{(\ell)} &{}\mathrel{+}= \text{MLP}(\text{LN}_2^{(\ell)}(X^{(\ell)})) \end{aligned} \qquad Z = \text{LN}_f(X^{(L)})\,W_U.$$

   The += structure is the residual stream: every sublayer reads from and writes small updates into a running state in $\mathbb{R}^{T\times d}$. (Lecture 5 develops why this matters for both optimization and interpretation.)

2. Parameter count. Derive, term by term, the parameter count of your model as a function of $(V, T_{\max}, d, L)$ with $d_{ff} = 4d$ and biases where your code has them. Show the blocks contribute $\approx 12 L d^2$ and identify exactly which matrices the $12$ counts. You will check this number against numel() to the last parameter.

3. LayerNorm. For $u \in \mathbb{R}^d$: $\text{LN}(u) = \gamma \odot \frac{u - \bar u\,\mathbf 1}{\sigma(u)} + \beta$. Show the normalization part is: orthogonal projection onto $\mathbf 1^\perp$, then radial projection onto the sphere of radius $\sqrt d$. So the stream is repeatedly pinned to a $(d{-}2)$-sphere — a strong geometric constraint to keep in mind whenever you picture “vectors in the residual stream.”

Tasks

5.1 Multi-head attention

$H$ heads with $d_{\text{head}} = d/H$, run in parallel, concatenated, then mixed by an output projection:

class MultiHeadAttention(nn.Module):
    def __init__(self, d, H, T_max):
        super().__init__()
        self.heads = nn.ModuleList(Head(d, d // H, T_max) for _ in range(H))  # generator expression; // = integer division
        self.proj = nn.Linear(d, d)

    def forward(self, x):
        return self.proj(torch.cat([h(x) for h in self.heads], dim=-1))  # list comprehension

print(7 / 2)          # ordinary division: always a float
print(7 // 2)         # integer division: rounds down
print(128 / 4)        # exact — but still a float!
print(128 // 4)       # an int, usable as a size or index
for _ in range(3):    # run 3 times, counter unused
    print("another head")

(This loop-over-heads version is deliberately transparent; the “going further” refactor batches it into one tensor op.)

5.2 The MLP block

Position-wise: hidden width $4d$, GELU nonlinearity, back to $d$. Two nn.Linears and an F.gelu; note it acts on each of the $T$ positions independently — all cross-position communication in the entire model happens inside attention, nowhere else.

5.3 The block and the model

class Block(nn.Module):
    def __init__(self, d, H, T_max):
        super().__init__()
        self.ln1, self.attn = nn.LayerNorm(d), MultiHeadAttention(d, H, T_max)  # parallel assignment (tuple unpacking)
        self.ln2, self.mlp  = nn.LayerNorm(d), MLP(d)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))     # read, transform, write back
        x = x + self.mlp(self.ln2(x))
        return x

Then GPT(nn.Module) holding: token embedding (V, d), positional embedding (T_max, d), L blocks, final LayerNorm, and the unembedding nn.Linear(d, V, bias=False). forward(idx, targets=None) returns logits (B, T, V) and, if targets are given, F.cross_entropy on the flattened logits. Also port your Step 3 generate (crop the context to the last T_max tokens each iteration — why is cropping required for this model when it wasn’t for the MLP?).

Course configuration (matches Step 6’s training budget):

V=65, T_max=64, d=128, H=4, L=4        # just over 0.8M parameters — your baby GPT

5.4 Verification suite

1. Parameter count: sum(p.numel() for p in model.parameters()) must equal your paper formula exactly. Chase any discrepancy to the guilty term (biases and LayerNorms are the usual suspects). Report the percentage of parameters in: embeddings, attention, MLPs.
2. Causality, end-to-end: Step 4’s perturbation test on the full model. It must still pass — residuals, LayerNorm, and MLPs are all position-wise, so the mask remains the only cross-position gate.
3. Initial loss: on a random batch, loss should land just above $\ln 65 = 4.174$ (see checkpoints for the exact figure and why it is not equal). If it’s 8, your init is too hot (logits too spread — softmax collapsed); GPT-2’s fix, which you should apply: scale the output projections in each block by $1/\sqrt{2L}$ (Lecture 6 derives why: $2L$ writes into the stream, variance adds).
4. A shape table: print x.shape after every stage for one forward pass — embeddings, each block, final LN, logits. This table is the architecture; keep it next to the formula from “on paper” 1.

Checkpoints

• Parameter formula matches numel() exactly. For the course config the total is 816,640, distributed as: embeddings 16,512 (2.0%), attention 262,656 (32.2%), MLPs 526,848 (64.5%), block LayerNorms 2,048, final LayerNorm + unembedding 8,576. Note the MLPs hold nearly twice the parameters of the attention layers — a fact worth remembering whenever someone calls the transformer “an attention architecture.”
• The $12Ld^2$ rule of thumb gives 786,432 against an actual block total of 791,552; the 5,120 difference is exactly the biases and LayerNorm parameters the rule ignores. Confirm that decomposition yourself.
• Causality and simplex tests pass on the full model.
• Initial loss ≈ 4.2–4.4 against $\ln 65 = 4.174$ (measured 4.30 on a real batch, 4.34 on random targets). It sits slightly above $\ln V$ because PyTorch’s default nn.Linear initialization gives the unembedding a std of $1/\sqrt d$, spreading the logits a little; GPT-2 uses std 0.02 throughout, which pulls it down to $\ln V$ almost exactly. Try both and explain the direction of the gap. If you see 8+, something is badly mis-scaled; after 50 quick gradient steps on one repeated batch it plummets (memorizing one batch is easy — a useful “the plumbing works” smoke test, and your first glimpse of how overparametrized this model is).
• Samples from the untrained model: uniform gibberish drawing all 65 characters. Save one — it’s the “before” picture.

If you’re stuck

• Loss nan at init: check the $-\infty$ mask meets softmax before any LayerNorm gets a zero-variance row (all-equal inputs at init).
• Parameter count off by exactly $2dL$ or $2d$: forgotten LayerNorm gammas/betas.
• Generation crashes past 64 tokens: you didn’t crop the context to T_max (positional embedding matrix has no row 65).

Going further (optional)

• Weight tying: set the unembedding equal to the token embedding transposed (lm_head.weight = tok_emb.weight), as GPT-2 does. Recount parameters; think about why input and output vocabulary geometry might reasonably coincide.
• Batched heads: refactor MultiHeadAttention to a single tensor of shape (B, H, T, d_head) with one matmul for all heads; assert numerical agreement with the loop version. This is what every production implementation does.
• Read your model in the QK/OV language of Elhage et al. (2021): per head, which products of your weight matrices form the “QK circuit” ($W_Q W_K^\top$) and “OV circuit”? Verify each has rank ≤ $d/H = 32$.

Catch-up

Run solutions/step-04.ipynb, whose Head class is this step’s only dependency; everything else is fresh construction from Lecture 5.