From e708a5b8869b76d347c024ed3c7d0aa9674410de Mon Sep 17 00:00:00 2001 From: rzen Date: Sun, 12 Jul 2026 10:48:55 -0400 Subject: [PATCH] Phase 3: decoder-only transformer (~4.3M params) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 | 25 +++++++ src/model.py | 162 ++++++++++++++++++++++++++++++++++++++++++++ tests/test_model.py | 92 +++++++++++++++++++++++++ 3 files changed, 279 insertions(+) create mode 100644 configs/v1.py create mode 100644 src/model.py create mode 100644 tests/test_model.py diff --git a/configs/v1.py b/configs/v1.py new file mode 100644 index 0000000..47ec585 --- /dev/null +++ b/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() diff --git a/src/model.py b/src/model.py new file mode 100644 index 0000000..b176963 --- /dev/null +++ b/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 diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..d71fae5 --- /dev/null +++ b/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()