Step 5 — Assembling the GPT

Released with Lecture 5 · ~3 hours · Starts from: Step 4 (or solutions/step-04.ipynb).

In what follows, there are blocks of code provided for you. It’s important that you can read the code. If you are not proficient with Python, then I propose the following workflow. For each block of code, keywords are provided for the python syntax which is being used, so you can look up python syntax and functions you don’t recognise. To help, beneath many code blocks you will find Python background foldouts: self-contained tutorials, with live Python boxes you can edit and run right on the page, for syntax that is new in this step (syntax already introduced is covered in earlier steps’ foldouts). Then I suggest that you first experiment with changing the code or printing out intermediate variables, until you have figured out what you think it does. Then point AI toward the page and step number, provide your precise explanation of what you believe each line of the code is doing, and ask for corrections. In this way you will quickly become a code reader (if not a code writer). — Kate

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.