Step 5 — Assembling the GPT¶

Solution notebook. Step 4's head, multiplied and wrapped into the complete architecture, with a parameter count verified to the integer against a formula derived on paper. No training — that is Step 6. Runs in seconds.

Setup: data and tokenizer (Step 1's solution, reproduced)¶

Every notebook in this series is self-contained: it re-creates what it needs from earlier steps in one compact cell, so you can run it top to bottom without opening the others. On Colab, uncomment the download.

In [ ]:
# !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
import torch

with open('input.txt') as f:
    text = f.read()
vocab = sorted(set(text))
V = len(vocab)
stoi = {ch: i for i, ch in enumerate(vocab)}
itos = {i: ch for i, ch in enumerate(vocab)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: ''.join(itos[i] for i in ids)
ids = torch.tensor(encode(text), dtype=torch.long)

assert len(text) == 1_115_394 and V == 65
print(f'{len(text):,} characters, vocab {V}')

5.1–5.3 The architecture¶

Pre-norm blocks (Lecture 5 §3): each sublayer reads a normalized copy of the stream and adds its output back, so the Jacobian is $I + \partial F$ and gradients flow. All cross-position communication happens inside Head; everything else is position-wise, which is what makes the causality proof survive assembly.

In [ ]:
import math
import torch.nn as nn
import torch.nn.functional as F

class Head(nn.Module):
    def __init__(self, d, dh, T):
        super().__init__()
        self.key   = nn.Linear(d, dh, bias=False)
        self.query = nn.Linear(d, dh, bias=False)
        self.value = nn.Linear(d, dh, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(T, T)))

    def forward(self, x):
        B, T, _ = x.shape
        q, k, v = self.query(x), self.key(x), self.value(x)
        s = q @ k.transpose(-2, -1) / k.shape[-1]**0.5
        s = s.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        A = F.softmax(s, dim=-1)
        self.A = A.detach()          # stashed for Step 8 interpretability
        return A @ v

class MHA(nn.Module):
    def __init__(self, d, H, T):
        super().__init__()
        self.heads = nn.ModuleList(Head(d, d // H, T) for _ in range(H))
        self.proj = nn.Linear(d, d)

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

class MLP(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.fc, self.proj = nn.Linear(d, 4 * d), nn.Linear(4 * d, d)

    def forward(self, x):
        return self.proj(F.gelu(self.fc(x)))

class Block(nn.Module):
    def __init__(self, d, H, T):
        super().__init__()
        self.ln1, self.attn = nn.LayerNorm(d), MHA(d, H, T)
        self.ln2, self.mlp  = nn.LayerNorm(d), MLP(d)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        return x + self.mlp(self.ln2(x))

class GPT(nn.Module):
    def __init__(self, V, T, d, H, L):
        super().__init__()
        self.T = T
        self.tok = nn.Embedding(V, d)
        self.pos = nn.Embedding(T, d)
        self.blocks = nn.ModuleList(Block(d, H, T) for _ in range(L))
        self.lnf = nn.LayerNorm(d)
        self.head = nn.Linear(d, V, bias=False)
        for blk in self.blocks:   # GPT-2 residual scaling (Lecture 6 §1)
            nn.init.normal_(blk.attn.proj.weight, std=0.02 / math.sqrt(2 * L))
            nn.init.normal_(blk.mlp.proj.weight,  std=0.02 / math.sqrt(2 * L))

    def forward(self, idx, targets=None):
        B, T = idx.shape
        x = self.tok(idx) + self.pos(torch.arange(T))
        for blk in self.blocks:
            x = blk(x)
        logits = self.head(self.lnf(x))
        loss = None if targets is None else F.cross_entropy(
            logits.view(-1, logits.size(-1)), targets.reshape(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, n_new, temperature=1.0):
        for _ in range(n_new):
            logits, _ = self(idx[:, -self.T:])
            probs = F.softmax(logits[:, -1, :] / temperature, dim=-1)
            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)
        return idx
In [ ]:
torch.manual_seed(1337)
model = GPT(V=65, T=64, d=128, H=4, L=4)
print(model.__class__.__name__, '| blocks:', len(model.blocks))

5.4 Verification 1 — the parameter count¶

Derive on paper first (Lecture 5, Thm 5.1), then check. Getting this exact means you have accounted for every matrix in the model.

In [ ]:
def count(mod):
    return sum(p.numel() for p in mod.parameters())

total = count(model)
emb   = model.tok.weight.numel() + model.pos.weight.numel()
attn  = sum(count(b.attn) for b in model.blocks)
mlps  = sum(count(b.mlp)  for b in model.blocks)
lns   = sum(count(b.ln1) + count(b.ln2) for b in model.blocks)
final = count(model.lnf) + model.head.weight.numel()

for name, v in [('embeddings (tok+pos)', emb), ('attention', attn),
                ('MLPs', mlps), ('block LayerNorms', lns),
                ('final LN + unembed', final)]:
    print(f'  {name:<22}{v:>9,}  ({100*v/total:5.1f}%)')
print(f'  {"TOTAL":<22}{total:>9,}')
assert total == 816_640

d, L = 128, 4
print(f'\n  12*L*d^2 rule of thumb: {12*L*d*d:,}')
print(f'  actual block total    : {attn+mlps+lns:,}')
print(f'  difference            : {attn+mlps+lns-12*L*d*d:,}'
      f'  (= biases + LayerNorms, which the rule ignores)')

The MLPs hold nearly twice the parameters of the attention layers (64.5% against 32.2%) — worth remembering whenever someone calls this 'an attention architecture'.

5.4 Verification 2 — causality survives assembly¶

Step 4's perturbation test, on the full model. It must still pass: residuals, LayerNorm, and MLPs are all position-wise, so the causal mask remains the only cross-position gate in the entire network.

In [ ]:
idx = ids[:32][None, :]
with torch.no_grad():
    o1, _ = model(idx)
    idx2 = idx.clone(); idx2[0, 20] = (idx2[0, 20] + 7) % 65
    o2, _ = model(idx2)
assert torch.equal(o1[:, :20, :], o2[:, :20, :])
assert not torch.allclose(o1[:, 20:, :], o2[:, 20:, :])
print('causality end-to-end          OK  (positions 0-19 bitwise identical)')

5.4 Verification 3 — initial loss¶

Should be just above $\ln 65 = 4.1744$: the model starts knowing only the vocabulary size. It sits slightly above because PyTorch's default nn.Linear init gives the unembedding a std of $1/\sqrt d$, spreading the logits a little; GPT-2's uniform std 0.02 would pull it to $\ln V$ almost exactly. A value near 8 means something is badly mis-scaled.

In [ ]:
import math
xb, yb = ids[:64][None, :], ids[1:65][None, :]
with torch.no_grad():
    _, loss0 = model(xb, yb)
print(f'initial loss on a real batch: {loss0.item():.4f}')
print(f'ln V                        : {math.log(65):.4f}')
assert 4.0 < loss0.item() < 4.5

Smoke test: can it memorize one batch?¶

50 steps on a single repeated batch. The loss should plummet — not because the model is good, but because it is enormously overparametrized relative to 4,096 tokens. This checks the plumbing (gradients reach every parameter) before we spend an hour on Step 6.

In [ ]:
import copy
probe = copy.deepcopy(model)
opt = torch.optim.AdamW(probe.parameters(), lr=1e-3)
xb, yb = ids[:64][None, :], ids[1:65][None, :]
for i in range(51):
    _, l = probe(xb, yb)
    opt.zero_grad(set_to_none=True); l.backward(); opt.step()
    if i % 25 == 0:
        print(f'  step {i:>2}  loss {l.item():.4f}')
assert l.item() < loss0.item() / 2

5.4 Verification 4 — the shape table¶

Print the stream's shape after every stage. This table is the architecture; keep it beside the five formulas of Lecture 5 §5.

In [ ]:
with torch.no_grad():
    idx = ids[:64][None, :]
    print(f'{"token indices":<28}{tuple(idx.shape)}')
    x = model.tok(idx) + model.pos(torch.arange(64))
    print(f'{"after embedding":<28}{tuple(x.shape)}')
    for i, blk in enumerate(model.blocks):
        x = blk(x)
        print(f'{f"after block {i}":<28}{tuple(x.shape)}')
    x = model.lnf(x)
    print(f'{"after final LayerNorm":<28}{tuple(x.shape)}')
    z = model.head(x)
    print(f'{"logits":<28}{tuple(z.shape)}   <- (B, T, V)')

The 'before' picture¶

Samples from the untrained model: uniform gibberish over all 65 characters. Save it — the contrast with Step 6 is the point.

In [ ]:
torch.manual_seed(0)
ctx = torch.tensor([[stoi['\n']]])
print(decode(model.generate(ctx, 300)[0].tolist()))

LayerNorm geometry (Lecture 5, Prop. 4.2)¶

The normalization step is orthogonal projection onto $\mathbf 1^\perp$ followed by radial projection onto the sphere of radius $\sqrt d$ — so its image is a $(d-2)$-sphere, and sublayers see direction, not magnitude.

In [ ]:
ln = nn.LayerNorm(128, elementwise_affine=False)
# Use magnitudes in [2.5, 7.5] so every row's variance is >> eps=1e-5;
# see the caveat below for why that matters.
u = torch.randn(1000, 128) * (2.5 + 5 * torch.rand(1000, 1))
y = ln(u)
print(f'mean of output rows  : {y.mean(-1).abs().max():.2e}   (in 1-perp)')
print(f'norms of output rows : {y.norm(dim=-1).min():.4f} .. '
      f'{y.norm(dim=-1).max():.4f}   (sqrt(d) = {128**0.5:.4f})')
assert torch.allclose(y.norm(dim=-1), torch.full((1000,), 128**0.5), atol=1e-3)
assert torch.allclose(ln(3.7 * u), y, atol=1e-4)
print('all rows on the sphere, and LN(cu) = LN(u)  OK')

A caveat the proof hides, and the code must not. Proposition 4.2 ignores the $\varepsilon$ in $\sqrt{\sigma^2+\varepsilon}$. That $\varepsilon = 10^{-5}$ is what keeps LayerNorm defined when a row is constant ($\sigma = 0$), but it means the sphere-radius and scale-invariance properties hold only approximately, and hold well precisely when $\sigma^2 \gg \varepsilon$ — which is why the cell above deliberately keeps every row's magnitude bounded away from zero. Rescale a row to near-zero variance (try u * 1e-3) and both assertions fail: the normalized vector falls short of the sphere and $\mathrm{LN}(cu) \neq \mathrm{LN}(u)$. In a trained model the residual stream has healthy variance, so this corner never bites — but it is the difference between the clean theorem and the floating-point object.

→ Continue with Step 6: train it.