Tiny Transformer

torch>=2.14.0 is the only runtime dependency: everything in model/ is pure PyTorch (nn.Linear, nn.Embedding, softmax, matmul). transformers==5.17.0 is pinned in dev deps and used solely by reference.py to build the LlamaForCausalLM answer key — pinned exact because even a minor HF RoPE or norm tweak would break the 1e-4 numerical match. REFERENCE_TRANSFORMERS in code mirrors this pin. pytest is the dev-only runner, and setting pythonpath to the project root is why from model.config import ... works when you run uv run pytest inside llm-lab/. No Kafka, no server, no GPU requirement — CPU is sufficient, the whole model is 26k params.

Taped to the fridge. Nothing computes here. It says: we speak 128 tokens, every token is 32 numbers, stack 2 layers, each layer has 4 question-heads sharing 2 key/value-heads. Every nn.Linear size below is arithmetic on these numbers — change d_model to 64 and the whole building reshapes. frozen=True means a struct you cannot accidentally mutate mid-run: config is truth, not state. d_model=32 is the width of the residual stream, the highway tokens ride. n_heads=4, n_kv_heads=2, head_dim=8 is the attention geometry: 4 query vectors of 8 numbers is exactly 32 (d_model), with only 2 key/value vectors — that is Grouped-Query Attention. d_ff=64 is the MLP workroom width, tokens detour 32→64→32. rope_theta=10_000.0 is the RoPE base frequency, bigger means slower-spinning hands. tie_embeddings=False means the input lookup table and the output vote matrix are separate, costing an extra 4k params (sharing would save it).

The guardrails: 4 % 2 == 0 means 2 query heads can cleanly share each KV head (n_rep = 2). RoPE rotates pairs of numbers like x/y coordinates — an odd dim leaves one number partnerless, so evenness is rejected in both config.py and rope.py.

RMSNorm runs 5 times per forward (twice per layer + once at the end). Think of it as a volume knob per token: it forces each token's 32 numbers to unit loudness, then applies 32 learned fader sliders (self.weight, starting all-ones). What it does NOT do: never looks across the sequence — token 3's loudness knows nothing about token 4 — and it does not subtract the mean (that is LayerNorm). Llama showed mean-centering is not worth the extra reduction.

Forward pass, step by step, with x of shape [2, 5, 32]: 1. Upcast: x.float() — math in fp32 even if the stream is fp16, precision where it matters. 2. Loudness: pow(2).mean(dim=-1, keepdim=True) gives [2, 5, 1], the mean-square per token. keepdim=True keeps [2,5,1] so it broadcasts against [2,5,32]. 3. Normalize: x * rsqrt(var + eps), that is x / sqrt(mean(x²)), with eps=1e-5 avoiding divide-by-zero on silent tokens. 4. Rescale and restore: weight * x32, cast back to the input dtype. Formula: $y = w \odot x/\mathrm{RMS}(x)$.

Why no bias, no mean? One reduction instead of two, fewer params, identical quality at this scale. Why the dtype dance? mean(pow(2)) overflows easily in fp16 — computing in fp32 then returning keeps the stream fast but the math safe. Why eps inside the sqrt? rsqrt(0) is inf, and eps guarantees a floor.

Attention is order-blind — shuffle inputs, get shuffled outputs. RoPE fixes it by rotating queries and keys by an angle proportional to position. Each pair of numbers is a clock hand: position 0 leaves hands alone, position 5 twists them. When query at m dots key at n, the relative twist m−n falls out of the dot product automatically. No position table to learn, and it generalizes past max_seq_len=64 better than absolute embeddings. Applied to q and k, never v — scores need position, while the content being averaged must not be scrambled.

Building blocks: with head_dim=8, arange(0,8,2) = [0,2,4,6] gives 4 pairs, and inv_freq = 1 / theta^(i/8). Pair 0 has 10_000^0 = 1 and spins fast at 1 rad per step; pair 3 has 10_000^0.75, roughly 1000, and spins slow at 0.001 rad per step — fine versus coarse position. freqs = position × inv_freq, so position 5 on pair 0 is 5.0 rad and on pair 3 is 0.005 rad. rotate_half([a|b]) = [-b|a] is a 90° rotation in halves-pairing: the first half pairs with the second half, (0,4),(1,5),(2,6),(3,7) for dim 8. This is Llama/GPT-NeoX style, *not* neighbor pairing (0,1),(2,3) — matching the HF pairing is why the error is 0.0. Finally x·cos + rotate(x)·sin is a true 2D rotation, vectorized.

Forward pass with positions [0,1,2,3,4] and head_dim=8: rotary_frequencies returns cos_half, sin_half each [5, 4]; rope_cos_sin duplicates them to [5, 8] each; in attention they are unsqueezed to [1,1,5,8] and broadcast against [B,H,5,8] — same twist for every batch and head, different twist per position. Position 0 is identity (cos=1, sin=0), so apply_rope(x) = x there. Why even-only? You cannot rotate a loner, and the config check plus this check double-guard it. Why theta=10_000? The Llama default — the slow hands do not wrap within 64 tokens. Why halves-pairing, not interleaved? Because that is what LlamaForCausalLM does; pick the other pairing and the reference check fails even with identical weights.

The only place tokens listen to each other. Meeting analogy: each token asks a question (query), everyone holds a name tag (key) and a memo (value). Query×key similarity decides listening weights; values are what is actually averaged. GQA means 4 askers share 2 tag printers — halving the KV cache at inference versus full MHA. It does NOT see the future (causal mask) and does NOT think per-token (that is the MLP).

Forward pass with x = [2,5,32]: 1. Remix. Each Linear is input @ W, no bias. q becomes [2,4,5,8], k,v become [2,2,5,8] — the view(...).transpose(1,2) splits the flat projection into heads. 2. Position. RoPE twists q,k, taking row 0 if a [B,T] batch of position ids is passed (uniform positions across the batch are assumed). 3. Expand groups. repeat_kv(..., n_rep=2) turns [2,2,5,8] into [2,4,5,8] via a new axis, expand, and reshape — each KV head serves 2 query heads, and n_rep==1 short-circuits to pure MHA. 4. Score. q @ kᵀ / sqrt(8) gives [2,4,5,5], where [b,h,i,j] reads token i listens to token j this much; the scale keeps dot variance near 1 so softmax does not saturate. 5. Mask. The tril() lower triangle is past-allowed, future entries go to finfo.min (\−3e38) and get softmax weight \0 — token 0 hears only itself. 6. Mix and merge. softmax(dim=-1, dtype=fp32).to(q.dtype), then @ v back to [2,4,5,8], transpose and merge to [2,5,32], and o_proj remixes to model width.

Why fp32 softmax then cast? exp() overflows in fp16, the rest stays cheap. Why bias=False on all four projections? Norms already carry scale, biases add params and complicate quantization for zero gain — and it matches attention_bias=False in the HF config. Why materialize the [T,T] mask? Fine at T=64; at T=8k you would switch to scaled_dot_product_attention(is_causal=True) or FlashAttention instead of allocating the triangle.

Attention mixes across tokens; this thinks within each token, blind to neighbors. With x of shape [2,5,32], gate and up each produce [2,5,64], then silu(gate) * up stays [2,5,64] before down squeezes back to [2,5,32]. silu(z) = z × sigmoid(z): smooth, slightly non-monotonic near zero, unbounded above — gentler than the hard kink of ReLU. The * is the GLU valve: the gate decides how much of each up feature survives. Stacked linears without this would collapse (W2(W1x) = (W2W1)x — no new power), so the SiLU plus multiply is what makes the 64-wide middle add capacity. Three matrices of 32×64 is 6,144 params per layer, versus 4,096 for a vanilla 2-matrix MLP — the 1.5× cost is the Llama-family trade that won, and it matches HF hidden_act=silu with mlp_bias=False exactly — one reason the 0.0 error is attainable.

Pre-norm decoder: embed, then twice (norm→attend→add, norm→think→add), then norm and un-embed. The x + sublayer(norm(x)) residual highway lets signal and gradients skip any block. Post-norm (Norm(x + Attn(x)), the original 2017 style) forces everything through a norm; pre-norm keeps the highway clean and trains stably — every modern decoder does it this way. Forward pass with input_ids = [2,5] ints: position_ids = arange(5) when the caller passes none; embed_tokens looks ints up in the [128,32] table to give [2,5,32] — pure dictionary, no math; two layers of the residual pair; final norm, then lm_head 32→128 gives logits [2,5,128], where logits[b,t,v] scores position t having vocab v as its next token.

The fixed counter: parameter_breakdown() groups by startswith('embed_tokens'), '.self_attn.' in name, '.mlp.' in name, startswith('lm_head'), else norm. The old embeded_tokens typo had hidden the 4,096 embedding params inside norm; with the key fixed, the measured result (counted from the module, not a formula) is: embed 4,096 (128×32); attention 6,144 — that is 2 × (32×32 + 32×16 + 32×16 + 32×32); mlp 12,288 — 2 × 3 × 32×64; norm 160 — 5 × 32; head 4,096 — 32×128; total 26,784. Smaller than a JPEG. tie_embeddings=True would share the embed and head storage (same Parameter object, counted once).

No new math — it is a translator. tiny_config() returns your defaults; hf_llama_config() maps each field to HF names (d_model→hidden_size, d_ff→intermediate_size, and so on). build_hf_reference(seed=0) constructs the HF model in eval mode, and copy_llama_weights() loads every tensor slot-for-slot: embed, both norms per layer, gate/up/down, q/k/v/o, final norm, head. Proof (recorded in reports/tiny-transformer.md, torch 2.14 CPU): same weights plus the same 2×7 probe ids gives hidden error 0.0 and logits error 0.0 at tolerance 1e-4. eager plus fp32 is deliberate — Flash and SDPA kernels reorder floating-point summation and would muddy a bit-exact check at this scale. __init__.py stays the front door, re-exporting RMSNorm and apply_rope.