05 Notebook: Real Attention with distilgpt2

April 2026 (296 Words, 2 Minutes)

Goal

Take the attention function you wrote in notebooks 04 and 05 and point it at a real, trained transformer. By the end of this notebook, your pure-NumPy attention block reproduces a layer of distilgpt2 to float precision — the exact computation a real LLM is running, modulo bookkeeping.

Open the Notebook

This notebook needs transformers and torch, which require Rust/C extensions and cannot run in JupyterLite. Run it on Google Colab.

Open in Google Colab — recommended, no setup required.

Loading distilgpt2 takes about 15 seconds; weights are cached after the first run.

What You Will Do

1. Load distilgpt2 — 82 M params, 6 layers, 12 heads, GPT-2 tokenizer
2. Tokenize “The cat sat on the mat”
3. Capture the attention output of layer 3 via a forward hook on model.transformer.h[3].attn
4. Extract layer-3 weights as NumPy: ln_1, c_attn, c_proj
5. The HuggingFace Conv1D gotcha — weight shape (d, 3d), forward is x @ W + b (no transpose)
6. Implement apply_conv1d — wrap the convention as a tiny helper
7. Implement split_qkv — split the (T, 3d) output of c_attn into Q, K, V
8. Build the attention sub-layer in pure NumPy — LayerNorm → c_attn → split → heads → causal attention → merge → c_proj
9. Verify against HuggingFace — np.allclose(custom, hf, atol=1e-5) should pass; expect max diff ≈ 4e-6
10. Implement numpy_attention_block — package the entire pipeline as one function; verify it matches HF the same way
11. Visualise all 12 attention heads — find the previous-token, BOS-attractor, copy heads
12. Bonus — sample next tokens — apply the weight-tied lm_head, temperature, top-k, watch the model’s predictions roll out

Deliverables (AI Diary entry — notebook §10)

• Conv1D vs nn.Linear — what would the weight shape and forward call look like with a regular nn.Linear?
• One head whose pattern caught your eye — describe what it attends to and what real-world structure that might be detecting
• What additional pieces (in order) would you need to implement to reproduce the full layer-3 output, not just the attention sub-layer?
• Weight-tying parameter accounting — for distilgpt2 (V ≈ 50 257, d = 768), how many parameters does tying save? Express as a percentage of the model’s total
• Temperature sweep — re-run the sampling cell with temperature ∈ {0.3, 1.0, 2.0} and describe what changes

Reference

• Lecture 5 — Inside the Transformer
• Chapter 7 — Understanding Attention
• Chapter 12 — The Complete Block, Stacked
• Chapter 14 — Picking the Next Token — for the temperature/top-k sampling story
• distilgpt2 model card on HuggingFace