Step 6 — Training the Baby GPT

Released with Lecture 6 · ~3 hours active + background training time · Starts from: Step 5 (or solutions/step-05.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

Train the model. This step is where the course’s optimization mathematics meets the empirical craft of making it work: you will write a real training loop (AdamW, warmup + cosine decay, gradient clipping), diagnose its behaviour with train/val curves, and run one controlled architecture experiment. At the end, your GPT writes its first Shakespeare.

Compute note: the course config (~0.8M params, T_max=64) trains to respectability in ~10–20 min on a laptop CPU and ~2 min on a free Colab T4 GPU. Use Colab this week if you can (model.to('cuda'), and move each batch too); but CPU works.

On paper first

1. AdamW, stated precisely. Write the update: first/second moment EMAs, bias correction, $\theta \mathrel{-}= \eta\, \hat m/(\sqrt{\hat v} + \varepsilon) + \eta\lambda\theta$. Show that (ignoring $\varepsilon$, at stationarity of the EMAs) each coordinate’s step size is $\approx \eta \cdot \operatorname{sign-like}$, invariant to rescaling that coordinate’s gradient — Adam is a diagonal preconditioner. Contrast with SGD’s dependence on curvature scale.
2. Why warmup + decay? From Lecture 6: early on, moment estimates are noise and the landscape is at its roughest — warmup keeps steps small; late, decaying $\eta$ trades exploration for convergence. Sketch the schedule you’ll implement: $\eta(t) = \eta_{\max}\cdot\min(t/t_w,\; \tfrac12(1+\cos(\pi (t-t_w)/(t_f-t_w))))$.
3. What can go wrong with minibatch estimates: the loss you plot is a noisy estimator of the true loss. How does batch size enter its variance, and what does that imply about judging progress from the jagged raw curve vs an EMA of it?

Tasks

6.1 Data pipeline

Train/val split of the token stream (90/10), and a batch sampler:

def get_batch(split, B=64, T=64):                        # default parameter values
    data = train_ids if split == 'train' else val_ids   # conditional (ternary) expression
    ix = torch.randint(len(data) - T - 1, (B,))          # (B,) = a one-element tuple
    x = torch.stack([data[i:i+T] for i in ix])           # list comprehension of tensor slices
    y = torch.stack([data[i+1:i+T+1] for i in ix])       # shifted by one
    return x.to(device), y.to(device)                    # returns a tuple

t = (5,)              # one-element tuple: the comma does the work
not_t = (5)           # just a parenthesized number
print(type(t), type(not_t))
print(len(t))
import numpy as np
print(np.zeros((3,)).shape)    # shape (3,): a vector of length 3
print(np.zeros((3, 2)).shape)  # shape (3, 2): a matrix

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = GPT(...).to(device)   # moves all parameters AND buffers
x, y = get_batch('train')     # already on device — get_batch handled it

Note y is x shifted: every position is a training example — one batch contains $B \times T = 4096$ next-token problems, and the causal mask is what makes computing them all in one forward pass legitimate. (Connect this to Lecture 4: this is why masking, not bidirectional attention, is what a language model needs.)

6.2 Evaluation done right

@torch.no_grad()                    # decorator
def estimate_loss(model, iters=100):
    model.eval()
    # dict comprehension over ('train','val'); *get_batch(s) unpacks a tuple into arguments
    out = {s: torch.mean(torch.tensor([model(*get_batch(s))[1].item()
                                       for _ in range(iters)])) for s in ('train','val')}
    model.train()
    return out

def shout(fn):                  # takes a function...
    def wrapper(name):
        return fn(name).upper() + "!"
    return wrapper              # ...returns a wrapped version (a closure)

@shout                          # greet = shout(greet)
def greet(name):
    return "hello " + name

print(greet("kate"))

def volume(w, h, d):
    return w * h * d
box = (2, 3, 4)
print(volume(*box))       # the tuple spread into three arguments
print(volume(2, 3, 4))    # exactly equivalent
def measurements():
    return 5, 6, 7        # a function returning a tuple...
print(volume(*measurements()))   # ...unpacked straight into another call

Averaging many batches: your answer to “on paper” 3, in code.

6.3 The training loop

Write it in full — this loop is an artifact you’ll reuse for years:

• torch.optim.AdamW(model.parameters(), lr=..., weight_decay=0.1, betas=(0.9, 0.99)) — permitted now that you’ve written SGD (Step 3) and derived Adam (paper 1).
• Your warmup + cosine schedule from paper 2, applied per step (either by hand-setting group['lr'] — recommended, it’s transparent — or via a scheduler). $t_w \approx 100$, $t_f = 5000$ steps, $\eta_{\max} \approx 3\times10^{-3}$ for the course config.
• Gradient clipping: torch.nn.utils.clip_grad_norm_(params, 1.0), and log the pre-clip norm every step — the spikes you’ll see are Lecture 6’s loss-landscape roughness made visible.
• Evaluate every 250 steps; keep (step, train, val, grad-norm, lr) records.

Train 5000 steps. While it runs, predict: where should the loss start (you know: $\ln 65$), and what’s the best it could conceivably reach (your Step 1 table: Shannon says ~1 bit/char ≈ 0.69 nats)?

6.4 Read the curves

Produce and annotate (in prose, in the notebook) four plots:

1. Train & val loss vs step — mark where they separate. Do they, at 0.8M params on 1M characters? Given ~0.8 params per training character, is the amount of overfitting you see surprising in the classical bias–variance picture? (Lecture 6’s double-descent discussion is exactly about this.)
2. Loss vs step with your lr schedule overlaid — find the warmup kink and the late-decay flattening.
3. Gradient norm vs step — spikes, and whether clipping engaged.
4. Sample evolution: generate 200 characters from checkpoints at step 0, 250, 1000, 5000. Watch the model discover, in order: characters’ frequencies, words, line structure, CHARACTER NAME: stage format. This ordering is a scaling-law fact you’ll meet again in Step 7.

6.5 One controlled experiment

Pick one knob and vary it with everything else fixed, 3–4 values, tabulated (val loss @ 5000 steps, wallclock, params):

• depth $L \in \{1, 2, 4, 8\}$ at fixed width, or
• width $d \in \{32, 64, 128, 256\}$ at fixed depth, or
• context $T_{\max} \in \{8, 16, 64, 256\}$, or
• heads $H \in \{1, 2, 4, 8\}$ at fixed $d$.

Write three sentences on the trend and whether it’s monotone. These tables get pooled in Lecture 7 as our homemade scaling-law data.

Checkpoints

• Loss starts at 4.30 (just above $\ln 65 = 4.174$) and ends at val 1.5178 = 2.190 bits/char for the course config at 5,000 steps. Reference trajectory (yours will differ in the last digits by seed):

  step	0	500	1000	2000	3000	4000	4500	5000
  train	4.306	1.656	1.471	1.348	1.271	1.204	1.185	1.166
  val	4.304	1.813	1.661	1.568	1.521	1.519	1.509	1.518

  Update your bits/char table (Step 1 bigram: 3.54; Step 3 MLP: 2.86; here: 2.19).

• Two things in that table are worth more than the headline number, and plot 1 should make you find them:

  1. Validation stops improving around step 3000 and the minimum is at 4500, not at the end. The last 2,000 steps drop train loss by 0.106 and move val by 0.003 in the wrong direction — roughly 40% of the run’s compute bought nothing. Your final model is therefore not your best model. Real pipelines checkpoint on validation and keep the best; ours doesn’t, which is itself the lesson. Add best-checkpointing if you like — it is five lines.
  2. Gradient norms rise late (0.298 at step 2500 to 0.414 at 4500) while the learning rate falls tenfold. Sharpness increases as the step size shrinks. See Lecture 6 §5.4 — this is edge-of-stability behaviour visible in your own run.

• Total wallclock: about 73 minutes on 16 CPU threads; a few minutes on a Colab T4. For reference, nanoGPT’s larger baby GPT ($d{=}384$, $L{=}6$, $T_{\max}{=}256$) reaches ≈1.48 on this dataset — a good Step 8 target if you have GPU time.

• Samples are unmistakably Shakespeare-shaped. Genuine output from the reference run at temperature 1.0:

  How doest thou lefthard woman, we'll be time:
  we are ta'en,--Citizens are every ruth.
  
  Citizen:
  Either farewell; here's none bload in his general.
  Had it no further than you will coward of my shoulder,
  That his fresh dead quarter death,
  Let them glad-hearts to be her bones;
  Who pleaded thousand Volsces, peace, obeys
  That Barnardine they have stripp'd again.
  
  WARWICK:
  Stands those way-dength in me, Montague,
  Nor I, the kind dark of that which he had

  Compare with Step 5’s untrained gibberish and Step 3’s LOET:. The play format is exact; the speakers (WARWICK, Citizen) are real characters, spelled correctly — which the three-character-context MLP of Step 3 could not do, since a name does not fit in its window. Even Montague and Barnardine are right, and Volsces is from Coriolanus. Attention buying long-range identity is visible in the output, exactly as promised at the end of Step 3. Meaning, of course, is still absent: clauses are locally grammatical and globally nonsense. Save your favourite for the Step 8 showcase.

• You can point at the step in your loop where each of AdamW / schedule / clipping acts, and say what lecture result each implements.

If you’re stuck

• Loss plateaus ≈ 2.4: model is behaving bigram-ly — usually the target shift is wrong (y not offset by one) or the mask is upper- instead of lower-triangular.
• Loss explodes mid-run: check plot 3 — if grad norms spiked first, lower $\eta_{\max}$ or check clipping is actually applied after backward() and before step().
• Colab: RuntimeError: expected ... cuda — a tensor (often the mask buffer or a fresh torch.arange) didn’t move to the device.

Going further (optional)

• Dropout (p=0.1 on attention weights and MLP output): retrain, compare the train/val gap. The course config barely needs it — find a config that does (hint: shrink the dataset ×10).
• Reproduce the loss spike → recover phenomenon: raise $\eta_{\max}$ until training visibly self-heals after spikes, then until it doesn’t. You are walking the edge of stability from Lecture 6.
• Estimate your model’s compute in FLOPs ($\approx 6 \cdot \text{params} \cdot \text{tokens}$ — derive the 6!) and place your run on a Kaplan-style loss-vs-compute plot alongside classmates’ Step 6.5 tables.

Catch-up

Run solutions/step-05.ipynb to get a verified model class, then do this step in full — training is the step least worth skipping, and needs only Lecture 6.