Notebook 09 · Watch it learn

Train it on Shakespeare, and the gibberish becomes a play.

The same four-line loop as always -- only the model in the middle is now a real Transformer. The loss falls from random (4.17) through the bigram ceiling (2.45) to 1.59, and the text sharpens with it.

This notebook needs PyTorch and a GPU A browser can't run PyTorch, so the cells below are shown with the outputs of a real run (Apple-GPU). Download the notebook at the bottom of the page to run and train it yourself.

Everything is in place. Now we train the model from notebook 08 on Shakespeare and watch it learn. The loop is the same four lines as always (notebook 08's bigram used them too); the only difference is that the model in the middle is now a real Transformer.

What to watch for: the loss falling from about 4.2 (random) towards about 1.6, and the generated text turning from noise into something resembling a play, with character names, line breaks, dialogue, and mostly real words, all from one model and a few minutes on a GPU.

python
# Colab setup -- fetch the files this notebook needs.
# (Does nothing when run locally in the course folder.)
import os, urllib.request
BASE = ("https://raw.githubusercontent.com/waze"
        "emlabs/llm-book-code/main/")
for f in ['gpt.py', 'data/input.txt']:
    if not os.path.exists(f):
        d = os.path.dirname(f)
        if d: os.makedirs(d, exist_ok=True)
        urllib.request.urlretrieve(BASE + f, f)
        print("downloaded", f)
Line by line: what each line does
  • import os, urllib.request: two standard-library toolboxes: checking files on disk, and downloading from the web.
  • BASE = (...): the web address of this course's folder on GitHub, split over two lines (Python glues adjacent strings together).
  • for f in [...]: loop over the file names this notebook needs.
  • if not os.path.exists(f): only download what is missing -- running locally, everything already exists, so nothing happens.
  • os.makedirs(d, exist_ok=True): create the folder for the file if it has one (exist_ok means don't complain if it's already there).
  • urllib.request.urlretrieve(...): download the file and save it under the same name here.
python
import torch, time, math
from pathlib import Path
from gpt import GPT, GPTConfig
torch.manual_seed(1337)
device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
print("training on:", device)
Line by line: what each line does
  • import torch, time, math: PyTorch, plus time for the stopwatch and math for ln(65).
  • from gpt import GPT, GPTConfig: the model classes you toured in notebook 08.
  • torch.manual_seed(1337): pin the randomness so training is reproducible.
  • The device line picks Apple GPU → NVIDIA → CPU, as before.
outputtraining on: mps

Data and batching

Encode the text (notebook 02), hold out the last 10% as a validation set, and draw random block_size chunks for each step.

Why hold out that 10%? It is the model's practice exam: text that it never trains on. We watch two losses, the train loss (on text it studies) and the validation loss (on the held-out text). If both fall together, the model is genuinely learning the language. If the train loss keeps dropping while the validation loss stalls or creeps up, the model is memorizing the training text instead of learning patterns that carry over. That failure is called overfitting, and the gap between the two numbers is how you spot it.

python
text = Path("data/input.txt").read_text()
chars = sorted(set(text)); vocab_size = len(chars)
stoi = {c:i for i,c in enumerate(chars)}; itos = {i:c for c,i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]; decode = lambda l: "".join(itos[i] for i in l)

data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9*len(data)); train_data, val_data = data[:n], data[n:]

block_size = 128     # context length
batch_size = 32      # sequences per step
def get_batch(split):
    d = train_data if split == "train" else val_data
    ix = torch.randint(len(d) - block_size, (batch_size,))
    x = torch.stack([d[i:i+block_size] for i in ix])
    y = torch.stack([d[i+1:i+1+block_size] for i in ix])
    return x.to(device), y.to(device)
print("train tokens:", len(train_data), "| val tokens:", len(val_data))
Line by line: what each line does
  • The tokenizer and split are identical to notebook 02 (read the file, build the tables, encode, hold out the last 10%), in PyTorch tensors.
  • block_size = 128: the context window is now 128 characters (the bigram saw 1).
  • batch_size = 32: 32 sequences per training step.
  • def get_batch(split):: same recipe as before: 32 random 128-long chunks as context x, the same chunks shifted one right as targets y, both shipped to the GPU.
outputtrain tokens: 1003854 | val tokens: 111540

Model, optimizer, and an honest loss estimate

We build the same model of about 0.8 million numbers, now with a little dropout (set to 0.1). Dropout randomly switches off 10% of the signals on each training step, which stops the model from leaning too heavily on any single path. It is a bit like studying with random pages of your notes covered, so that you learn the material itself rather than memorizing one route to the answer. It is a standard guard against the overfitting just described, and it is active only during training.

estimate_loss averages the loss over many batches (150 of them here) before printing. A single batch is a noisy reading, because some stretches of text are simply easier than others, so averaging gives an honest number that can be compared fairly from one step to the next.

python
cfg = GPTConfig(vocab_size=vocab_size, block_size=block_size,
                n_layer=4, n_head=4, n_embd=128, dropout=0.1)
model = GPT(cfg).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
print(f"parameters: {model.num_params():,}")

eval_iters = 150
@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ["train", "val"]:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            x, y = get_batch(split)
            _, loss = model(x, y)
            losses[k] = loss.item()
        out[split] = losses.mean().item()
    model.train()
    return out
Line by line: what each line does
  • cfg = GPTConfig(..., dropout=0.1): the same model as notebook 08, now with 10% dropout switched on for training.
  • model = GPT(cfg).to(device) / optimizer = torch.optim.AdamW(..., lr=1e-3): build it on the GPU and hand its knobs to AdamW.
  • @torch.no_grad(): decorator: everything inside runs without recording gradients (it's measurement, not learning).
  • model.eval() / model.train(): flip the model's mode: eval turns dropout off for an honest measurement, train turns it back on. Forgetting this pair is a classic bug.
  • losses = torch.zeros(eval_iters): a 150-slot tensor; losses[k] = loss.item() fills slot k with a plain number.
  • out[split] = losses.mean().item(): store each split's average loss in a dictionary: {"train": ..., "val": ...}.
outputparameters: 824,897

The training loop

The same five moves as always: sample a batch, forward, zero_grad, backward, step. Every eval_interval steps we pause to estimate the train and validation losses and print them. Reading the output below, both numbers drop quickly at first and then slow down, and the validation loss sits a little above the train loss (which is normal: the practice exam is always a touch harder than the homework). It runs in a couple of minutes on an Apple GPU; raise max_iters (and the model size in cfg) for sharper text.

python
max_iters = 3000
eval_interval = 500
history = []

t0 = time.time()
for it in range(max_iters + 1):
    if it % eval_interval == 0:
        l = estimate_loss()
        history.append((it, l["train"], l["val"]))
        print(f"step {it:4d} | train {l['train']:.3f} | val {l['val']:.3f} | {time.time()-t0:5.1f}s")
    x, y = get_batch("train")
    _, loss = model(x, y)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
print(f"\ndone in {time.time()-t0:.1f}s")
Line by line: what each line does
  • max_iters = 3000; eval_interval = 500; history = []: how many steps to train, how often to measure, and a list to record the loss curve.
  • t0 = time.time(): start the stopwatch.
  • for it in range(max_iters + 1):: 3001 passes; the + 1 makes the final evaluation at step 3000 happen.
  • if it % eval_interval == 0:: every 500th step (% is remainder), pause and measure train/val loss properly.
  • history.append((it, l["train"], l["val"])): record a (step, train, val) triple for the plot.
  • print(f"step {it:4d} | train {l['train']:.3f} | ..."): formatted print: 4d pads the step number, .3f fixes 3 decimals.
  • x, y = get_batch("train") → forward → zero_grad → backward → step: the same four-line heartbeat as every notebook. Only the model in the middle changed since the bigram.
outputstep 0 | train 4.177 | val 4.174 | 1.4s step 500 | train 1.971 | val 2.045 | 11.1s step 1000 | train 1.674 | val 1.830 | 19.4s step 1500 | train 1.546 | val 1.717 | 28.1s step 2000 | train 1.473 | val 1.654 | 36.9s step 2500 | train 1.421 | val 1.623 | 45.8s step 3000 | train 1.382 | val 1.593 | 54.7s done in 54.7s
python
import matplotlib.pyplot as plt
steps = [h[0] for h in history]
plt.figure(figsize=(6,4))
plt.plot(steps, [h[1] for h in history], "o-", label="train")
plt.plot(steps, [h[2] for h in history], "o-", label="val")
plt.axhline(math.log(vocab_size), ls="--", c="gray", label="random baseline ln(vocab)")
plt.axhline(2.452, ls=":", c="crimson", label="bigram optimum (nb 03)")
plt.xlabel("step"); plt.ylabel("loss"); plt.legend(); plt.title("GPT training: well below the bigram ceiling"); plt.show()
Line by line: what each line does
  • import matplotlib.pyplot as plt: the plotting toolbox.
  • steps = [h[0] for h in history]: pull the step numbers out of the recorded triples; h[1] / h[2] pull train / val losses.
  • plt.plot(steps, [h[1] for h in history], "o-", label="train"): the train-loss curve, dots joined by lines.
  • plt.axhline(math.log(vocab_size), ls="--", ...): a horizontal reference line at ln(65) = 4.17 (random guessing). ls is line style.
  • plt.axhline(2.452, ls=":", ...): the bigram optimum from notebook 03. Watching the GPT curve dive under that line is the whole point of the picture.
outputoutput figure

What the numbers mean, and then hearing it speak

A handy way to feel a loss value is to compute e^loss, a quantity called the perplexity. Roughly, it answers "how many next-characters is the model still torn between?"

  • random guessing: loss 4.17, so e^4.17 is about 65 choices (all of them, since it knows nothing),
  • the bigram ceiling: loss 2.45, so about 12 choices,
  • our GPT: validation loss 1.59, so about 5 choices.

So training took the model from "could be any of 65 characters" down to "one of about five." That is what breaking through the bigram ceiling looks like: attention let the model use 128 characters of context instead of just one.

Now we sample from it. Compare the result with the bigram's output in notebook 03 and the untrained gibberish in notebook 08. It is the same machinery; it has simply learned.

4.17
random baseline (ln 65) · perplexity 65
2.45
bigram optimum (nb 03) · perplexity ~12
1.59
our GPT, val · perplexity ~5
Interactive · watch it learn

The loss falls, the text sharpens

Training drives the loss down from the random baseline (4.17), through the bigram ceiling (2.45), toward 1.59 — and the sample underneath sharpens from noise into almost-Shakespeare as it goes. Press Play.

python
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(model.generate(context, max_new_tokens=500)[0].tolist()))
Line by line: what each line does
  • torch.zeros((1, 1), dtype=torch.long, device=device): the prompt: one sequence containing one token, id 0 (newline). A neutral start.
  • model.generate(context, max_new_tokens=500): 500 dice-rolls of the trained model, each fed back in.
  • [0].tolist(): take the first (only) sequence and convert ids to a plain list for decode.
output SICINGHAM: My lord, you shall palt breate; Nor be neight your cousin y soverey to The you often us. A say light Farewell your most penfitly love Retuor fortune with merry grued choose Pent his -chost nonal! Second Maken: I shall bous coffpent the wearty unating. When that wear how is that? WARD: I do have you, Lord Aladies, 'tis, though thinks a bree of sleet from of your broughins lord. PAULINA: Lenford, that? BRUTUS: Master you I do not hence, Justiculza me felp in thy seaths, I fifthes i

Look inside: what the trained heads learned

Notebooks 05 and 06 drew this exact picture for untrained, random heads, and every head looked alike. Now the model has trained. Recompute block 0's four attention maps on a real line of text and they have clearly specialised: one head stays on the diagonal (attend to the current character), others lean to earlier characters or lock onto the line's opening. Each head has learned a different way to gather context -- and this structure is what turns the gibberish into Shakespeare.

python
# Recompute block 0's attention on a real snippet, one map per head. Notebooks
# 05-06 drew this for random weights; here the weights are trained.
model.eval()
snippet = "First Citizen:\nWe are"
ids = torch.tensor([[stoi[c] for c in snippet]], device=device)
Tn = len(snippet); labs = [c.replace("\n", "\\n") for c in snippet]

fig, axes = plt.subplots(2, 2, figsize=(8.5, 8))
for H, ax in enumerate(axes.flat):
    with torch.no_grad():
        x = model.token_embedding(ids) + model.position_embedding(torch.arange(Tn, device=device))
        blk = model.blocks[0]; xn = blk.ln1(x); h = blk.sa.heads[H]
        k, q = h.key(xn), h.query(xn)
        w = (q @ k.transpose(-2, -1)) * k.shape[-1] ** -0.5
        w = w.masked_fill(h.tril[:Tn, :Tn] == 0, float("-inf"))
        A = torch.softmax(w, dim=-1)[0].cpu().numpy()
    ax.imshow(A, cmap="viridis")
    ax.set_xticks(range(Tn)); ax.set_xticklabels(labs, fontsize=6)
    ax.set_yticks(range(Tn)); ax.set_yticklabels(labs, fontsize=6)
    ax.set_title(f"block 0, head {H}", fontsize=10)
fig.suptitle("what the trained heads attend to (row = current char, col = the char it looks at)", fontsize=11)
plt.tight_layout(); plt.show()
Line by line: what each line does
  • model.eval(): turn off dropout so the attention we read is the clean, deterministic version.
  • For each of block 0's four heads it recomputes the attention by hand: embed the snippet, LayerNorm it, then softmax(q @ k.T / sqrt(head_size)) with the future masked -- exactly the formula from notebook 05.
  • Each heatmap row is a character; the bright cells along that row are the earlier characters it attends to.
  • Unlike the near-identical random heads in notebook 06, these have specialised: compare the four -- diagonal (the current character), one-step-back, and longer-range -- each head learned a different job.
  • Stack four such heads across four blocks and train for 3000 steps, and this learned structure is what turns gibberish into Shakespeare.
outputoutput figure

Save the trained model

We store the weights, the config, and the tokenizer, so that notebook 10 can load this exact model and explore sampling without retraining it.

python
import os
os.makedirs("checkpoints", exist_ok=True)
torch.save({"model": model.state_dict(), "config": cfg, "stoi": stoi, "itos": itos},
           "checkpoints/gpt_shakespeare.pt")
print("saved -> checkpoints/gpt_shakespeare.pt")
Line by line: what each line does
  • os.makedirs("checkpoints", exist_ok=True): create the folder; the flag means "fine if it already exists."
  • model.state_dict(): every learned weight, as a dictionary of tensors. This is the trained model.
  • torch.save({"model": ..., "config": cfg, "stoi": stoi, "itos": itos}, "checkpoints/gpt_shakespeare.pt"): bundle weights + config + tokenizer into one file, so notebook 10 can reload exactly this model without retraining.
outputsaved -> checkpoints/gpt_shakespeare.pt

Recap

You trained a GPT that you built from scratch. The loss fell from about 4.2 to well below the bigram ceiling, and the samples look like Shakespeare because predicting the next character well forced the model to learn spelling, character names, and the shape of dialogue. That is the course's central idea, made real.

Next, notebook 10 covers how to control the way it generates (temperature, top-k, and top-p), and offers a tour of what separates this from a frontier model.

⬇ Download this lesson as a notebook — 09_train_gpt.ipynb