--- layout: post title: "Build a Mini LLM from Scratch in NumPy: RoPE, GQA, SwiGLU Visual Guide" date: 2026-07-16 17:00:00 +0530 tags: ["LLM", "NumPy", "Transformers", "From Scratch", "RoPE", "GQA", "SwiGLU", "RMSNorm", "Chat"] category: "ai" description: "How to build a mini LLM from scratch in NumPy — RoPE, GQA, QK-Norm, SwiGLU, tied embeddings, and a 3.87M chat companion with full architecture visuals." keywords: "build mini LLM from scratch numpy, numpy language model from scratch, RoPE GQA SwiGLU explained, grouped query attention, RMSNorm Pre-LN, weight tied embedding, from scratch transformer numpy, small language model architecture" image: /assets/images/og-build-mini-llm-numpy.png order: 15 --- I already knew how to [build GPT-2 in PyTorch](/blog/build-gpt2-from-scratch/). That was useful. It was also a cheat: `torch.nn` hid the autograd, the backend, and half the decisions modern small models actually make. So I rebuilt the whole stack in **pure NumPy** — custom reverse-mode autograd, byte-BPE, training loop, KV-cache, sampler — as a tiny **3.87M-parameter** decoder that only does short English small-talk. No code generation. No tools. No retrieval. Deliberately narrow, so every formula had somewhere to land. This post is the visual architecture walkthrough of that model. If you want the older GPT-2 recipe (absolute position embeddings, MHA, GELU MLP, LayerNorm), start with [that walkthrough](/blog/build-gpt2-from-scratch/). If you want the arithmetic of one training step past attention, keep [GPT math beyond attention](/blog/gpt-math-beyond-attention/) open beside this. Here I'm documenting what I actually shipped: **RoPE + RMSNorm (Pre-LN) + GQA + QK-Norm + SwiGLU + weight-tied head**. ![Model A overview — from-scratch mini chat companion architecture tower with RoPE, RMSNorm, GQA, SwiGLU](/assets/images/model-a/01-model-a-overview.png) --- ## What did I actually build? A decoder-only language model that maps token IDs → next-token logits. Same objective as GPT. Different modern-tiny choices. | Piece | Choice | Why I picked it | |---|---|---| | Positions | **RoPE** (rotary) | Relative positions without a learned `wpe` table | | Norm | **RMSNorm**, Pre-LN | Cheaper than LayerNorm; Pre-LN trains sanely | | Attention | **GQA** (8 query / 2 KV heads) + **QK-Norm** | Fewer KV params; stabilize scores before RoPE | | MLP | **SwiGLU** (hidden 704) | Gated MLP used by modern LLaMA-style stacks | | Head | **Weight-tied** to token embedding | Cuts ~1M params on a 4k vocab | | Framework | Pure NumPy (+ optional CuPy swap) | No PyTorch/JAX; every op owns its backward | | Scope | Companion chat only | Keeps the eval honest — refusal of coding prompts is a feature | ```mermaid flowchart TB IDS["Token IDs (B, T)"] --> EMB["TiedEmbedding gather\nW ∈ R^{4096×256}"] EMB --> B1["Block 1: RMSNorm → Attn → + → RMSNorm → SwiGLU → +"] B1 --> B2["Block 2"] B2 --> B3["Block 3"] B3 --> B4["Block 4"] B4 --> FN["Final RMSNorm"] FN --> HEAD["logits = H Wᵀ"] HEAD --> OUT["Logits (B, T, 4096)"] ``` --- ## What are the exact hyperparameters? I kept the config small enough to train on CPU and still feel like a real transformer: | Hyperparameter | Value | |---|---:| | `vocab_size` | 4096 | | `d_model` | 256 | | `n_layers` | 4 | | `n_heads` | 8 | | `n_kv_heads` | 2 | | `head_dim` | 32 | | `swiglu_hidden` | 704 | | `seq_len` | 256 | | Total parameters | **3,869,184** | Sanity check on head math: $8 \times 32 = 256 = d_{\text{model}}$, and GQA repeats each KV head $8/2 = 4$ times so query heads still match. ![Model A config at a glance — vocab 4096, d_model 256, 4 layers, GQA 8Q/2KV, SwiGLU 704](/assets/images/model-a/02-model-a-config.png) ### How do the 3.87M parameters break down? Tied embedding (shared with the LM head): $$ 4096 \times 256 = 1{,}048{,}576 $$ One attention block (GQA + QK-Norm weights): $$ \begin{align*} W_Q &: 256 \times 256 = 65{,}536 \\ W_K &: 256 \times 64 = 16{,}384 \\ W_V &: 256 \times 64 = 16{,}384 \\ W_O &: 256 \times 256 = 65{,}536 \\ \text{Q/K norms} &: 32 + 32 = 64 \\ \text{subtotal} &= 163{,}904 \end{align*} $$ One SwiGLU: $$ 3 \times (256 \times 704) = 540{,}672 $$ Two RMSNorms per block ($256 + 256$), times 4 blocks, plus final RMSNorm: $$ \begin{align*} 4 \times (163{,}904 + 540{,}672 + 512) &= 2{,}820{,}352 \\ +\; 1{,}048{,}576 + 256 &= 3{,}869{,}184 \end{align*} $$ That number is not marketing copy — I print `model.n_params()` and it returns exactly `3869184`. --- ## How does one forward pass look end to end? Text → byte-BPE → IDs → embed → 4 Pre-LN blocks → final RMSNorm → tied projection → logits → sample. ![One forward pass pipeline from chat text through tokenizer, tied embedding, four blocks, and logits](/assets/images/model-a/03-forward-pass.png) Shapes, for a batch of size $B$ and sequence length $T$: | Stage | Shape | |---|---| | Token IDs | $(B, T)$ | | Embeddings | $(B, T, 256)$ | | After each block | $(B, T, 256)$ | | Logits | $(B, T, 4096)$ | Next-token training uses the usual shift: predict token $t+1$ from positions $0..t$, with [cross-entropy](/blog/cross-entropy-loss-explained/) on the logits. Padding / user-turn tokens in SFT use an ignore index so only assistant tokens pull the gradient. --- ## What happens inside one transformer block? Pre-LN residual. Normalize *before* the sublayer, add the residual *after*. $$ \begin{align*} x &\leftarrow x + \mathrm{Attn}(\mathrm{RMSNorm}(x)) \\ x &\leftarrow x + \mathrm{SwiGLU}(\mathrm{RMSNorm}(x)) \end{align*} $$ ![Transformer block Pre-LN residual — RMSNorm then attention, then RMSNorm then SwiGLU](/assets/images/model-a/04-transformer-block.png) I used to treat Post-LN vs Pre-LN as trivia. Then I tried training a tiny stack without Pre-LN and watched early layers thrash. Pre-LN is the boring choice that lets a 4-layer toy actually converge. RMSNorm itself is LayerNorm without the mean subtraction: $$ \mathrm{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{d}\sum_i x_i^2 + \varepsilon}} \odot \gamma $$ No bias. One gain vector $\gamma \in \mathbb{R}^{d}$. That's it. --- ## How does attention work with GQA, QK-Norm, and RoPE? This is the part that differs hardest from the GPT-2 post. ![Attention path — GQA 8 query heads vs 2 KV heads, QK-Norm, RoPE, causal mask, softmax](/assets/images/model-a/05-attention-gqa.png) ### What is grouped-query attention? Full MHA would give every head its own $K$ and $V$. GQA keeps **8 query heads** but only **2 key/value heads**, then repeats KV: $$ n_{\text{rep}} = \frac{n_{\text{heads}}}{n_{\text{kv heads}}} = \frac{8}{2} = 4 $$ So each KV head is shared by four query heads. You pay less for the KV cache at inference and fewer parameters in $W_K$, $W_V$. For a chat toy with `seq_len=256`, that matters less than it does at 32k context — but the code path is the one I wanted to own. ### Why QK-Norm before RoPE? I RMSNorm $Q$ and $K$ **per head** (dim 32) before rotating. Attention scores then look like: $$ \mathrm{scores} = \frac{Q K^\top}{\sqrt{d_h}} + \mathrm{causal\_mask}, \quad d_h = 32 $$ The scale is $1/\sqrt{32} \approx 0.1768$. Causal mask writes $-10^9$ into future positions so softmax zeros them. Same idea as [causal self-attention](/blog/attention-in-transformers-explained/), just with modern extras bolted on. ### How does RoPE encode position? Instead of adding a learned position vector (GPT-2's `wpe`), RoPE rotates pairs of dimensions in $Q$ and $K$ by an angle that depends on token index $m$: $$ \theta_i = 10000^{-2i / d_h}, \qquad (\tilde{q}_m)_{2i:2i+1} = R(m\theta_i)\, q_{2i:2i+1} $$ Relative position falls out of the dot product $q_m^\top k_n$ depending on $m-n$. No separate position table. If absolute sine/cosine encodings are what you grew up with, [positional encoding explained](/blog/positional-encoding-explained/) is the classic version; RoPE is the one this model uses. Worked fragment for $d_h=4$, position $m=1$, base $10000$: $$ \theta_0 = 10000^{0} = 1, \quad \theta_1 = 10000^{-1/2} = 0.01 $$ Cosine/sine tables store $\cos(m\theta)$ and $\sin(m\theta)$ for every position up to `seq_len=256`. The forward applies the rotate-and-mix form $x\cos + \mathrm{rot}(x)\sin$ on the last dimension. --- ## What is SwiGLU and why 704? Classic GPT MLP is $d \to 4d \to d$ with GELU. SwiGLU is gated: $$ \mathrm{SwiGLU}(x) = \bigl(\mathrm{SiLU}(x W_g) \odot (x W_u)\bigr) W_d $$ ![SwiGLU MLP — gate path SiLU(xWg), value path xWu, multiply, project with Wd](/assets/images/model-a/06-swiglu-mlp.png) SiLU is $x \cdot \sigma(x)$. The $2/3$ rule shrinks the hidden width so the three matrices roughly match the parameter count of a two-matrix $4d$ MLP: $$ \frac{2}{3} \times 4 \times 256 \approx 682.67 \;\rightarrow\; \text{I rounded to } 704 $$ Three matrices of shape $256\times704$ / $704\times256$ dominate the block. On this model, MLP params per layer ($540{,}672$) are about $3.3\times$ the attention params ($163{,}904$). Same story as every modern decoder: the feed-forward is where most of the weight lives. --- ## Why tie the embedding to the LM head? One matrix $W \in \mathbb{R}^{4096 \times 256}$ does both jobs: $$ \mathrm{embed}(\mathrm{ids}) = W[\mathrm{ids}], \qquad \mathrm{logits} = H W^\top $$ ![Tied embedding — same weight matrix for token lookup and vocabulary projection](/assets/images/model-a/08-tied-embedding.png) Without tying I'd pay another $1{,}048{,}576$ parameters for a separate `lm_head`. With a 4k vocab that's ~27% of the model. GPT-2 does the same trick; see the weight-tying section in [Build GPT-2 from scratch](/blog/build-gpt2-from-scratch/). I also wrote a dedicated note on [token embeddings](/blog/token-embeddings-explained/) if the gather semantics still feel abstract. --- ## What does the full system stack look like? The model file is only one layer of the cake. Below it: NumPy/CuPy backend, autograd `Tensor`, functional ops. Above it: loss, AdamW, schedules, train/SFT scripts, KV-cache generate, chat runtime. ![Model A full system stack from NumPy backend through autograd, model, train, KV-cache, and chat](/assets/images/model-a/07-system-stack.png) If you've built [autograd from scratch](/blog/build-autograd-from-scratch/), the Tensor layer will feel familiar — matmul, reshape, gather, broadcast-aware backward. I lean on [numerical gradient checking](/blog/numerical-gradient-checking-explained/) before I trust any new op. Optimizer is [AdamW](/blog/adam-optimizer-explained/) with grad clipping; LR is warmup then cosine decay. ### How was it trained? **Pretrain** on everyday dialogue (DailyDialog), formatted with chat specials ``, `<|user|>`, `<|assistant|>`, ``: - ~1.62M tokens · 2000 steps · batch 32 × block 256 - AdamW · cosine LR $3\times10^{-4} \rightarrow 3\times10^{-5}$ - Val perplexity $\approx 25$ **SFT** on warm supportive dialogues (EmpatheticDialogues) with **loss masking** — only assistant turns get supervised: - 19.4k dialogues · 800 steps · fresh AdamW · cosine $1\times10^{-4} \rightarrow 1\times10^{-5}$ - Assistant-masked val perplexity $\approx 33$ · next-token accuracy $0.325$ - Coding-prompt refusal check: $3/3$ answered as chit-chat (no code emitted) Tokenizer is byte-level BPE, vocab 4096, specials reserved before the byte base — same family of ideas as [BPE from scratch](/blog/byte-pair-encoding-from-scratch/), just smaller. ### How does inference work? Sampling supports greedy, temperature, top-$k$, top-$p$. KV-cache stores past $K,V$ so each new token is an incremental forward; I test that cached logits match a full recompute. Chat runtime renders turns, generates until ``, optional system persona string. Hosting (`/chat` API, Docker) is still on the checklist. The learning artifact is the trainable NumPy model + SFT checkpoint, not a production companion. At 3.87M params it will mismatch emotion, forget turns, and invent nonsense. That's expected — the point was owning every op, not beating a 7B chat model. --- ## Can I run it? Conceptually: ```python from mla.checkpoint import load_checkpoint from mla.tokenizer import Tokenizer from mla.chat import ChatSession tok = Tokenizer.load("data/tokenizer/tokenizer.json") model, _, _ = load_checkpoint("checkpoints/sft_final.npz") chat = ChatSession(model, tok, temperature=0.8, top_k=40, top_p=0.9) print(chat.reply("I had a rough day today.")) ``` Backend switch for GPU is an env flag (`ZYN_BACKEND=cuda` → CuPy). Gradcheck and overfit runs stay on float64; GPU train/infer use float32. --- ## Where does this sit in the series? Read order that worked for me: 1. [How GPT actually works](/blog/what-is-a-gpt-visual-intro/) — big picture 2. [Attention: Q, K, V](/blog/attention-in-transformers-explained/) — intuition 3. [Token embeddings](/blog/token-embeddings-explained/) · [positional encoding](/blog/positional-encoding-explained/) · [BPE](/blog/byte-pair-encoding-from-scratch/) 4. [Autograd](/blog/build-autograd-from-scratch/) · [grad checking](/blog/numerical-gradient-checking-explained/) 5. [Cross-entropy](/blog/cross-entropy-loss-explained/) · [AdamW](/blog/adam-optimizer-explained/) 6. [Build GPT-2 in PyTorch](/blog/build-gpt2-from-scratch/) — classic stack 7. [GPT math beyond attention](/blog/gpt-math-beyond-attention/) — one full training step 8. **This post** — modern-tiny NumPy companion (RoPE / GQA / SwiGLU) I used to think "from scratch" meant reimplementing GPT-2's 2019 choices in a new file. Rebuilding with RoPE, GQA, and SwiGLU forced me to confront what actually changed in the last few years — and how little of that change shows up if you only ever call `nn.TransformerDecoder`.