Project 04 — Tiny Transformer Engine

Load real GPT-2 weights into your own code and generate text — deliberately the slow way.

1. What you build#

A decoder-only transformer forward pass written on top of your Project 01 operators, loaded with the public GPT-2 small (124M) weights, plus a greedy generation loop. No KV cache yet — each new token re-runs the whole sequence. That inefficiency is the point; Project 05 removes it and you measure the difference.

Diagram — What you are building#

flowchart LR
  T["Token ids"] --> E["Token + position embeddings"]
  E --> BLK
  subgraph BLK["x 12 blocks"]
    direction TB
    L1["LayerNorm"] --> AT["Causal self-attention"] --> R1(("+"))
    R1 --> L2["LayerNorm"] --> ML["MLP"] --> R2(("+"))
  end
  BLK --> LF["Final LayerNorm"] --> HD["Output head<br/>d_model x vocab"]
  HD --> AM["argmax"]
  AM -.->|"append and re-run EVERYTHING"| T

  class AT,ML,HD compute
  class E,L1,L2,LF neutral
  class AM queue
  class T neutral

2. Why it matters#

This is the moment “the transformer” stops being a diagram. You will know exactly which matmuls run, in what order, with what shapes — which is the prerequisite for understanding every optimization in Sections V and VII.

3. Read first#

• III.08 — Attention from first principles
• III.09 — The transformer
• IV.01 — Tracing one request
• V.02 — Tokenization
• V.04 — Autoregressive generation

4. Spec#

Model:     GPT-2 small — 12 layers, 12 heads, d_model 768, vocab 50257, ctx 1024
Weights:   from Hugging Face `gpt2` (convert once to .npz)
Tokenizer: use `tiktoken` or HF tokenizer — do not write BPE yourself here
API:       logits = model(ids)          # ids: (B, T) -> (B, T, vocab)
           ids    = generate(prompt_ids, max_new_tokens)   # greedy

GPT-2 specifics that bite: learned absolute position embeddings, pre-LN blocks, a final LayerNorm, weight tying between wte and the output head, and HF’s Conv1D layers whose weights are stored (in, out) — the transpose of nn.Linear.

5. Milestones#

1. Convert weights. Dump the HF state_dict to .npz. Print every name and shape; map each to a line in your forward pass.
2. One block. Implement causal self-attention + MLP for a single layer. Compare its output to HF’s output_hidden_states at layer 1.
3. Full forward. Match final logits to HF within 1e-3 (float32 accumulation over 12 layers) and identical argmax at every position.
4. Generate. Greedy decode 50 tokens. Text must equal HF’s generate(do_sample=False).
5. Instrument. Time each generated token. Plot per-token latency against position.

6. Starter skeleton#

def attention(x, p, n_head):
    B, T, C = x.shape
    qkv = x @ p["c_attn.w"] + p["c_attn.b"]           # (B, T, 3C) — Conv1D: (in, out)
    q, k, v = np.split(qkv, 3, axis=-1)
    split = lambda t: t.reshape(B, T, n_head, C // n_head).transpose(0, 2, 1, 3)
    q, k, v = split(q), split(k), split(v)             # (B, H, T, D)
    att = q @ k.transpose(0, 1, 3, 2) / np.sqrt(C // n_head)
    att = att + np.triu(np.full((T, T), -1e10, dtype=np.float32), k=1)   # causal mask
    y = softmax(att) @ v                               # (B, H, T, D)
    y = y.transpose(0, 2, 1, 3).reshape(B, T, C)
    return y @ p["c_proj.w"] + p["c_proj.b"]

def block(x, p, n_head):
    x = x + attention(layernorm(x, *p["ln_1"]), p["attn"], n_head)   # pre-LN
    x = x + mlp(layernorm(x, *p["ln_2"]), p["mlp"])
    return x

def generate(ids, n_new):
    for _ in range(n_new):
        logits = forward(ids[:, -1024:])       # recomputes EVERYTHING each step
        nxt = logits[:, -1].argmax(-1, keepdims=True)
        ids = np.concatenate([ids, nxt], axis=1)
    return ids

7. What to measure#

8. Done when#

• Logits match Hugging Face; greedy text is identical.
• You can draw the shape of every tensor in one block from memory.
• You have a plot showing per-token latency growing with position, and you can state what work is being repeated.
• You computed the output head for all positions and then realized you only need the last one — and measured what that saves.

9. Common pitfalls#

The Conv1D transpose. If your logits are garbage, this is the first suspect.

Mask of -inf in float32 producing NaN via inf - inf. Use a large negative finite value or apply the mask carefully (IV.12).

Forgetting the final ln_f.

Position embeddings off by one when you later truncate context.

Comparing against HF with dropout active. .eval().

10. Stretch goals#

• Add temperature, top-k, top-p sampling with a seed (V.13) and match HF for a fixed seed.
• Port the forward pass to PyTorch tensors on GPU while keeping your own code structure.
• Swap in a RoPE + RMSNorm + GQA model (e.g. SmolLM2-360M). Note every line that changes.
• Batch two prompts of different lengths with padding and an attention mask.

11. Interview questions this project answers#

1. Walk through the tensor shapes in one attention layer.
2. Why is naive generation quadratic (or worse) in sequence length?
3. What is the causal mask and where exactly is it applied?
4. Which single matmul in GPT-2 small has the most parameters?

12. Next#

Project 05 — KV cache