Step 6 — Training the Baby GPT

Solution notebook. The one with the long cell: the 5,000-step training run takes ~1 hour on a 16-thread CPU and a few minutes on a Colab T4 GPU. Every other cell is seconds. The reference numbers quoted in the prose are from the run that produced the course's official checkpoint values; sampling and initialization are seeded, so your numbers should land within noise of them (minibatch order differs across hardware, so expect the last digits to move).

At the end this notebook saves step6_model.pt, which Steps 7 and 8 load.

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.

# !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}')

The model (Step 5's solution, reproduced)

Identical to the Step 5 notebook, including the $1/\sqrt{2L}$ scaling of the residual-writing projections — Lecture 6 §1 derives why: the stream receives $2L$ writes, and variances add.

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

torch.manual_seed(1337)
model = GPT(V=65, T=64, d=128, H=4, L=4)
n_params = sum(p.numel() for p in model.parameters())
print(f'{n_params:,} parameters')
assert n_params == 816_640    # Step 5's checkpoint, to the integer

6.1 Data pipeline

Train/val split, then a batch sampler. y is x shifted by one: every position of every sequence is a training example, so one batch of shape (64, 64) contains 4,096 next-token problems. The causal mask (Lecture 4, Prop. 5.1) is what makes computing them all in one forward pass legitimate — position $i$'s prediction provably never saw positions $> i$.

n_split = int(0.9 * len(ids))
train_ids, val_ids = ids[:n_split], ids[n_split:]
B, T = 64, 64

def get_batch(split):
    data = train_ids if split == 'train' else val_ids
    ix = torch.randint(len(data) - T - 1, (B,))
    x = torch.stack([data[i:i+T] for i in ix])
    y = torch.stack([data[i+1:i+T+1] for i in ix])
    return x, y

xb, yb = get_batch('train')
assert torch.equal(xb[:, 1:], yb[:, :-1])   # y really is x shifted
print(xb.shape, yb.shape)

6.2 Evaluation done right

One minibatch loss is a noisy estimate (Lecture 3, Prop. 5.2 — variance $\propto 1/B$). We average many batches, in eval mode, under no_grad.

@torch.no_grad()
def estimate_loss(iters=100):
    model.eval()
    out = {}
    for split in ('train', 'val'):
        losses = [model(*get_batch(split))[1].item() for _ in range(iters)]
        out[split] = sum(losses) / len(losses)
    model.train()
    return out

6.3 The training loop

The three ingredients, each earned in lecture:

• AdamW (Lecture 6 §2): diagonal preconditioning so rare-feature coordinates still move — the fix for exactly the pathology measured in Step 2 — with decoupled weight decay.
• Warmup + cosine decay (Lecture 6 §3), set by hand on the optimizer's param group so nothing is hidden.
• Gradient clipping (Lecture 6 §4) at norm 1.0 — and we log the pre-clip norm every evaluation, because whether the safeguard ever fires is data worth having. (Spoiler from the reference run: it never does; the norm stays in 0.30–0.41. Instrumenting a safeguard to discover it never fires beats assuming it was load-bearing.)

import math

MAX_STEPS, WARMUP = 5000, 100
LR_MAX, LR_MIN = 3e-3, 3e-4

def lr_at(t):
    if t < WARMUP:
        return LR_MAX * t / WARMUP
    r = (t - WARMUP) / (MAX_STEPS - WARMUP)
    return LR_MIN + 0.5 * (1 + math.cos(math.pi * r)) * (LR_MAX - LR_MIN)

opt = torch.optim.AdamW(model.parameters(), lr=LR_MAX,
                        weight_decay=0.1, betas=(0.9, 0.99))

import time
history = []          # (step, train, val, grad_norm, lr)
t0, gnorm = time.time(), float('nan')

for step in range(MAX_STEPS + 1):
    if step % 500 == 0 or step == MAX_STEPS:
        e = estimate_loss(50 if step < MAX_STEPS else 200)
        history.append((step, e['train'], e['val'], gnorm, lr_at(step)))
        print(f"{step:>5}  train {e['train']:.4f}  val {e['val']:.4f}"
              f"  ({e['val']/math.log(2):.4f} bits/char)"
              f"  gnorm {gnorm:.3f}  lr {lr_at(step):.2e}"
              f"  {time.time()-t0:.0f}s", flush=True)
    if step == MAX_STEPS:
        break
    for g in opt.param_groups:
        g['lr'] = lr_at(step)
    _, loss = model(*get_batch('train'))
    opt.zero_grad(set_to_none=True)
    loss.backward()
    gnorm = float(torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0))
    opt.step()

torch.save(model.state_dict(), 'step6_model.pt')
print('saved step6_model.pt')

Reference trajectory (validation): 4.304 → 1.813 → 1.661 → 1.609 → 1.568 → 1.538 → 1.521 → 1.520 → 1.519 → 1.509 → 1.518. Final: 1.5178 nats = 2.190 bits/char. Your bits/char table gains its biggest single improvement: 3.54 (bigram) → 2.86 (MLP) → 2.19.

6.4 Read the curves

Four plots, and the prose you attach to them matters more than the plots.

import matplotlib.pyplot as plt
steps, tr, va, gn, lrs = zip(*history)
fig, ax = plt.subplots(2, 2, figsize=(11, 7))
ax[0,0].plot(steps, tr, label='train'); ax[0,0].plot(steps, va, label='val')
ax[0,0].axhline(math.log(65), ls=':', c='gray', label='ln V')
ax[0,0].set_title('loss'); ax[0,0].legend()
ax[0,1].plot(steps, [v-t for t,v in zip(tr,va)], c='crimson')
ax[0,1].set_title('generalization gap (val − train)')
ax[1,0].plot(steps[1:], gn[1:], c='darkorange')
ax[1,0].axhline(1.0, ls=':', c='gray', label='clip threshold')
ax[1,0].set_title('pre-clip gradient norm'); ax[1,0].legend()
ax[1,1].plot(steps, lrs, c='seagreen'); ax[1,1].set_title('learning rate')
for a in ax.flat: a.set_xlabel('step')
plt.tight_layout()

What the reference run shows, panel by panel:

1. Loss. Starts at $\ln 65$ (the model knows only the vocabulary size), drops to 1.81 within 500 steps, then grinds. Roughly linear in $\log(\text{step})$ — a within-run shadow of Lecture 7's scaling laws.
2. Gap. Widens monotonically, 0.158 → 0.352. At ~0.8 parameters per training character the model can memorize, and is beginning to. Also look closely at val alone: it stops improving near step 3000 and its minimum is at 4500, not the end — the last 40% of the run bought ~nothing, and the final model is not the best model. Real pipelines checkpoint on validation for exactly this reason.
3. Gradient norm. Never approaches the clip threshold — clipping was pure insurance here. But it rises late (0.298 → 0.414) while the learning rate falls tenfold: the iterate is settling into sharper curvature as steps shrink. That is edge-of-stability behaviour (Lecture 6 §5.4), visible on a laptop.
4. Learning rate. The warmup spike is invisible at this resolution (100 steps); the cosine decay is not.

Samples

The payoff cell. Reference output includes real, correctly spelled speaker names — WARWICK:, Citizen:, Montague, Barnardine, Volsces all occur in the corpus (99, 98, 46, 15, 18 times). Compare Step 3's LOET:, a name-shaped fake — three characters of context cannot hold an identity, 64 can. That difference is attention, visible in the output.

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

6.5 One controlled experiment

The reference choice: depth $L \in \{1, 2, 4, 8\}$ at fixed width. This cell takes several times longer than the main run — leave it for a GPU session, or shorten MAX_STEPS. (Left unexecuted here; pool your results with the class for Step 7's scaling plot.)

for L in (1, 2, 4, 8):
    torch.manual_seed(1337)
    model = GPT(V=65, T=64, d=128, H=4, L=L)
    # ... identical training loop ...

Whatever knob you choose, hold everything else fixed, tabulate (val loss, params, wallclock), and write three sentences on whether the trend is monotone — Step 3's $k$-sweep should have taught you not to assume it is.

→ Continue with Step 7: a real tokenizer and a real sampler for this trained model.