Build a Mini LLM from Scratch in NumPy: RoPE, GQA, SwiGLU Visual Guide

I already knew how to build GPT-2 in PyTorch. 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. If you want the arithmetic of one training step past attention, keep 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

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

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

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

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 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

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

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, 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 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

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

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. I also wrote a dedicated note on token embeddings 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

If you’ve built autograd from scratch, the Tensor layer will feel familiar — matmul, reshape, gather, broadcast-aware backward. I lean on numerical gradient checking before I trust any new op. Optimizer is AdamW with grad clipping; LR is warmup then cosine decay.

How was it trained?

Pretrain on everyday dialogue (DailyDialog), formatted with chat specials <bos>, <|user|>, <|assistant|>, <eos>:

~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, 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 <eos>, 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:

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:

How GPT actually works — big picture
Attention: Q, K, V — intuition
Token embeddings · positional encoding · BPE
Autograd · grad checking
Cross-entropy · AdamW
Build GPT-2 in PyTorch — classic stack
GPT math beyond attention — one full training step
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.