diff --git a/README.md b/README.md index 9ac0ae6..a93e5fc 100644 --- a/README.md +++ b/README.md @@ -166,6 +166,38 @@ experiment-llm/ (v1: grammatical, story-arc coherent; plot logic wobbles as expected) - [x] Every phase re-runnable from a clean clone with documented commands +## Degradation experiments (iteration 1.5) + +A controlled study of *how* a model gets worse: start from the v1 baseline and +change exactly **one** variable per run, so any difference in the output is +attributable to that variable. Each variant config is literally `v1` with one +field replaced (see `configs/*.py`); `tests/test_configs.py` enforces that only +the intended field changes. + +| config | one change from v1 | isolates | expected failure mode | +|---|---|---|---| +| `v1` | — (baseline) | — | grammatical, coherent story arcs | +| `small` | 2L/64d, ~0.38M params | capacity | generic, repetitive, forgets rare words | +| `short` | `max_iters` 20k→2k | training time | half-learned grammar, weaker structure | +| `ctx` | `context_length` 256→64 | long-range coherence | fine locally, plot threads break sooner | +| `data` | `max_train_tokens`→1M | data quantity | overfits: memorizes, novel prompts collapse | + +Run each (they all train faster than v1; `short` is ~minutes): + +```sh +python src/train.py --config small # then short, ctx, data +python scripts/plot_loss.py --config small +``` + +Then compare generated text across every trained variant, side by side (same +prompt, same RNG seed, so differences reflect the model): + +```sh +python scripts/compare.py # default prompt, all configs +python scripts/compare.py --suite # every suite prompt -> reports/compare.md +python scripts/compare.py --configs v1,small # a subset +``` + ## Iteration 2 candidates (out of scope for now) - Scale toward the coherence sweet spot (~15–30M params, longer training) diff --git a/configs/ctx.py b/configs/ctx.py new file mode 100644 index 0000000..a86a121 --- /dev/null +++ b/configs/ctx.py @@ -0,0 +1,17 @@ +"""Degradation experiment: same as v1, but a smaller context window. + +Isolates *long-range coherence*. Only context_length changes from v1 +(256 -> 64). Local fluency should survive; callbacks and plot threads that span +more than a few sentences should break sooner. + +Honest caveat: shrinking context also lowers tokens/step (batch * context), so +this run sees fewer total tokens than v1 over the same number of steps. The +qualitative coherence effect is still the dominant, visible change. +""" + +from dataclasses import replace + +import v1 + +model = replace(v1.model, context_length=64) +train = v1.train diff --git a/configs/data.py b/configs/data.py new file mode 100644 index 0000000..319bfaf --- /dev/null +++ b/configs/data.py @@ -0,0 +1,17 @@ +"""Degradation experiment: same as v1, but trained on a tiny slice of data. + +Isolates *data quantity*. Only max_train_tokens changes from v1: the model sees +~1M training tokens instead of 555M. Expect overfitting — train loss falls while +val loss (measured on the full held-out set) climbs, and novel prompts collapse +as the model regurgitates memorized stories. + +Tip: this overfits within a few thousand steps. Watch the val loss turn upward +in the log and stop early (Ctrl-C) rather than running all 20,000 steps. +""" + +from dataclasses import replace + +import v1 + +model = v1.model +train = replace(v1.train, max_train_tokens=1_000_000) diff --git a/configs/short.py b/configs/short.py new file mode 100644 index 0000000..db2dbcc --- /dev/null +++ b/configs/short.py @@ -0,0 +1,13 @@ +"""Degradation experiment: same as v1, but far fewer training steps. + +Isolates *training time*. Only max_iters changes from v1 (20,000 -> 2,000). +The cosine LR schedule anneals fully within these 2,000 steps, so this is the +best model achievable in that budget — not v1 halted mid-schedule. +""" + +from dataclasses import replace + +import v1 + +model = v1.model +train = replace(v1.train, max_iters=2000) diff --git a/configs/small.py b/configs/small.py new file mode 100644 index 0000000..bb0aa5c --- /dev/null +++ b/configs/small.py @@ -0,0 +1,14 @@ +"""Degradation experiment: same as v1, but a much smaller model. + +Isolates *capacity*. Only the model dimensions change; context length and the +training schedule match v1 exactly. ~0.38M params vs v1's 4.26M — and only +~0.1M of those are non-embedding (the 4,096-token embedding table dominates a +model this small, so the actual 'reasoning' capacity shrinks even harder). +""" + +from dataclasses import replace + +import v1 + +model = replace(v1.model, n_layers=2, d_model=64, n_heads=2, d_ff=256) +train = v1.train diff --git a/configs/v1.py b/configs/v1.py index 88e80e7..5a3f260 100644 --- a/configs/v1.py +++ b/configs/v1.py @@ -25,6 +25,7 @@ class ModelConfig: class TrainConfig: # Data / batching. tokens per step = batch_size * context_length = 16,384. batch_size: int = 64 + max_train_tokens: int | None = None # None = whole train set; set to starve data # AdamW optimiser. learning_rate: float = 6e-4 # peak LR (reached after warmup) diff --git a/scripts/compare.py b/scripts/compare.py new file mode 100644 index 0000000..3419a3b --- /dev/null +++ b/scripts/compare.py @@ -0,0 +1,90 @@ +"""Compare generated text across the degradation experiments. + +Samples the same prompt(s) from every config that has a trained checkpoint, +using the same RNG seed for each so differences reflect the model, not sampling +luck. Prints the results side by side and writes reports/compare.md. + +Usage: + python scripts/compare.py # default prompt, all configs + python scripts/compare.py --prompt "The dog and the cat" + python scripts/compare.py --suite # every suite prompt + python scripts/compare.py --configs v1,small,ctx # a subset +""" + +import argparse +import sys +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT / "src")) + +from sample import ( # noqa: E402 + SUITE_PROMPTS, + TOK_PATH, + load_model, + pick_device, + sample_text, +) +from tokenizer import BPETokenizer # noqa: E402 + +# Baseline first, then the four one-parameter degradations. +DEFAULT_CONFIGS = ["v1", "small", "short", "ctx", "data"] + +LABELS = { + "v1": "v1 (baseline)", + "small": "small (less capacity)", + "short": "short (less training)", + "ctx": "ctx (smaller context)", + "data": "data (less data)", +} + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--configs", default=",".join(DEFAULT_CONFIGS)) + ap.add_argument("--prompt", default="Once upon a time,") + ap.add_argument("--suite", action="store_true") + ap.add_argument("--max-new-tokens", type=int, default=200) + ap.add_argument("--temperature", type=float, default=0.8) + ap.add_argument("--top-k", type=int, default=200) + ap.add_argument("--seed", type=int, default=1234) + args = ap.parse_args() + + configs = [c.strip() for c in args.configs.split(",") if c.strip()] + prompts = SUITE_PROMPTS if args.suite else [args.prompt] + + device = pick_device() + tok = BPETokenizer.load(TOK_PATH) + + # Load each available model once; warn about any not yet trained. + models = {} + for name in configs: + if (ROOT / "checkpoints" / name / "ckpt.pt").exists(): + models[name] = load_model(name, device) + else: + print(f"[skip {name}: no checkpoint — train it first]") + if not models: + sys.exit("no trained checkpoints found among: " + ", ".join(configs)) + + out = ["# Degradation comparison", ""] + for prompt in prompts: + print(f"\n########## prompt: {prompt!r} ##########") + out += [f"## Prompt: {prompt}", ""] + for name, model in models.items(): + torch.manual_seed(args.seed) # same RNG state per model = fair compare + text = sample_text(model, tok, prompt, device, + args.max_new_tokens, args.temperature, args.top_k) + label = LABELS.get(name, name) + print(f"\n----- {label} -----\n{text}") + out += [f"### {label}", "", f"> {text}", ""] + + out_path = ROOT / "reports" / "compare.md" + out_path.parent.mkdir(exist_ok=True) + out_path.write_text("\n".join(out)) + print(f"\nsaved {out_path}") + + +if __name__ == "__main__": + main() diff --git a/src/train.py b/src/train.py index 12020d0..8a7b3bf 100644 --- a/src/train.py +++ b/src/train.py @@ -116,6 +116,10 @@ def main() -> None: tokens_dir = ROOT / "data" / "tokens" train_data = load_tokens(tokens_dir / "train.bin") val_data = load_tokens(tokens_dir / "val.bin") + # Data-starvation experiments restrict training to a slice; val stays full so + # the generalization gap (overfitting) is visible. + if tc.max_train_tokens is not None: + train_data = train_data[: tc.max_train_tokens] print(f"train tokens: {len(train_data):,} | val tokens: {len(val_data):,}") model = GPT(mc).to(device) diff --git a/tests/test_configs.py b/tests/test_configs.py new file mode 100644 index 0000000..8f1e1ec --- /dev/null +++ b/tests/test_configs.py @@ -0,0 +1,72 @@ +"""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()