Build a Transformer from Scratch — The Capstone

The capstone is a small decoder-only transformer trained end-to-end on a character-level language modeling task. It reads Shakespeare and generates new Shakespeare. It is small enough to train on a laptop in under 10 minutes, and correct enough that a bigger dataset and longer training get you a real LM.

You've implemented attention, multi-head splits, positional encodings, encoder and decoder blocks, BERT and GPT losses, MoE, and KV cache. Here they have to work together on a real task.

This is the "nanoGPT" of the course. It isn't original: Karpathy's 2023 nanoGPT tutorial is the reference implementation, and a lot of students write it at least once. We lift the shape and retool it around what we've covered.

Here is the architecture, annotated:

input tokens (B, N)
   │
   ▼
token embedding + positional embedding  ◀── Lesson 04 (RoPE option)
   │
   ▼
┌──── block × L ────────────────────┐
│  RMSNorm                          │  ◀── Lesson 05
│  CausalSelfAttention (causal)     │  ◀── Lesson 03 + 07 (causal mask)
│  residual                         │
│  RMSNorm                          │
│  SwiGLU FFN                       │  ◀── Lesson 05
│  residual                         │
└────────────────────────────────── ┘
   │
   ▼
final RMSNorm
   │
   ▼
lm_head (tied to token embedding)
   │
   ▼
logits (B, N, V)
   │
   ▼
shift-by-one cross-entropy            ◀── Lesson 07

What we ship

• Hyperparameters: plain variables at the top of try_train() (block_size, d_model, n_heads, n_layers, and so on). There is no config class.
• CausalSelfAttention: causal and batched, with the mask-and-softmax written out by hand so you can see every step.
• SwiGLUFFN: a modern FFN.
• Block: pre-norm, residual-wrapped attention + FFN.
• GPT: embeddings, stacked blocks, LM head, generate().
• Training loop with AdamW (lr 3e-4, betas 0.9/0.95, weight decay 0.1) and gradient clipping at 1.0. The learning rate is constant.
• Char-level tokenizer on Shakespeare text, plus a parameter-count preview table printed before training.

What we don't ship

• RoPE: implemented conceptually in Lesson 04. Here we use learned positional embeddings for simplicity, and the exercises ask you to swap in RoPE.
• KV cache during generation: each generation step recomputes attention over the full prefix. That's slower but simpler, and the exercises ask you to add a KV cache.
• Flash Attention: main.py builds the attention matrix explicitly instead of calling F.scaled_dot_product_attention. Swapping it in is an easy extension, and PyTorch 2.0+ dispatches to a fused kernel when the inputs allow.
• A learning-rate schedule and bf16 autocast: the loop runs at a constant LR in default precision. Exercise 6 adds both.
• MoE: each block has a single FFN. You saw MoE in Lesson 11.

Target metrics

main.py as shipped is a smoke test: 3 layers, 4 heads, d_model=64, block_size=64, batch size 16, 500 steps. Its docstring says 4 layers, d_model=128, seq_len=128 and about 2 minutes, which disagrees with the code. The code is what runs, so trust the config block inside try_train(). The larger configuration below is the one the exercises scale you up to, and the numbers come from that configuration, not from the shipped one.

These numbers come from actually running the script. The larger configuration is n_layers=4, d_model=128 (803,584 parameters), everything else as shipped (batch 16, block 64, constant LR 3e-4), on the full tinyshakespeare.txt placed next to main.py, on CPU. Seeds are fixed, so a rerun of the same configuration prints the same numbers on the same machine.

• Step 0 prints a loss near 88, not the ln 65 ≈ 4.2 you would expect from a uniform guess. The token embedding is initialised from N(0, 1) and the LM head is tied to it, so the first logits are huge. It falls below 3.3 by step 100, so ignore the first row.
• After 2,000 steps (about 1.5 minutes on a multi-core CPU) train loss is about 2.1 and val loss about 2.2. After 5,000 steps they are about 1.7 and 1.8.
• Sampled output has Shakespeare's rhythm: line breaks, capitalised speaker tags (MARTIUS: shows up at 5,000 steps), and mostly word-shaped fragments with few real words. Do not expect fluent text from this budget.
• Val loss (held-out final 10% of text) stays within about 0.2 of train loss through 5,000 steps and is still falling, so overfitting is not the limit at this size and budget.

You need PyTorch for this one. Install torch (a CPU build is fine) and open code/main.py. The script handles:

• Reading tinyshakespeare.txt if it sits next to the script, otherwise falling back to a 999-character Shakespeare excerpt embedded in main.py. It does not download anything.
• A character tokenizer built from sorted(set(text)): one id per distinct character, no bytes, no BPE.
• Train/val split at 90/10.
• A training loop in default precision (no autocast).
• Sampling after training completes.

Step 1: data

text = open("tinyshakespeare.txt").read()
chars = sorted(set(text))
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 xs: "".join(itos[x] for x in xs)

With the full tinyshakespeare.txt that gives 65 unique characters, a tiny vocabulary with no BPE and no tokenizer drama. With the embedded excerpt you get 46, because the excerpt contains fewer distinct characters. The script prints vocab_size so you can see which case you are in.

Step 2: model

The model is in code/main.py. The block is textbook from Lesson 05: pre-norm, RMSNorm, SwiGLU, causal multi-head attention (the CausalSelfAttention class). The parameter-count preview the script prints first puts the 4-layer, 4-head, d_model=128 config at ~800K parameters. The shipped 3-layer, d_model=64 config is much smaller, around 155K.

Step 3: training loop

Get a random batch of block_size-token windows (64 as shipped). Forward. Shift-by-one cross-entropy. Backward. Clip. AdamW step. Log every 100 steps. Repeat. AdamW with betas (0.9, 0.95) and weight decay 0.1 is the usual transformer recipe: the lower second beta reacts faster to noisy gradients, and the decay keeps weights from growing unchecked.

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)
 
for step in range(max_steps):
    x, y = get_batch("train")
    _, loss = model(x, y)  # shift-by-one cross-entropy happens inside forward
    opt.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()

Step 4: sample

Given a prompt, run a forward pass, keep only the top-k logits, sample from the softmax, append the token, and repeat. main.py prompts with "First", uses temperature=0.9 and top_k=10, and stops after 200 new tokens. There is no top-p sampling. generate() feeds the model only the last block_size tokens each step, so with block_size=64 anything past 64 tokens is generated from a sliding window of recent text, and a 500-token generation exceeds the model's context by a wide margin.

Step 5: read the output

Here is the start of the sample after 2,000 steps on the larger config, from a real run:

Firste atre whit haves as wirse,
Or heres you swer hom st her thare wive ore
Tountimim,
As inor ine th yod talais ometinces me, wee,

That isn't Shakespeare, but the line lengths, capitalisation and word shapes are right, which is about what ~800K parameters and 2,000 steps buy you. At 5,000 steps a few real words appear ("my", "mouse", "bring"). The sampler is seeded, so the same configuration reproduces this text on the same machine, but another machine or PyTorch build will print different text.

The shipped 500-step config on the embedded excerpt prints a step-0 loss near 41, train loss 1.95 and val loss 2.88 at step 500. Val loss bottoms out at step 300 (2.81) and then creeps back up while train loss keeps falling: the excerpt is 999 characters, so the model memorises it. That gap is why Exercise 1 sends you to the full file.

The capstone is a reference architecture. Three extensions take it toward something real:

1. Swap the tokenizer. Use BPE (e.g. tiktoken.get_encoding("cl100k_base")). Vocab size jumps from 65 to ~50,000, and model capacity needs to scale up to compensate.
2. Train on a bigger corpus. Use OpenWebText or fineweb-edu (HuggingFace). 10B tokens on a single A100 takes ~24 hours for a 125M-param GPT.
3. Add RoPE + KV cache + Flash Attention. The exercises below walk you through each.

Do all three and you end up with a 125M-parameter GPT that generates fluent English. It's not a frontier model, but the same code path, just bigger, is what Karpathy, EleutherAI, and the Allen Institute use to train research checkpoints in 2026.

The skill in outputs/skill-transformer-review.md reviews a transformer-from-scratch implementation for correctness across all 13 prior lessons.

1. Easy. Run code/main.py as shipped and keep the printed train and val losses as your baseline. Then scale up: put the full tinyshakespeare.txt next to the script, set n_layers=4, d_model=128, and max_steps=2000. Expect a final-step validation loss around 2.2 at 2,000 steps (2.16 and 2.19 in two of our runs). Then change max_steps to 5,000 and check that val loss keeps improving and drops under 2.0 (1.84 in ours). The embedded excerpt is far too small for a meaningful val loss.
2. Medium. Replace learned positional embeddings with RoPE. Apply the rotation to Q and K inside CausalSelfAttention. Train and verify val loss is at least as low.
3. Medium. Implement a KV cache in the sampling loop. Generate 500 tokens with and without the cache. Note that generate() crops to the last block_size tokens and the learned positional table has only block_size rows, so 500 tokens exceeds the shipped 64-token context. Either raise block_size and retrain, or give both versions the same sliding-window rule, so the timing comparison is fair. Expect the speedup to grow with sequence length (a 5-20x gain on a laptop is a rough estimate, not something main.py measures).
4. Hard. Add a second head to the model that predicts the next-plus-one token (MTP, Multi-Token Prediction from DeepSeek-V3). main.py has no MTP head, so you write it. Train jointly. Does it help?
5. Hard. Replace the single FFN per block with a 4-expert MoE. Use a router with top-2 routing. See how val loss changes at matched active parameters.
6. Medium. Add a linear warmup followed by cosine decay to 10% of peak LR, and wrap the forward pass in bf16 autocast on supported hardware. main.py uses neither. Compare val loss against your baseline.

• The Annotated Transformer (Harvard NLP): the classic annotated implementation.