rzen
e708a5b886
- 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.
93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
"""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()
|