RMSNorm + RoPE + GQA + SwiGLU + KV cache — every modern LLM component, in one inspectable codebase.
Not a competitive model. A decoder-only transformer small enough to fully understand.
The goal is not to match GPT. It is to build something small enough to trace — attention weights, residual stream norms, loss curves, and KV cache — and understand every design choice by building it yourself.
pip install -r requirements.txt
python src/train.py --data data/train.txt --model_dim 256 --n_layers 6 --n_heads 8| File | What you learn |
|---|---|
src/tokenizer.py |
Character-level tokenizer vs BPE — why subword matters |
src/model.py |
RMSNorm, RoPE, Grouped-Query Attention, SwiGLU activation |
src/train.py |
AdamW + cosine LR schedule + gradient clipping |
src/generate.py |
Temperature sampling + top-p (nucleus) sampling |
src/kv_cache.py |
Why inference is O(T) not O(T²) with caching |
| Choice | This repo | Vanilla transformer |
|---|---|---|
| Norm | RMSNorm (no centering) | LayerNorm |
| Position | RoPE (rotary, no learned params) | Learned absolute |
| Attention | GQA (fewer KV heads) | MHA (n_kv = n_q) |
| Activation | SwiGLU | ReLU / GELU |
Each of these is a decision from a production model (LLaMA, Mistral, Gemma). This lab lets you swap them and observe the difference.
| Parameter | Default | Notes |
|---|---|---|
model_dim |
256 | Hidden dimension |
n_layers |
6 | Transformer blocks |
n_heads |
8 | Attention heads |
n_kv_heads |
2 | KV heads (GQA: 4 query heads per KV head) |
ffn_mult |
4 | MLP expansion ratio |
max_seq_len |
512 | Context length |
vocab_size |
256 | Character-level default |
lr |
3e-4 | Peak learning rate |
- Train with and without RoPE — compare position-sensitivity at long context
- Vary KV heads: MHA (
n_kv=n_heads) vs GQA (n_kv=2) vs MQA (n_kv=1) — measure memory vs quality - Compare character-level vs BPE tokenization on the same corpus
- Insert a fact at different context positions and measure retrieval accuracy (use
long-context-bench) - Visualize attention weights at step 0, 500, 5000 — watch heads specialise
- Loss curve: decreases, plateaus, resumes as model moves from frequency → syntax → semantics
- Attention weights: visualize with
notebooks/attention_viz.ipynb— watch for head collapse early in training - Gradient norm: should stay stable; spikes indicate LR too high or corrupted batches
- Sample quality: inspect every 500 steps; loss alone is not sufficient
- quantization — what to do after training: 4× smaller at 20 dB SNR cost
- microssm — the O(T) alternative to O(T²) attention
- Article: Building a Tiny Transformer from Scratch
MIT