Phase 3: decoder-only transformer (~4.3M params)

- src/model.py: GPT written module by module — RMSNorm, explicit causal
  self-attention (scores/mask/softmax/weighted-sum), GELU MLP, pre-norm
  residual blocks, learned positions, weight-tied head. GPT-2 style init.
- configs/v1.py: frozen ModelConfig dataclass (4L/256d/4h/256ctx/4096vocab)
- tests/test_model.py: exact param count (4,262,144), forward shapes, init
  loss ~= ln(vocab), weight tying, causal-masking check (5 tests, forward-only)
- Verified forward pass runs on MPS.

configs/v1.py

@@ -0,0 +1,25 @@
+"""Experiment v1 configuration.
+
+One experiment = one config file. Iteration 2 is a copy of this file with
+different numbers, not a code change.
+
+v1 is deliberately small (~4.3M parameters): the goal is a complete, legible
+pipeline, not a strong model. It sits below TinyStories' coherence sweet spot
+(~30M params), so some incoherence in the output is expected.
+"""
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ModelConfig:
+    vocab_size: int = 4096       # matches the trained tokenizer
+    context_length: int = 256    # tokens of history the model attends over
+    n_layers: int = 4
+    n_heads: int = 4
+    d_model: int = 256           # residual-stream width
+    d_ff: int = 1024             # MLP hidden width (4 x d_model)
+
+
+# The single source of truth other modules import.
+model = ModelConfig()

src/model.py

@@ -0,0 +1,162 @@
+"""A decoder-only (GPT-style) transformer, written module by module.
+
+Architecture (pre-norm, as in modern GPTs):
+
+    tokens ─► token embedding + learned position embedding
+           ─► N x [ RMSNorm ─► causal self-attention ─► + residual
+                    RMSNorm ─► MLP                    ─► + residual ]
+           ─► final RMSNorm ─► linear head (weight-tied to token embedding)
+
+Attention is implemented explicitly (scores → causal mask → softmax → weighted
+sum of values) rather than via a fused kernel, because seeing the mechanism is
+the point here. `torch.nn.functional.scaled_dot_product_attention` is the
+faster production alternative and an easy iteration-2 swap.
+
+The model takes any config object exposing: vocab_size, context_length,
+n_layers, n_heads, d_model, d_ff.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RMSNorm(nn.Module):
+    """Root-mean-square layer norm: rescale each vector to unit RMS, then a
+    learned per-channel gain. Cheaper than LayerNorm (no mean-centering, no bias)."""
+
+    def __init__(self, dim: int, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return x * rms * self.weight
+
+
+class CausalSelfAttention(nn.Module):
+    """Multi-head self-attention where each position may only attend to itself
+    and earlier positions (the causal constraint that makes this a language model)."""
+
+    def __init__(self, cfg) -> None:
+        super().__init__()
+        assert cfg.d_model % cfg.n_heads == 0
+        self.n_heads = cfg.n_heads
+        self.d_head = cfg.d_model // cfg.n_heads
+        # One matrix produces query, key and value for every head at once.
+        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
+        self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        # Lower-triangular mask, precomputed once. Not a parameter, so register
+        # it as a buffer (moves with .to(device), excluded from grads).
+        causal = torch.tril(torch.ones(cfg.context_length, cfg.context_length))
+        self.register_buffer("mask", causal.view(1, 1, cfg.context_length, cfg.context_length))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+        q, k, v = self.qkv(x).split(C, dim=2)
+        # (B, T, C) -> (B, n_heads, T, d_head): split channels across heads.
+        q = q.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
+        k = k.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
+
+        # Attention scores, scaled so their variance doesn't grow with d_head.
+        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)  # (B, nh, T, T)
+        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
+        att = F.softmax(att, dim=-1)
+        y = att @ v  # (B, nh, T, d_head): each position = weighted sum of past values
+
+        y = y.transpose(1, 2).contiguous().view(B, T, C)  # merge heads back
+        return self.proj(y)
+
+
+class MLP(nn.Module):
+    """Position-wise feed-forward: expand, GELU, project back."""
+
+    def __init__(self, cfg) -> None:
+        super().__init__()
+        self.fc = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
+        self.proj = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.proj(F.gelu(self.fc(x)))
+
+
+class Block(nn.Module):
+    """One transformer block: attention then MLP, each with a pre-norm and a
+    residual connection so gradients flow straight down the stack."""
+
+    def __init__(self, cfg) -> None:
+        super().__init__()
+        self.norm1 = RMSNorm(cfg.d_model)
+        self.attn = CausalSelfAttention(cfg)
+        self.norm2 = RMSNorm(cfg.d_model)
+        self.mlp = MLP(cfg)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class GPT(nn.Module):
+    def __init__(self, cfg) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.pos_emb = nn.Embedding(cfg.context_length, cfg.d_model)
+        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
+        self.norm_f = RMSNorm(cfg.d_model)
+        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        # Weight tying: the output head reuses the token-embedding matrix. Saves
+        # ~1M params and ties "recognising a token" to "predicting it".
+        self.head.weight = self.tok_emb.weight
+
+        self.apply(self._init_weights)
+        # GPT-2 trick: shrink residual-projection init by 1/sqrt(2*n_layers) so
+        # the residual stream doesn't blow up with depth.
+        std = 0.02 / math.sqrt(2 * cfg.n_layers)
+        for name, p in self.named_parameters():
+            if name.endswith("proj.weight"):
+                nn.init.normal_(p, mean=0.0, std=std)
+
+    @staticmethod
+    def _init_weights(module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(
+        self, idx: torch.Tensor, targets: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        B, T = idx.shape
+        assert T <= self.cfg.context_length, "sequence longer than context_length"
+        pos = torch.arange(T, device=idx.device)
+        x = self.tok_emb(idx) + self.pos_emb(pos)  # broadcast pos over batch
+        for block in self.blocks:
+            x = block(x)
+        x = self.norm_f(x)
+        logits = self.head(x)  # (B, T, vocab_size)
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)), targets.reshape(-1)
+            )
+        return logits, loss
+
+    def num_params(self, non_embedding: bool = False) -> int:
+        """Total trainable parameters. Weight tying means the token embedding /
+        head matrix is counted once. `non_embedding` excludes the embedding tables."""
+        n = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n -= self.tok_emb.weight.numel()
+            n -= self.pos_emb.weight.numel()
+        return n

tests/test_model.py

@@ -0,0 +1,92 @@
+"""Model tests: parameter count, output shapes, init loss, and causality.
+
+All forward-pass only (no training). Run: `.venv/bin/python tests/test_model.py`
+"""
+
+import math
+import sys
+from pathlib import Path
+
+import torch
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT / "src"))
+sys.path.insert(0, str(ROOT / "configs"))
+
+from model import GPT  # noqa: E402
+from v1 import ModelConfig  # noqa: E402
+
+
+def test_param_count_matches_hand_estimate():
+    model = GPT(ModelConfig())
+    # tok_emb 4096*256 + pos_emb 256*256
+    # + 4 blocks * (qkv 256*768 + attn_proj 256*256 + 2 norms*256
+    #               + mlp_fc 256*1024 + mlp_proj 1024*256)
+    # + final norm 256   (head is weight-tied, so no extra params)
+    expected = (
+        4096 * 256 + 256 * 256
+        + 4 * (256 * 768 + 256 * 256 + 2 * 256 + 256 * 1024 + 1024 * 256)
+        + 256
+    )
+    assert expected == 4_262_144
+    assert model.num_params() == expected
+
+
+def test_weight_tying():
+    model = GPT(ModelConfig())
+    # Head and token embedding must be the very same parameter tensor.
+    assert model.head.weight is model.tok_emb.weight
+
+
+def test_forward_shapes():
+    cfg = ModelConfig()
+    model = GPT(cfg)
+    idx = torch.randint(0, cfg.vocab_size, (2, 16))
+    logits, loss = model(idx)
+    assert logits.shape == (2, 16, cfg.vocab_size)
+    assert loss is None
+    _, loss = model(idx, targets=idx)
+    assert loss.ndim == 0  # scalar
+
+
+def test_init_loss_near_uniform():
+    # A well-initialised model predicts ~uniformly at first, so the
+    # cross-entropy loss should be close to ln(vocab_size). Targets are drawn
+    # independently of the input: with weight tying, using the input as its own
+    # target would leak (each position's residual already holds its embedding).
+    cfg = ModelConfig()
+    torch.manual_seed(0)
+    model = GPT(cfg)
+    idx = torch.randint(0, cfg.vocab_size, (8, 64))
+    targets = torch.randint(0, cfg.vocab_size, (8, 64))
+    _, loss = model(idx, targets=targets)
+    assert abs(loss.item() - math.log(cfg.vocab_size)) < 0.5, loss.item()
+
+
+def test_causality():
+    # Changing a token at position t must not affect logits at positions < t.
+    cfg = ModelConfig()
+    torch.manual_seed(0)
+    model = GPT(cfg)
+    model.eval()
+    idx = torch.randint(0, cfg.vocab_size, (1, 32))
+    with torch.no_grad():
+        base, _ = model(idx)
+        changed = idx.clone()
+        changed[0, 20] = (changed[0, 20] + 1) % cfg.vocab_size
+        after, _ = model(changed)
+    # Positions before the edit are identical; the edited position differs.
+    assert torch.allclose(base[0, :20], after[0, :20], atol=1e-5)
+    assert not torch.allclose(base[0, 20], after[0, 20])
+
+
+def main() -> None:
+    tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
+    for t in tests:
+        t()
+        print(f"ok  {t.__name__}")
+    print(f"\n{len(tests)} passed")
+
+
+if __name__ == "__main__":
+    main()