model_explained.md

Model Folder — Plain Language Explanation

The model/ folder builds a GPT-style decoder-only transformer from scratch, piece by piece. Each file is one component. Here's how they stack:

tokens (integers)
    │
    ▼
┌─────────────┐
│  Embedding  │  config.py defines the shape of everything
└──────┬──────┘
       │
       ▼  ×N layers
┌──────────────────────────────────────┐
│         TransformerBlock             │  block.py
│                                      │
│   ┌──────────┐    ┌──────────────┐   │
│   │ RMSNorm  │    │  RMSNorm     │   │  norm.py
│   └────┬─────┘    └──────┬───────┘   │
│        │                 │           │
│   ┌────▼─────┐    ┌──────▼───────┐   │
│   │Attention │    │  SwiGLU MLP  │   │  attention.py / mlp.py
│   │  + RoPE  │    │              │   │  rope.py
│   └────┬─────┘    └──────┬───────┘   │
│        │  (+residual)    │ (+residual)│
└────────┼─────────────────┼───────────┘
         │                 │
         └────────┬────────┘
                  │
                  ▼
           ┌──────────┐
           │ RMSNorm  │  final norm
           └────┬─────┘
                │
           ┌────▼─────┐
           │  LM Head │  Linear → vocab_size logits
           └──────────┘

---

1. config.py — The Blueprint

What it does: Stores all the numbers that define the model size. Nothing computes anything here — it's just a settings object.

@dataclass
class ModelConfig:
    vocab_size     = 32_000   # how many tokens exist
    context_length = 1024     # max sequence length
    d_model        = 1024     # width of every vector throughout the model
    n_heads        = 16       # how many attention heads
    n_layers       = 9        # how many transformer blocks stacked
    d_ff           = 2816     # width of the MLP hidden layer (auto-computed)

Why these numbers?

• d_model is the "resolution" of the model — bigger = more expressive but more memory
• n_heads splits each attention layer into parallel sub-attentions
• head_dim = d_model / n_heads = 64 — each head sees 64-dim slices
• d_ff for SwiGLU = round_256( 2/3 × 4 × d_model ) — compensates for having 3 matrices instead of 2

Presets defined here:

SLLM_100M: d=768,  h=12, l=12  →  109.5M params
SLLM_150M: d=1024, h=16, l=9   →  148.4M params

---

2. norm.py — RMSNorm

What it does: Normalizes vectors so they don't explode or vanish during training. Used before every attention and MLP layer.

Standard LayerNorm (GPT-2):

1. Compute mean of x
2. Subtract mean  (centering)
3. Divide by std
4. Scale by learned weight
5. Add learned bias

RMSNorm (LLaMA / our model):

1. Compute RMS = sqrt( mean(x²) )   ← no mean subtraction!
2. Divide by RMS
3. Scale by learned weight           ← no bias!

Why simpler is better:

• No mean subtraction → ~40% faster
• No bias → fewer parameters
• Works just as well in practice
• LLaMA, Mistral, Gemma all use it

# What it computes:
output = (x / sqrt(mean(x²) + 1e-6)) * weight
#          ↑ normalize           ↑ rescale with learned gain

The weight starts at all-ones (no change at init) and is learned during training.

---

3. rope.py — Rotary Position Embedding (RoPE)

The problem it solves: Transformers have no built-in sense of position. Without position encoding, "cat sat on mat" and "mat on sat cat" look identical.

How older models solved it (GPT-2): Added a fixed learned vector to each token: token[i] += position_embedding[i] Problem: can't generalize beyond the training length.

What RoPE does instead: Instead of adding position info to token vectors, it rotates the Query and Key vectors in attention by an angle that depends on their position.

Token at position 3 → rotate Q and K by angle θ₃
Token at position 7 → rotate Q and K by angle θ₇

When you compute attention score Q·K, the rotation cancels out in a way that encodes relative distance between tokens, not absolute positions.

Why this is better:

• No extra parameters (pure math, no learned table)
• Works beyond training length (extrapolates)
• Used in LLaMA, Mistral, GPT-4 (likely), Gemma

How the code works:

# Step 1: precompute a table of cos/sin values for every position
cos, sin = precompute_rope_freqs(head_dim=64, max_seq_len=1024)
# cos/sin shape: (1024, 64)

# Step 2: at forward time, rotate Q and K
q_rotated = q * cos + rotate_half(q) * sin
k_rotated = k * cos + rotate_half(k) * sin

# rotate_half(x): splits x in half, negates second half, swaps
# [a, b, c, d] → [-c, -d, a, b]

V (values) are not rotated — only Q and K get position encoding.

---

4. attention.py — Causal Self-Attention

What it does: Lets every token look at all previous tokens and decide which ones are relevant to predict the next token.

The full flow:

Input x: (Batch, Tokens, d_model)
         e.g. (2, 1024, 1024)
    │
    ▼
QKV projection: one big Linear(d_model → 3×d_model)
    │
    ├─── Q: (2, 1024, 1024)  — "what am I looking for?"
    ├─── K: (2, 1024, 1024)  — "what do I contain?"
    └─── V: (2, 1024, 1024)  — "what do I send if attended to?"
    │
    ▼
Reshape to heads: (2, 16_heads, 1024, 64_head_dim)
    │
    ▼
Apply RoPE to Q and K  ← position encoding happens here
    │
    ▼
Scaled Dot-Product Attention:
    scores = Q @ K^T / sqrt(64)    # how much does each token attend to each other
    mask   = causal mask            # can only look LEFT (past), not right (future)
    weights = softmax(scores + mask)
    out    = weights @ V            # weighted sum of values
    │
    ▼
Reshape back: (2, 1024, 1024)
    │
    ▼
Output projection: Linear(d_model → d_model)

Causal mask — this is what makes it a language model (predicts next token):

Position:  0  1  2  3
Token 0:  [✓  ✗  ✗  ✗]   can only see itself
Token 1:  [✓  ✓  ✗  ✗]   can see 0,1
Token 2:  [✓  ✓  ✓  ✗]   can see 0,1,2
Token 3:  [✓  ✓  ✓  ✓]   can see all

Flash Attention: We use F.scaled_dot_product_attention(..., is_causal=True) which is PyTorch 2.0's built-in Flash Attention — it never materializes the full O(T²) attention matrix in memory. Much faster and uses far less VRAM.

---

5. mlp.py — SwiGLU Feed-Forward Network

What it does: After attention (which mixes between tokens), the MLP transforms each token independently — it's where most of the model's "knowledge" is stored.

Standard MLP (GPT-2):

out = W2 @ GELU(W1 @ x)   # 2 matrices

SwiGLU (LLaMA / our model):

gate   = W_gate @ x         # linear
up     = W_up   @ x         # linear
hidden = SiLU(gate) * up    # element-wise gate  ← the key difference
out    = W_down @ hidden     # 3 matrices total

What is SiLU?

SiLU(x) = x × sigmoid(x)

It's a smooth version of ReLU — never exactly zero, has a small negative region.

Why gating matters:

• SiLU(gate) acts as a learned on/off switch for each hidden dimension
• The model learns to activate only the neurons relevant to each input
• Empirically outperforms GELU at the same parameter count
• Used in LLaMA, PaLM, Mistral

The d_ff formula:

d_ff = round_up_256( int(2/3 × 4 × d_model) )

For 150M: round_up_256( int(2/3 × 4 × 1024) ) = round_up_256(2730) = 2816

The 2/3 factor compensates for having 3 matrices instead of 2 — keeps total parameter count equal to a standard 4× FFN.

---

6. block.py — TransformerBlock

What it does: Wraps attention + MLP into one reusable block. The model is just N copies of this block stacked.

def forward(x):
    # Attention sub-layer
    x = x + attention( rmsnorm(x) )   # pre-norm + residual
    
    # MLP sub-layer  
    x = x + mlp( rmsnorm(x) )         # pre-norm + residual
    
    return x

Two key ideas:

1. Pre-norm (normalize BEFORE the sublayer):

Pre-norm (LLaMA):   x → norm → attention → + original x
Post-norm (GPT-2):  x → attention → + original x → norm

Pre-norm is more stable at large scale — gradients flow more cleanly.

2. Residual connections (x + sublayer(x)): The output of each sublayer is added back to the input, not replacing it. This means:

• Gradients can skip directly to earlier layers during backprop
• The model learns corrections to the input, not transformations from scratch
• Allows stacking many layers without vanishing gradients

---

7. model.py — SLLM (The Full Model)

What it does: Assembles everything into the complete language model.

tokens: (B, T)  ← integer IDs like [423, 1829, 55, ...]
    │
    ▼
token_emb: Embedding(32000 → 1024)
    │   converts each integer to a 1024-dim vector
    ▼
blocks[0]: TransformerBlock   ─┐
blocks[1]: TransformerBlock    │  9 blocks for 150M
...                            │
blocks[8]: TransformerBlock   ─┘
    │
    ▼
norm: RMSNorm(1024)   ← final stabilization
    │
    ▼
lm_head: Linear(1024 → 32000)
    │   produces a score for each possible next token
    ▼
logits: (B, T, 32000)   ← unnormalized scores

Weight tying: The token_emb matrix and lm_head matrix share the same weights.

self.lm_head.weight = self.token_emb.weight

• Same matrix used for: embedding lookup (input) AND output projection
• Saves 32M parameters (32000 × 1024)
• Works because: if token X has a similar embedding to the current hidden state, it should also score highly as the next token prediction

Loss computation:

# Cross-entropy: at each position, predict the NEXT token
# Input:  [The, cat, sat, on]   → predicts [cat, sat, on, mat]
# targets = input shifted by 1
loss = cross_entropy(logits.view(-1, 32000), targets.view(-1))

Gradient checkpointing (enable_gradient_checkpointing()): Normally PyTorch saves all intermediate activations during forward pass to use in backprop. For 9 layers with batch_size=2 and seq_len=1024, that's ~1.5GB.

With gradient checkpointing:

• Activations are NOT saved during forward pass
• During backward pass, they are recomputed on-the-fly
• Result: ~40% less VRAM, ~30% slower training
• Essential for fitting 150M on a 4GB GPU

Weight initialization:

# All Linear and Embedding weights: Normal(mean=0, std=0.02)
# Residual projections (o_proj, mlp.down): scaled down by 1/sqrt(2 × n_layers)

The residual scaling prevents the residual stream from growing too large at initialization when many layers add to it.

---

How it all fits together — One forward pass

"The cat sat" → tokenizer → [423, 1829, 55]

token_emb:  [423]→[0.1,-0.3,...] (1024 floats)
            [1829]→[0.8, 0.2,...] (1024 floats)
            [55]  →[-0.1,0.4,...] (1024 floats)

Block 0:
  norm → Q,K,V projections → RoPE rotation → Flash Attention → output proj → + residual
  norm → gate,up projections → SiLU(gate)*up → down proj → + residual

Block 1..8: same

Final norm → LM head → 32000 scores per position

softmax → probabilities → sample next token

Total parameters (150M):

Embedding:   32000 × 1024          =  32.8M
Per block:   attn(4.2M) + mlp(8.6M) + norms(~0M)  =  12.85M
9 blocks:    9 × 12.85M            = 115.6M
Final norm:  1024                  = ~0M
LM head:     TIED to embedding     =   0M  (reuses same weights)
─────────────────────────────────────────
TOTAL:       148.4M params