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.

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