rzen
9fd56b6063
Controlled study of how the model degrades. Four variants, each = v1 with a single field replaced (dataclasses.replace), so differences are attributable: - small: capacity down (2L/64d, ~0.38M params) - short: training time down (max_iters 20k->2k) - ctx: context window down (256->64) - data: data quantity down (max_train_tokens->1M; new TrainConfig knob + train.py slice) - scripts/compare.py: sample same prompt across all trained configs with a shared seed, print side by side, write reports/compare.md - tests/test_configs.py: enforces one-parameter-at-a-time (only intended fields differ from v1) + small param count (3 tests). Full suite: 29 passing.
73 lines
2.4 KiB
Python
73 lines
2.4 KiB
Python
"""Guardrail: each degradation config changes exactly one thing from v1.
|
|
|
|
Enforces the controlled-experiment discipline — if a variant accidentally
|
|
drifts from the baseline in an unintended field, this fails. Forward-only.
|
|
Run: `.venv/bin/python tests/test_configs.py`
|
|
"""
|
|
|
|
import importlib
|
|
import sys
|
|
from dataclasses import asdict
|
|
from pathlib import Path
|
|
|
|
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
|
|
|
|
# For each variant: which model fields and which train fields may differ from v1.
|
|
EXPECTED_DIFFS = {
|
|
"small": ({"n_layers", "d_model", "n_heads", "d_ff"}, set()),
|
|
"short": (set(), {"max_iters"}),
|
|
"ctx": ({"context_length"}, set()),
|
|
"data": (set(), {"max_train_tokens"}),
|
|
}
|
|
|
|
|
|
def _diff_keys(a, b):
|
|
return {k for k in a if a[k] != b[k]}
|
|
|
|
|
|
def test_each_variant_changes_only_intended_fields():
|
|
v1 = importlib.import_module("v1")
|
|
base_model, base_train = asdict(v1.model), asdict(v1.train)
|
|
for name, (exp_model, exp_train) in EXPECTED_DIFFS.items():
|
|
cfg = importlib.import_module(name)
|
|
model_diffs = _diff_keys(asdict(cfg.model), base_model)
|
|
train_diffs = _diff_keys(asdict(cfg.train), base_train)
|
|
assert model_diffs == exp_model, f"{name}: model diffs {model_diffs} != {exp_model}"
|
|
assert train_diffs == exp_train, f"{name}: train diffs {train_diffs} != {exp_train}"
|
|
|
|
|
|
def test_each_variant_builds_a_model():
|
|
for name in EXPECTED_DIFFS:
|
|
cfg = importlib.import_module(name)
|
|
model = GPT(cfg.model)
|
|
assert model.num_params() > 0
|
|
|
|
|
|
def test_small_param_count():
|
|
cfg = importlib.import_module("small")
|
|
model = GPT(cfg.model)
|
|
# tok_emb 4096*64 + pos_emb 256*64
|
|
# + 2*(qkv 64*192 + proj 64*64 + 2 norms*64 + mlp_fc 64*256 + mlp_proj 256*64)
|
|
# + final norm 64
|
|
expected = 4096 * 64 + 256 * 64 + 2 * (64 * 192 + 64 * 64 + 2 * 64 + 64 * 256 + 256 * 64) + 64
|
|
assert expected == 377_152
|
|
assert model.num_params() == expected
|
|
# The embedding table dominates: non-embedding capacity is tiny.
|
|
assert model.num_params(non_embedding=True) == 98_624
|
|
|
|
|
|
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()
|