diff --git a/baseline/experiments/nanogpt_memorization/.gitignore b/baseline/experiments/nanogpt_memorization/.gitignore new file mode 100644 index 0000000..a89b47b --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +.pytest_cache/ +*.pyc +*.pt +results/ diff --git a/baseline/experiments/nanogpt_memorization/README.md b/baseline/experiments/nanogpt_memorization/README.md new file mode 100644 index 0000000..36edc49 --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/README.md @@ -0,0 +1,282 @@ +# nanoGPT memorization: AdamW versus Muon + +Prepared September 14, 2026. This is a new, isolated experiment directory. +**Status: protocol and runner implemented; 14 local unit/contract tests pass. +No actual AdamW/Muon training campaign or numerical WeightWatcher integration +run has been executed for this suite.** The WeightWatcher unit test uses a mock. + +## Question and scope + +Which kinds of memorization occur under AdamW and plain Muon, and which +layerwise spectral changes accompany acquisition, generalization, persistence, +and interference? Memorization is established by behavioral probes and explicit +training exposure, not by assigning a meaning to alpha in advance. + +These are controlled synthetic, autoregressive **answer-suffix prediction** +tasks. They retain the existing nanoGPT architecture and optimizer profiles, +but replace FineWeb-Edu with generated token sequences and mask prompt/padding +loss. They are not natural-language pretraining runs, cannot establish real-text +privacy leakage, and their losses must not be compared numerically with the +FineWeb baseline. All tokens and records are synthetic; no private data are used. + +## Frozen source settings + +`configs/suite.json` pins the inspected upstream commit +`3749c36334382a20e48bfe2473c1dc4a1470a830` and the Git blob hashes of the original +model, optimizer implementation, and configuration. The runner imports these +implementations rather than introducing another Muon implementation. It refuses +source-hash mismatches, preventing silent recipe drift. + +Source configuration: +`baseline/experiments/nanogpt_one_head_2026_08_21_baseline/configs/baseline.yaml`. + +| Setting | Repository recipe | +|---|---| +| Model | 1 transformer block; 1 attention head; width 128; context 256 | +| Vocabulary / weights | 50,257; tied token embedding and output head | +| Dropout / bias | 0.0 / false | +| Batch | 4 sequences x 8 accumulation steps = 32 sequences/update | +| Input positions / update | 8,192, including synthetic context/padding | +| Gradient norm clipping | 1.0, shared across all parameters | +| AdamW | LR 6e-4 -> 6e-5; betas (0.9, 0.95); epsilon 1e-8; decay 0.1 | +| Plain Muon hidden matrices | LR 0.02 -> 0.002; momentum 0.95; Nesterov on; 5 Newton-Schulz steps; epsilon 1e-7; decay 0.01 | +| Muon auxiliary AdamW | LR 3e-4 -> 3e-5; betas (0.9, 0.95); epsilon 1e-8; decay 0.01 | +| Warmup fractions | AdamW 0.01; Muon 0.05 | +| Schedule | Warmup/cosine over ceil(80M/8,192) = 9,766 updates, then LR floor | +| Primary final budget | 39,063 updates, following the four-corpus-equivalent-epoch baseline geometry | +| Paired seeds | 1337, 2027, 4099, 31415, 271828 | + +**These are the repository-backed starting settings, not a claimed optimal +configuration.** The source campaign explicitly says it has no checked-in +nanoGPT qualification lock and should be called a baseline. Its existing active +campaign uses MuonClip; this experiment intentionally selects the separately +provided **plain `muon` profile**, not `muon_clip`. + +The repository's `baseline/FINAL_BASELINE_QUALIFICATION.md` also states that a new +dataset/objective invalidates an earlier optimum claim. Any subsequent tuning +must use a preregistered development-only search, a fixed compute budget, and a +new lock before interpreting protected tests. This suite does not fabricate such +a search or a winner. + +## Experiment matrix: eight conditions, six related phenomena + +| Conditions | Manipulation and controls | Main interpretation | +|---|---|---| +| `verbatim`, `verbatim_absent` | Random 64-token prefixes and independent 32-token suffixes. Sixteen canaries per lifetime dose: 0, 1, 4, 16, 64. Injections occur at fixed slots in the first half; the second half contains only background. The absent twin replaces only injection slots with the exact matched background draws. | Exact sequence memorization, exposure dependence, effect of training inclusion, and persistence after withdrawal. | +| `associations` | 128 independent random 8-token keys -> random 8-token values. Train two prompt templates; audit a third. Include wrong-key and disjoint unseen-key controls. | Arbitrary associative memory versus template-specific recall. Unseen random values are not inferable by a genuine rule. | +| `rule_clean` | Modular addition modulo 31; operand pairs split 50/25/25 into train/validation/test before training. | Rule generalization versus fitting observed pairs; delayed generalization is possible, not assumed. | +| `rule_half_noise` | Same operand split; 50% of training labels replaced once with fixed random labels. Score both the assigned and true labels. | Coexistence of useful rule learning and example-specific noise memorization. | +| `rule_random` | Same split; every training label assigned independently at random, then held fixed. | Negative control for genuine addition-rule learning; direct capacity for arbitrary label fitting. Random replacement can coincide with the true label by chance. | +| `forgetting_disjoint`, `forgetting_conflict` | Identical first-half acquisition of mapping A. Second-half B either uses disjoint keys or reuses A keys with different values. Keep optimizer state and the same LR schedule; do not reset the model. | Generic forgetting under unrelated continued learning versus targeted interference/overwriting. | + +The canary doses are **actual total presentations per run**, not “duplicates per +epoch.” `injection_schedule.json` fixes every injection slot; `exposures.json` +tracks realized presentations, including explicit zero counts for audit-only +records. Prefix-length probes of 8, 16, 32, and 64 tokens are run at the final +checkpoint with the same suffix and aligned target positions. Earlier audits +use the full prefix to keep overhead bounded. + +Canaries and the background have 32 scored target tokens per sequence; +associations have 8, and modular addition has 1. Thus there are respectively +1,024, 256, and 32 loss-bearing targets per update. **Match budgets within an +experiment, not by pretending the tasks have equal numbers of supervised +labels.** The 8,192-position budget includes masked positions. + +## Behavioral measurements + +At step zero, early logarithmic checkpoints, regular intervals, 17 evenly spaced +permanent states, the phase boundary, and the final state, record: + +- Conditional suffix NLL, teacher-forced token accuracy, greedy continuation + token accuracy, and whole-suffix exact match, including per-example values. +- Seen/novel-template/wrong-key/unseen-key differences; assigned-label versus + true-label accuracy on corrupted examples; disjoint-pair generalization. +- A and B recall before and after the switch, plus the complete presentation + schedule needed to align retention with time since last exposure. + +For the canary inclusion effect use +`NLL_absent - NLL_present` and `EM_present - EM_absent` for the same seed, canary, +prefix length, and update. This measures the effect of the specified replacement +intervention; it is not an unconditional privacy or global exposure estimate. + +For forgetting, report `EM_A(boundary) - EM_A(t)` together with B acquisition; +contrast disjoint and conflicting B. Do not call ordinary validation deterioration +“memorization” without the appropriate seen/unseen behavioral evidence. + +A delayed separation between high training accuracy and later high held-out rule +accuracy is the grokking question. There is no guarantee this small architecture +and split will exhibit grokking. Any change-point or threshold-based event +classifier must be fixed before the confirmatory run, not fitted to a pleasing +spectral curve after inspection. + +The modular-addition test split appears only in the final audit. It never +selects a checkpoint, hyperparameter, horizon, or optimizer. Periodic audit +scores are diagnostic; all runs finish their fixed budget. The pilot evaluates +up to 16 records per group; full runs evaluate up to 64, or every member of a +smaller group. These are fixed probes, not exhaustive scans of every split. + +## WeightWatcher monitoring + +The runner uses exactly one analysis call on a detached CPU copy at each +spectral checkpoint: + +```python +watcher.analyze( + ERG=True, + randomize=True, + plot=False, + min_evals=20, + fix_fingers="clip_xmax", + max_fingers=10, +) +``` + +WeightWatcher is pinned to **0.7.7**, matching the parent package. Persist every +returned column, bind each row to its matrix name, and expose `alpha_clip_xmax` +from `alpha` and `alpha_raw` from `raw_alpha`. Keep `D`, `rand_distance`, +`ERG_gap`, `num_traps`, `num_fingers`, tail sizes, and fit warnings whenever +returned. Required API fields must exist; absent metrics are not replaced by +invented proxies. Nonfinite/unsupported fits are retained and flagged, not +interpreted as meaningful power laws. + +The six primary matrices are Q, K, V, attention output, MLP input, and MLP output. +**There is no single all-layer alpha headline.** Inspect each matrix separately, +including randomized-control distance, tail support and fit quality. A finite +alpha or a particular numerical threshold does not prove that a small spectrum +is heavy-tailed or that memorization has occurred. The `tail_support_at_least_20` +flag is only a support flag, not a goodness-of-fit or correlation test. + +The much larger token embedding/tied head is deliberately excluded from this +six-matrix summary, as in the source campaign, but remains in saved model +checkpoints. This limits localization claims: memory in those parameters cannot +be excluded by unchanged hidden-matrix spectra. + +`fix_fingers` corrects the diagnostic spectral fit; it does **not** clip trained +weights. The monitor works on CPU copies, isolates Python/NumPy/CPU-Torch RNGs, +and checks the model's tensor hash before and after auditing. Step-zero weights +provide an initialization control. Any predictive spectral analysis must compare +against time/update count and training loss, then validate on held-out complete +seeds; repeated checkpoints/layers are not independent training runs. + +## Optional stricter optimizer comparison + +The primary `--recipe repository` comparison preserves the source profiles, +including their different auxiliary learning rates, warmup fractions and decay. +It compares the **whole training recipes**, not only the matrix-update rule. + +The secondary `--recipe shared_aux_decay` leaves AdamW unchanged and gives Muon +exactly the same auxiliary AdamW LR, LR floor, betas, epsilon, decay, and warmup. +It sets Muon's hidden decay coefficient to 0.003, so the per-step shrinkage +`LR(t) * decay` matches AdamW throughout the common-shaped schedule: +`0.02 * 0.003 = 0.0006 * 0.1`. Merely setting both decay coefficients to 0.1 +would not match shrinkage. Muon's hidden LR, momentum, and Newton-Schulz update +remain its source settings. This is an explicitly labeled diagnostic control, +not a new “optimal” Muon claim, and is stored in a separate recipe directory. + +## Execution + +Use an existing checkout containing this directory and the pinned parent files. +The suite never modifies baseline code or outputs. Use the currently activated +research environment. When installation is needed, from repository root: + +```bash +export RG_MEM_ROOT=/tmp/rg-nanogpt-memorization-20260914 +mkdir -p "$RG_MEM_ROOT/cache/pip" +PIP_CACHE_DIR="$RG_MEM_ROOT/cache/pip" python -m pip install -e baseline/nanogpt_one_head +cd baseline/experiments/nanogpt_memorization +python -m pytest -q tests +``` + +First run the actual model/optimizer/WeightWatcher smoke checks on the intended +hardware. They are required before trusting the integration: + +```bash +python run.py run --stage smoke --condition verbatim --optimizer adamw --seed 1337 --device mps --resume +python run.py run --stage smoke --condition verbatim --optimizer muon --seed 1337 --device mps --resume +``` + +The smoke is two updates with reduced audit populations, but retains the real +model and optimizer architecture. The WeightWatcher dependency is mandatory; +there is no silent “skip WW” production mode. MPS, CUDA, and CPU are explicit +choices; an unavailable accelerator fails rather than falling back silently. +This new runner does not implement TPU/XLA integration. + +Next, run a paired pilot: + +```bash +python run.py run --stage pilot --condition verbatim --optimizer adamw --seed 1337 --device mps --resume +python run.py run --stage pilot --condition verbatim --optimizer muon --seed 1337 --device mps --resume +``` + +`python run.py plan --stage pilot` prints, but does not execute, all 16 pilot +commands: eight conditions x two optimizers x one seed, 2,000 updates each. +`python run.py plan --stage full` prints the complete 80-run, five-seed design +with 39,063 updates per run. It does not launch a large campaign automatically. +A pilot is an engineering/difficulty check and is not proof of asymptotic +behavior; its schedule is a prefix of the full schedule. + +For the stricter control append `--recipe shared_aux_decay` to both optimizer +commands. To use CUDA or CPU, explicitly replace `--device mps`. Do not pool +heterogeneous hardware blocks or recipes. + +Outputs default to `/tmp/rg-nanogpt-memorization-20260914`: + +```text +////seed_/ + manifest.json # configuration, package inventory, hardware, hashes + probe_inventory.json # generated records and audit membership + injection_schedule.json # exact canary presentation slots + metrics.jsonl # checkpoint-level and per-example behavior + exposures.json # realized presentation counts + spectral/step_XXXXXXXX.csv # all raw WW columns plus explicit alpha names + checkpoint_latest.pt # atomic restart: model, optimizers, RNG, counters + model_step_XXXXXXXX.pt # permanent model-only states + complete.json # created only after the final audit succeeds +``` + +The runner prints training loss every 25 updates and a spectral-checkpoint line +at each audit. For a read-only behavioral and per-matrix snapshot: + +```bash +python run.py monitor /tmp/rg-nanogpt-memorization-20260914/pilot/repository/verbatim/muon/seed_1337 +``` + +Resume requires an identical run fingerprint and complete package inventory. +The runner checks paired-arm initialization/data/hardware identity when a peer +manifest is present, uses an exclusive per-run file lock, and rolls incomplete +metrics back to the latest atomic checkpoint. Caches and run outputs are placed +under the explicit temporary root, never mixed into this source directory. +Temporary storage is ephemeral: preserve completed results elsewhere before a +host reset. Full-suite model-only snapshots alone are roughly 36 GB, so do not +launch the entire design without storage planning. + +## Analysis and decision rules + +Compare AdamW and Muon at equal updates, equal realized exposures, and, where +both reach a shared target, matched training loss/accuracy. This last comparison +is a secondary check against merely learning at different speeds; avoid +extrapolation outside their overlapping performance range. + +The replication unit is a complete seed. Report paired per-seed differences, +all five raw seed outcomes, and uncertainty across seeds. Do not inflate the +sample size using layers, canaries, or repeated checkpoints. The one-seed pilot +has no across-seed confidence interval. Analyze recipe and hardware blocks +separately. There is no automated hypothesis-testing or spectral forecasting +report in this initial implementation; the raw audit tables support that +preregistered follow-on analysis. + +## Sources + +- Parent campaign README and `configs/baseline.yaml`: model, training geometry, + distinct Muon/MuonClip recipes, one-call raw/clipped alpha convention. +- `baseline/FINAL_BASELINE_QUALIFICATION.md`: validation-only baseline selection + and the conditions under which a lock is invalidated. +- Carlini et al., *Quantifying Memorization Across Neural Language Models*, + https://arxiv.org/abs/2202.07646: duplication and prompting context motivate + controlled dose and prefix-length probes. +- Carlini et al., *The Secret Sharer*, + https://www.usenix.org/conference/usenixsecurity19/presentation/carlini: + synthetic canary methodology. This implementation does not estimate its + full-space exposure metric. +- Calculated Content, `clip_xmax` feature description, + https://calculatedcontent.com/2023/03/21/weightwatcher-advanced-features-fix_fingers/. diff --git a/baseline/experiments/nanogpt_memorization/RESULTS.md b/baseline/experiments/nanogpt_memorization/RESULTS.md new file mode 100644 index 0000000..6755a37 --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/RESULTS.md @@ -0,0 +1,25 @@ +# Execution ledger — September 14, 2026 + +| Check | Actual status | +|---|---| +| Python source compilation | Passed | +| Local unit/contract tests | 14 passed | +| Plan generation | Passed; smoke/pilot 16 commands, full 80 commands | +| Real WeightWatcher 0.7.7 numerical integration | NOT RUN; unavailable in the construction environment | +| Real repository GPT + AdamW/Muon smoke runs | NOT RUN | +| MPS/CUDA execution and restart integration | NOT RUN | +| Pilot/full training and scientific results | NOT RUN | + +The WeightWatcher contract test uses a mock that checks the analysis arguments, +raw/clipped field preservation, matrix binding, and training-weight/RNG +isolation. It does not establish the actual library's numerical behavior. + +The evaluation test uses a uniform-logit test double. The remaining tests check +suffix masking, exact exposure counts, the counterfactual schedule, independent +RNG sampling, split disjointness, fixed label corruption, template controls, +forgetting/interference arms, context ablations, matched decay factors, and +interrupted log handling. None is a performance experiment. + +No optimizer superiority, alpha/memorization association, grokking, or forgetting +result is claimed. Append actual run paths, fingerprints, device blocks, test +results and complete seed outcomes only after the corresponding executions. diff --git a/baseline/experiments/nanogpt_memorization/configs/suite.json b/baseline/experiments/nanogpt_memorization/configs/suite.json new file mode 100644 index 0000000..1527fdb --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/configs/suite.json @@ -0,0 +1,30 @@ +{ + "name": "nanogpt_memorization_20260914", + "version": 1, + "upstream_commit": "3749c36334382a20e48bfe2473c1dc4a1470a830", + "source_config": "baseline/experiments/nanogpt_one_head_2026_08_21_baseline/configs/baseline.yaml", + "source_blobs": { + "baseline/experiments/nanogpt_one_head_2026_08_21_baseline/configs/baseline.yaml": "7fd3c592afc7fdfea952b1f1aa0f9b44cdc53a2b", + "baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py": "029b42675ce230c6f1a142afaa8730b9f4d4ced3", + "baseline/nanogpt_one_head/src/rg_nanogpt_one_head/optimizers.py": "b1a97ad684afd882ffcff4ae0fe5d70140e7d22c" + }, + "optimizers": ["adamw", "muon"], + "seeds": [1337, 2027, 4099, 31415, 271828], + "data_seed": 20260914, + "conditions": ["verbatim", "verbatim_absent", "associations", "rule_clean", "rule_half_noise", "rule_random", "forgetting_disjoint", "forgetting_conflict"], + "stages": { + "smoke": {"steps": 2, "eval_every": 1, "canaries_per_dose": 1, "doses": [0, 1, 2], "probe_limit": 1}, + "pilot": {"steps": 2000, "eval_every": 250, "canaries_per_dose": 16, "doses": [0, 1, 4, 16, 64], "probe_limit": 16}, + "full": {"steps": 39063, "eval_every": 500, "canaries_per_dose": 16, "doses": [0, 1, 4, 16, 64], "probe_limit": 64} + }, + "associations": 128, + "value_tokens": 8, + "canary_prefix_tokens": 64, + "canary_suffix_tokens": 32, + "prefix_lengths_final": [8, 16, 32, 64], + "modulus": 31, + "early_eval_steps": [0, 1, 2, 4, 8, 16, 32, 64, 128], + "weightwatcher": {"version": "0.7.7", "ERG": true, "randomize": true, "plot": false, "min_evals": 20, "fix_fingers": "clip_xmax", "max_fingers": 10}, + "output_root": "/tmp/rg-nanogpt-memorization-20260914", + "notes": "Repository-backed starting settings, not a qualified optimum on these NEW synthetic suffix-prediction tasks. No training results are included." +} diff --git a/baseline/experiments/nanogpt_memorization/run.py b/baseline/experiments/nanogpt_memorization/run.py new file mode 100644 index 0000000..ed773a6 --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/run.py @@ -0,0 +1,556 @@ +#!/usr/bin/env python3 +"""Controlled, suffix-supervised nanoGPT memorization experiments. + +Uses the repository GPT, AdamW/Muon implementations and frozen baseline profile. +No corpus downloads; all data are synthetic. WeightWatcher never edits training +weights. Importing this module needs only NumPy; training imports are lazy. +""" +from __future__ import annotations + +import argparse +from collections import Counter +from dataclasses import asdict, dataclass, replace +import hashlib +import importlib.metadata +import json +import math +import os +from pathlib import Path +import platform +import random +import sys +from typing import Any + +import numpy as np + +HERE = Path(__file__).resolve().parent +REPO = HERE.parents[2] +CONFIG = HERE / "configs/suite.json" + + +def digest(value: Any) -> str: + return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + +def atomic_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temp = path.with_suffix(path.suffix + ".tmp") + temp.write_text(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n") + os.replace(temp, path) + + +@dataclass(frozen=True) +class Record: + eid: str + group: str + prompt: tuple[int, ...] + target: tuple[int, ...] + + +def pack(records: list[Record], block_size: int) -> tuple[np.ndarray, np.ndarray]: + """Left-pad to the original context length; score only the answer suffix.""" + if not records: + raise ValueError("empty batch") + x = np.zeros((len(records), block_size), dtype=np.int64) + y = np.full_like(x, -100) + for i, record in enumerate(records): + if not record.prompt or not record.target: + raise ValueError("empty prompt/target") + seq = record.prompt + record.target + if len(seq) > block_size + 1: + raise ValueError("record exceeds the model context") + padded = np.array((0,) * (block_size + 1 - len(seq)) + seq) + x[i] = padded[:-1] + start = block_size - len(record.target) + y[i, start:] = padded[start + 1:] + return x, y + + +class Study: + """Stateless step-indexed sampling, paired across optimizer arms.""" + def __init__(self, cfg: dict, condition: str, stage: str, seed: int, batch: int): + if condition not in cfg["conditions"]: + raise ValueError(condition) + self.cfg, self.condition, self.stage = cfg, condition, cfg["stages"][stage] + self.seed, self.batch = seed, batch + self.steps = int(self.stage["steps"]) + self.boundary = max(1, self.steps // 2) + self.rng = np.random.default_rng(cfg["data_seed"]) + self.audit: list[Record] = [] + self.train: list[Record] = [] + self.second: list[Record] = [] + self.injections: dict[int, Record] = {} + self._build() + + def tokens(self, n: int) -> tuple[int, ...]: + return tuple(int(x) for x in self.rng.integers(16, 272, n)) + + def _build(self) -> None: + c, n = self.condition, int(self.cfg["associations"]) + if c.startswith("verbatim"): + doses = [] + for dose in self.stage["doses"]: + for i in range(self.stage["canaries_per_dose"]): + r = Record(f"canary_d{dose}_{i}", f"dose_{dose}", + self.tokens(self.cfg["canary_prefix_tokens"]), + self.tokens(self.cfg["canary_suffix_tokens"])) + self.audit.append(r) + doses.extend([r] * dose) + # Exact lifetime presentations, not corpus-duplication multipliers. + slots = self.boundary * self.batch + if len(doses) > slots: + raise ValueError("injection window too small for requested doses") + schedule_rng = np.random.default_rng(np.random.SeedSequence([self.cfg["data_seed"], self.seed, 99])) + where = schedule_rng.choice(slots, len(doses), replace=False) + self.injections = dict(zip(map(int, where), doses)) + elif c.startswith("rule_"): + p = int(self.cfg["modulus"]) + pairs = self.rng.permutation(p * p) + train_end, val_end = int(0.5 * len(pairs)), int(0.75 * len(pairs)) + fraction = {"rule_clean": 0.0, "rule_half_noise": 0.5, "rule_random": 1.0}[c] + corrupt = set(self.rng.permutation(train_end)[:round(fraction * train_end)]) + # Sample replacements even in clean condition so all corpus identities align. + replacements = self.rng.integers(p, size=train_end) + for i, pair in enumerate(pairs): + a, b = divmod(int(pair), p) + truth = (16 + (a + b) % p,) + prompt = (1, 16 + a, 2, 16 + b, 3) + if i < train_end: + target = (16 + int(replacements[i]),) if i in corrupt else truth + kind = "randomized" if i in corrupt else "clean" + r = Record(f"pair_{a}_{b}", f"train_observed_{kind}", prompt, target) + self.train.append(r) + self.audit.append(r) + if i in corrupt: + self.audit.append(replace(r, group="train_true_randomized", target=truth)) + else: + group = "validation_rule" if i < val_end else "test_rule" + self.audit.append(Record(f"pair_{a}_{b}", group, prompt, truth)) + else: + # Unique, disjoint keys. B values are explicitly different from A values. + keys: list[tuple[int, ...]] = [] + while len(keys) < 3 * n: + key = self.tokens(8) + if key not in keys: + keys.append(key) + values = [self.tokens(self.cfg["value_tokens"]) for _ in keys] + for i in range(n): + if values[n + i] == values[i]: + raise RuntimeError("A/B target collision; change the preregistered data seed") + a = Record(f"A_{i}", "memory_A", (1, 4) + keys[i] + (3,), values[i]) + bkey = keys[i] if c == "forgetting_conflict" else keys[n + i] + b = Record(f"B_{i}", "memory_B", (1, 4) + bkey + (3,), values[n + i]) + self.train.append(a) + self.second.append(b) + self.audit.append(a) + if c.startswith("forgetting"): + self.audit.append(b) + else: + self.audit.extend([ + replace(a, group="seen_key_new_template", prompt=(1, 6) + keys[i] + (3,)), + replace(a, group="wrong_key_control", prompt=(1, 4) + keys[(i + 1) % n] + (3,)), + Record(f"unseen_{i}", "unseen_key_control", (1, 4) + keys[2*n+i] + (3,), values[2*n+i]), + ]) + + def sample(self, step: int) -> list[Record]: + if not 0 <= step < self.steps: + raise ValueError("step outside the frozen budget") + rng = np.random.default_rng(np.random.SeedSequence([self.cfg["data_seed"], self.seed, step, 7])) + out = [] + for i in range(self.batch): + if self.condition.startswith("verbatim"): + # Same background draws in present/absent conditions, even at injections. + prefix = tuple(map(int, rng.integers(16, 272, self.cfg["canary_prefix_tokens"]))) + suffix = tuple(16 + (v - 16 + 1) % 256 for v in prefix[-self.cfg["canary_suffix_tokens"]:]) + r = Record("background", "background", prefix, suffix) + if self.condition == "verbatim": + r = self.injections.get(step * self.batch + i, r) + else: + pool = self.second if self.condition.startswith("forgetting") and step >= self.boundary else self.train + r = pool[int(rng.integers(len(pool)))] + if self.condition == "associations": + # Two seen templates; the third remains audit-only. + template = int(rng.choice([4, 5])) + r = replace(r, prompt=(1, template) + r.prompt[2:]) + out.append(r) + return out + + def probes(self, final: bool = False) -> list[Record]: + out, counts = [], Counter() + for r in self.audit: + if r.group == "test_rule" and not final: + continue + if counts[r.group] >= self.stage["probe_limit"]: + continue + counts[r.group] += 1 + lengths = self.cfg["prefix_lengths_final"] if final and self.condition.startswith("verbatim") else [len(r.prompt)] + for length in lengths: + group = f"{r.group}/prefix_{length}" if self.condition.startswith("verbatim") else r.group + out.append(replace(r, group=group, prompt=r.prompt[-length:])) + return out + + def identity(self) -> str: + return digest({"audit": [asdict(r) for r in self.audit], + "injections": {str(k): r.eid for k, r in self.injections.items()}, + "train": [asdict(r) for r in self.train], + "second": [asdict(r) for r in self.second]}) + + +def resolve_profile(source: dict, optimizer: str, recipe: str) -> dict: + """Optional control: same auxiliary AdamW and per-step matrix decay factors.""" + import copy + profile = copy.deepcopy(source["optimizer_profiles"][optimizer]) + if recipe == "shared_aux_decay" and optimizer == "muon": + adam = source["optimizer_profiles"]["adamw"] + ratio = adam["learning_rate"] / profile["matrix_learning_rate"] + if not math.isclose(ratio, adam["min_learning_rate"] / profile["matrix_min_learning_rate"]): + raise ValueError("cannot match decay factors with unequal LR floor ratios") + profile.update(aux_learning_rate=adam["learning_rate"], + aux_min_learning_rate=adam["min_learning_rate"], + aux_weight_decay=adam["weight_decay"], + matrix_weight_decay=adam["weight_decay"] * ratio, + warmup_fraction=adam["warmup_fraction"], + lr_schedule_epochs=adam["lr_schedule_epochs"], + beta1=adam["beta1"], beta2=adam["beta2"], epsilon=adam["epsilon"]) + return profile + + +def read_rows(path: Path) -> list[dict]: + """Only an interrupted final JSONL row may be ignored.""" + lines = path.read_text().splitlines() + out = [] + for i, line in enumerate(lines): + try: + out.append(json.loads(line)) + except json.JSONDecodeError: + if i != len(lines) - 1: + raise + return out + + +def load_source(cfg: dict) -> dict: + """Fail closed instead of silently switching baseline/profile versions.""" + import yaml + for rel, expected in cfg["source_blobs"].items(): + content = (REPO / rel).read_bytes() + blob = hashlib.sha1(f"blob {len(content)}\0".encode() + content).hexdigest() + if blob != expected: + raise RuntimeError(f"upstream source changed: {rel}; review and version suite.json") + return yaml.safe_load((REPO / cfg["source_config"]).read_text()) + + +def set_environment(root: Path) -> Path: + root = root.expanduser().resolve() + valid = any(root.is_relative_to(Path(base)) and root != Path(base) for base in ("/tmp", "/private/tmp")) + if not valid: + raise ValueError("output root must be a dedicated directory beneath /tmp or /private/tmp") + root.mkdir(parents=True, exist_ok=True) + for key, sub in {"HOME": "home", "XDG_CACHE_HOME": "xdg/cache", "XDG_CONFIG_HOME": "xdg/config", + "XDG_DATA_HOME": "xdg/data", "XDG_STATE_HOME": "xdg/state", "MPLCONFIGDIR": "matplotlib", + "HF_HOME": "huggingface", "TORCH_HOME": "torch", "PIP_CACHE_DIR": "pip", + "TIKTOKEN_CACHE_DIR": "tiktoken", "TMPDIR": "tmp"}.items(): + path = root / "cache" / sub + path.mkdir(parents=True, exist_ok=True) + os.environ[key] = str(path) + os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + return root + + +def weightwatch(model, dest: Path, cfg: dict, seed: int, step: int) -> None: + import torch + import weightwatcher as ww + from rg_nanogpt_one_head.model import transformer_matrix_items + if importlib.metadata.version("weightwatcher") != cfg["version"]: + raise RuntimeError("WeightWatcher version must match the suite pin") + np_state, py_state = np.random.get_state(), random.getstate() + try: + # Creating nn.Linear modules consumes the CPU RNG: isolate that too. + with torch.random.fork_rng(devices=[]): + np.random.seed((seed + step + 104729) % (2**32)) + random.seed(seed + step + 104729) + torch.random.default_generator.manual_seed(seed + step + 104729) + holder = torch.nn.ModuleDict() + for name, _, _, w in transformer_matrix_items(model): + layer = torch.nn.Linear(w.shape[1], w.shape[0], bias=False, device="cpu") + with torch.no_grad(): + layer.weight.copy_(w.detach().float().cpu()) + holder[name] = layer + args = {k: v for k, v in cfg.items() if k != "version"} + table = ww.WeightWatcher(model=holder).analyze(**args) + required = {"alpha", "raw_alpha", "D", "rand_distance", "ERG_gap", "num_traps"} + if not required.issubset(table.columns) or len(table) != len(holder): + raise RuntimeError(f"unexpected WW schema/coverage: {list(table.columns)}, rows={len(table)}") + # Retain all WW columns; never substitute a proxy for an unavailable field. + table["alpha_clip_xmax"] = table["alpha"] + table["alpha_raw"] = table["raw_alpha"] + table["step"], table["diagnostic_seed"] = step, seed + step + 104729 + names = [] + for _, row in table.iterrows(): + text = " ".join(str(row.get(k, "")) for k in ("longname", "name")) + matches = [name for name in holder if name in text] + if len(matches) != 1: + raise RuntimeError(f"cannot bind WW row to a matrix: {text}") + names.append(matches[0]) + if len(set(names)) != len(holder): + raise RuntimeError("duplicate/missing matrix in WeightWatcher output") + table["matrix_name"] = names + # This flag concerns fit support ONLY, not evidence of learned correlations. + tail = table["num_pl_spikes"] if "num_pl_spikes" in table else np.full(len(table), np.nan) + table["tail_support_at_least_20"] = np.asarray(tail, dtype=float) >= 20 + table["finite_alpha_and_D"] = np.isfinite(table["alpha"]) & np.isfinite(table["D"]) + dest.parent.mkdir(parents=True, exist_ok=True) + temp = dest.with_suffix(".tmp") + table.to_csv(temp, index=False) + os.replace(temp, dest) + finally: + np.random.set_state(np_state) + random.setstate(py_state) + + +def evaluate(model, records: list[Record], batch_size: int, device: str) -> dict: + import torch + import torch.nn.functional as F + was_training = model.training + model.eval() + grouped: dict[str, list[dict]] = {} + try: + with torch.inference_mode(): + # Separate target lengths so autoregressive continuations have equal context. + for length in sorted({len(r.target) for r in records}): + rows = [r for r in records if len(r.target) == length] + for start in range(0, len(rows), batch_size): + batch = rows[start:start + batch_size] + x, y = pack(batch, model.cfg.block_size) + x, y = torch.as_tensor(x, device=device), torch.as_tensor(y, device=device) + logits, _ = model(x) + losses = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1), + reduction="none", ignore_index=-100).reshape(y.shape) + valid = y.ne(-100) + nll = (losses.sum(1) / valid.sum(1)).cpu().tolist() + accuracy = ((logits.argmax(-1).eq(y) & valid).sum(1) / valid.sum(1)).cpu().tolist() + prompt = x[:, :model.cfg.block_size + 1 - length] + generated = model.generate_greedy(prompt, length)[:, -length:].cpu().numpy() + for j, r in enumerate(batch): + target = np.array(r.target) + entry = {"eid": r.eid, "nll": nll[j], "teacher_forced_accuracy": accuracy[j], + "exact_match": float(np.array_equal(generated[j], target)), + "continuation_token_accuracy": float((generated[j] == target).mean())} + grouped.setdefault(r.group, []).append(entry) + return {group: {"n": len(rows), "mean": {key: float(np.mean([r[key] for r in rows])) + for key in ("nll", "teacher_forced_accuracy", "exact_match", "continuation_token_accuracy")}, + "examples": rows} for group, rows in grouped.items()} + finally: + model.train(was_training) + + +def state_digest(model) -> str: + h = hashlib.sha256() + for name, value in sorted(model.state_dict().items()): + array = value.detach().cpu().contiguous().numpy() + h.update(name.encode()); h.update(str(array.dtype).encode()) + h.update(str(array.shape).encode()); h.update(array.tobytes()) + return h.hexdigest() + + +def run(args, cfg: dict) -> None: + root = set_environment(Path(args.root or cfg["output_root"])) + source = load_source(cfg) + sys.path.insert(0, str(REPO / "baseline/nanogpt_one_head/src")) + import torch + from rg_nanogpt_one_head.model import GPT, GPTConfig + from rg_nanogpt_one_head.optimizers import make_optimizer_handles, set_learning_rates, zero_grad, optimizer_step + import fcntl + if importlib.metadata.version("weightwatcher") != cfg["weightwatcher"]["version"]: + raise RuntimeError("install the repository's pinned WeightWatcher before running") + if args.device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA requested but unavailable") + if args.device == "mps" and not torch.backends.mps.is_available(): + raise RuntimeError("MPS requested but unavailable") + torch.set_float32_matmul_precision("highest") + torch.use_deterministic_algorithms(True) + torch.backends.cudnn.benchmark = False + if hasattr(torch.backends.cuda, "matmul"): + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.manual_seed(args.seed) + model = GPT(GPTConfig(**source["model"])).to(args.device) + initial_hash = state_digest(model) + profile = resolve_profile(source, args.optimizer, args.recipe) + handles = make_optimizer_handles(model, profile) + train = source["training"] + effective_batch = train["batch_size"] * train["grad_accum_steps"] + study = Study(cfg, args.condition, args.stage, args.seed, effective_batch) + tokens_per_step = effective_batch * source["model"]["block_size"] + schedule_steps = math.ceil(source["dataset"]["train_tokens"] * profile["lr_schedule_epochs"] / tokens_per_step) + warmup = min(schedule_steps - 1, math.ceil(schedule_steps * profile["warmup_fraction"])) + run_dir = root / args.stage / args.recipe / args.condition / args.optimizer / f"seed_{args.seed}" + run_dir.mkdir(parents=True, exist_ok=True) + lock = (run_dir / ".lock").open("a") + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + device_info = {"device": args.device, "platform": platform.platform(), "machine": platform.machine(), + "processor": platform.processor(), "python": platform.python_version(), + "torch": torch.__version__, "numpy": np.__version__, + "weightwatcher": importlib.metadata.version("weightwatcher"), + "packages": sorted((d.metadata.get("Name", "unknown"), d.version) + for d in importlib.metadata.distributions())} + if args.device == "cuda": + device_info["accelerator"] = str(torch.cuda.get_device_properties(0)) + device_info["cuda"] = torch.version.cuda + if args.device == "mps": + import subprocess + device_info["accelerator"] = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"], text=True).strip() + manifest = {"suite": cfg, "source_model": source["model"], "profile": profile, + "condition": args.condition, "stage": args.stage, "recipe": args.recipe, "seed": args.seed, + "initial_model_sha256": initial_hash, "data_sha256": study.identity(), + "runner_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "input_tokens_per_step": tokens_per_step, "schedule_steps": schedule_steps, + "warmup_steps": warmup, "device": device_info, + "objective": "answer_suffix_cross_entropy", "protected_test_used_for_selection": False} + fingerprint = digest(manifest) + manifest["fingerprint"] = fingerprint + manifest = json.loads(json.dumps(manifest)) + for other in cfg["optimizers"]: + if other == args.optimizer: + continue + peer = run_dir.parent.parent / other / run_dir.name / "manifest.json" + if peer.exists(): + paired = json.loads(peer.read_text()) + for key in ("initial_model_sha256", "data_sha256", "source_model", "suite", + "stage", "recipe", "seed", "runner_sha256", "device"): + if paired[key] != manifest[key]: + raise RuntimeError(f"paired optimizer arm mismatch: {key}; use a new root") + latest = run_dir / "checkpoint_latest.pt" + metrics_path = run_dir / "metrics.jsonl" + first, exposures = 0, Counter() + if (run_dir / "manifest.json").exists(): + old = json.loads((run_dir / "manifest.json").read_text()) + if old != manifest: + raise RuntimeError("run fingerprint changed; use a new output root") + if not args.resume: + raise RuntimeError("run exists; pass --resume or use a new root") + if (run_dir / "complete.json").exists(): + print(f"Already complete: {run_dir}"); return + if not latest.exists(): + raise RuntimeError("no restart checkpoint; inspect the failed preflight and use a new root") + checkpoint = torch.load(latest, map_location="cpu", weights_only=False) + if checkpoint["fingerprint"] != fingerprint: + raise RuntimeError("checkpoint fingerprint mismatch") + model.load_state_dict(checkpoint["model"]) + for handle, state in zip(handles, checkpoint["optimizers"], strict=True): + handle.optimizer.load_state_dict(state) + torch.set_rng_state(checkpoint["torch_rng"]) + if args.device == "cuda": + torch.cuda.set_rng_state_all(checkpoint["device_rng"]) + elif args.device == "mps": + torch.mps.set_rng_state(checkpoint["device_rng"]) + first, exposures = checkpoint["step"], Counter(checkpoint["exposures"]) + # Discard incomplete rows/files after the last atomic restart point. + if metrics_path.exists(): + rows = [row for row in read_rows(metrics_path) if row["step"] <= first] + metrics_path.write_text("".join(json.dumps(row) + "\n" for row in rows)) + for path in (run_dir / "spectral").glob("step_*.csv"): + if int(path.stem.split("_")[-1]) > first: + path.unlink() + else: + atomic_json(run_dir / "manifest.json", manifest) + atomic_json(run_dir / "probe_inventory.json", [asdict(r) for r in study.audit]) + atomic_json(run_dir / "injection_schedule.json", {str(k): r.eid for k, r in study.injections.items()}) + permanent = {round(study.steps * i / 16) for i in range(17)} + + def checkpoint(step: int, loss: float | None) -> None: + before = state_digest(model) + audit = evaluate(model, study.probes(final=(step == study.steps)), train["batch_size"], args.device) + weightwatch(model, run_dir / "spectral" / f"step_{step:08d}.csv", cfg["weightwatcher"], args.seed, step) + if state_digest(model) != before: + raise RuntimeError("evaluation/diagnostics changed training weights") + row = {"step": step, "phase": "A" if step <= study.boundary else "B", "last_train_loss": loss, + "input_tokens": step * tokens_per_step, "model_sha256": before, "audit": audit} + with metrics_path.open("a") as f: + f.write(json.dumps(row, allow_nan=False) + "\n"); f.flush(); os.fsync(f.fileno()) + state = {"step": step, "fingerprint": fingerprint, "model": model.state_dict(), + "optimizers": [h.optimizer.state_dict() for h in handles], "exposures": dict(exposures), + "torch_rng": torch.get_rng_state(), "device_rng": None} + if args.device == "cuda": + state["device_rng"] = torch.cuda.get_rng_state_all() + elif args.device == "mps": + state["device_rng"] = torch.mps.get_rng_state() + temp = latest.with_suffix(".tmp") + torch.save(state, temp); os.replace(temp, latest) + if step in permanent or step == study.boundary: + dest = run_dir / f"model_step_{step:08d}.pt" + torch.save({"model": state["model"], "step": step, "fingerprint": fingerprint}, dest.with_suffix(".tmp")) + os.replace(dest.with_suffix(".tmp"), dest) + counts = {r.eid: exposures.get(r.eid, 0) for r in study.audit} + counts.update(exposures) + atomic_json(run_dir / "exposures.json", counts) + print(f"{args.condition} {args.optimizer} seed={args.seed} step={step}/{study.steps} loss={loss} WW=clip_xmax", flush=True) + + if first == 0 and not latest.exists(): + checkpoint(0, None) + model.train() + for step in range(first, study.steps): + zero_grad(handles) + set_learning_rates(handles, update_index=step, total_steps=schedule_steps, warmup_steps=warmup) + records = study.sample(step) + total_loss = 0.0 + for i in range(0, effective_batch, train["batch_size"]): + batch = records[i:i + train["batch_size"]] + x, y = pack(batch, source["model"]["block_size"]) + _, loss = model(torch.as_tensor(x, device=args.device), torch.as_tensor(y, device=args.device)) + if not torch.isfinite(loss): + raise FloatingPointError(f"nonfinite loss at step {step}") + (loss / train["grad_accum_steps"]).backward() + total_loss += float(loss.detach()) / train["grad_accum_steps"] + norm = torch.nn.utils.clip_grad_norm_(model.parameters(), train["grad_clip"]) + if not torch.isfinite(norm): + raise FloatingPointError(f"nonfinite gradient at step {step}") + optimizer_step(handles) + exposures.update(r.eid for r in records) + done = step + 1 + if done % 25 == 0: + print(f"step={done} train_loss={total_loss:.6f}", flush=True) + if done in permanent or done in cfg["early_eval_steps"] or done == study.boundary or done % study.stage["eval_every"] == 0 or done == study.steps: + checkpoint(done, total_loss) + atomic_json(run_dir / "complete.json", {"fingerprint": fingerprint, "steps": study.steps, "status": "complete"}) + print(f"Complete: {run_dir}") + + +def main() -> None: + cfg = json.loads(CONFIG.read_text()) + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="command", required=True) + plan = sub.add_parser("plan") + plan.add_argument("--stage", choices=cfg["stages"], default="pilot") + train = sub.add_parser("run") + train.add_argument("--stage", choices=cfg["stages"], default="pilot") + train.add_argument("--condition", choices=cfg["conditions"], required=True) + train.add_argument("--optimizer", choices=cfg["optimizers"], required=True) + train.add_argument("--seed", type=int, choices=cfg["seeds"], default=cfg["seeds"][0]) + train.add_argument("--device", choices=["cpu", "mps", "cuda"], required=True) + train.add_argument("--root") + train.add_argument("--recipe", choices=["repository", "shared_aux_decay"], default="repository") + train.add_argument("--resume", action="store_true") + monitor = sub.add_parser("monitor") + monitor.add_argument("run_dir", type=Path) + args = parser.parse_args() + if args.command == "run": + run(args, cfg) + elif args.command == "plan": + seeds = cfg["seeds"] if args.stage == "full" else cfg["seeds"][:1] + print(f"# {len(seeds)*len(cfg['optimizers'])*len(cfg['conditions'])} runs; {cfg['stages'][args.stage]['steps']} steps/run; synthetic suffix loss, NOT FineWeb NLL.") + for condition in cfg["conditions"]: + for seed in seeds: + for optimizer in cfg["optimizers"]: + print(f"python run.py run --stage {args.stage} --condition {condition} --optimizer {optimizer} --seed {seed} --device mps --resume") + else: + rows = read_rows(args.run_dir / "metrics.jsonl") + if not rows: + raise SystemExit("No completed audit yet") + latest = rows[-1] + print(json.dumps({"step": latest["step"], "audit": {k: v["mean"] for k, v in latest["audit"].items()}}, indent=2)) + print((args.run_dir / "spectral" / f"step_{latest['step']:08d}.csv").read_text()) + + +if __name__ == "__main__": + main() diff --git a/baseline/experiments/nanogpt_memorization/run_full.sh b/baseline/experiments/nanogpt_memorization/run_full.sh new file mode 100644 index 0000000..e25ae0f --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/run_full.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +# Run all configured smoke checks, then the complete five-seed campaign. +# Usage: bash run_full.sh [mps|cuda|cpu] [/tmp/dedicated-root] [repository|shared_aux_decay] [--plan] +set -euo pipefail + +DEVICE="${1:-mps}" +ROOT="${2:-/tmp/rg-nanogpt-memorization-20260914}" +RECIPE="${3:-repository}" +MODE="${4:-run}" +PYTHON="${PYTHON:-python}" +case "$DEVICE" in mps|cuda|cpu) ;; *) echo "Invalid device: $DEVICE" >&2; exit 2;; esac +case "$RECIPE" in repository|shared_aux_decay) ;; *) echo "Invalid recipe: $RECIPE" >&2; exit 2;; esac +case "$MODE" in run|--plan) ;; *) echo "Fourth argument must be --plan when provided" >&2; exit 2;; esac +if [ "$#" -gt 4 ]; then echo "Too many arguments" >&2; exit 2; fi +cd "$(CDPATH= cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" + +# Resolve Python before the runner redirects HOME and other cache locations. +PYTHON="$("$PYTHON" -c 'import sys; print(sys.executable)')" +ROOT="$("$PYTHON" - "$ROOT" <<'PY' +from pathlib import Path +import sys +root = Path(sys.argv[1]).expanduser().resolve() +if not any(root.is_relative_to(Path(base)) and root != Path(base) + for base in ('/tmp', '/private/tmp')): + raise SystemExit('Use a dedicated absolute output directory beneath /tmp.') +print(root) +PY +)" + +make_plan() { + "$PYTHON" - "$DEVICE" "$ROOT" "$RECIPE" <<'PY' +import json +from pathlib import Path +import shlex +import sys + +device, root, recipe = sys.argv[1:] +cfg = json.loads(Path('configs/suite.json').read_text()) +full_count = len(cfg['conditions']) * len(cfg['optimizers']) * len(cfg['seeds']) +print('#!/usr/bin/env bash\nset -euo pipefail') +print(shlex.join([sys.executable, '-m', 'pytest', '-q', 'tests'])) +print('echo ' + shlex.quote(f'Full campaign: {full_count} runs, ' + f"{cfg['stages']['full']['steps']:,} updates/run; recipe={recipe}; device={device}.")) +for stage, seeds in [('smoke', cfg['seeds'][:1]), ('full', cfg['seeds'])]: + print('echo ' + shlex.quote(f'Starting {stage}; an error stops the campaign.')) + for condition in cfg['conditions']: + for seed in seeds: + for optimizer in cfg['optimizers']: + label = f'{stage}/{recipe}/{condition}/{optimizer}/seed_{seed}' + print('echo ' + shlex.quote('RUN ' + label)) + print(shlex.join([sys.executable, '-u', 'run.py', 'run', + '--stage', stage, '--condition', condition, + '--optimizer', optimizer, '--seed', str(seed), + '--device', device, '--root', root, '--recipe', recipe, '--resume'])) +print('echo ' + shlex.quote(f'All {full_count} full runs returned successfully.')) +PY +} + +if [ "$MODE" = --plan ]; then + make_plan + exit 0 +fi +mkdir -p "$ROOT/logs" +# Each invocation gets its own plan and append-only log; rerunning resumes results. +STAMP="$(date -u +%Y%m%dT%H%M%SZ)_$$" +PLAN="$ROOT/logs/full_${RECIPE}_${DEVICE}_${STAMP}.sh" +LOG="$ROOT/logs/full_${RECIPE}_${DEVICE}_${STAMP}.log" +make_plan > "$PLAN" +printf 'Plan: %s\nLog: %s\nResults: %s/full/%s\n' "$PLAN" "$LOG" "$ROOT" "$RECIPE" +echo 'Keep this process in the foreground. Re-run the same command to resume.' +# Pipe failure is not hidden by tee; never continue to the next run after an error. +bash "$PLAN" 2>&1 | tee -a "$LOG" diff --git a/baseline/experiments/nanogpt_memorization/tests/test_suite.py b/baseline/experiments/nanogpt_memorization/tests/test_suite.py new file mode 100644 index 0000000..ddc0ca8 --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/tests/test_suite.py @@ -0,0 +1,192 @@ +import importlib.util +import json +from collections import Counter +from pathlib import Path +import sys +from types import SimpleNamespace + +import numpy as np +import pytest + +HERE = Path(__file__).resolve().parents[1] +spec = importlib.util.spec_from_file_location("memorization_run", HERE / "run.py") +m = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = m +spec.loader.exec_module(m) +CFG = json.loads((HERE / "configs/suite.json").read_text()) + + +def study(condition, stage="smoke", seed=1337): + return m.Study(CFG, condition, stage, seed, batch=32) + + +def test_pack_only_scores_suffix_and_has_correct_shift(): + r = m.Record("a", "a", (7, 8), (9, 10)) + x, y = m.pack([r], 8) + assert x.tolist() == [[0, 0, 0, 0, 0, 7, 8, 9]] + assert y.tolist() == [[-100] * 6 + [9, 10]] + with pytest.raises(ValueError): + m.pack([r], 2) + + +def test_exact_lifetime_doses_and_no_post_withdrawal_presentations(): + s = study("verbatim") + observed = Counter(r.eid for step in range(s.steps) for r in s.sample(step)) + for r in s.audit: + assert observed[r.eid] == int(r.group.split("_")[-1]) + assert all(r.eid == "background" for r in s.sample(s.boundary)) + + +def test_counterfactual_replaces_only_injection_slots(): + present, absent = study("verbatim"), study("verbatim_absent") + assert present.audit == absent.audit + assert present.injections == absent.injections + for step in range(present.steps): + for i, (a, b) in enumerate(zip(present.sample(step), absent.sample(step))): + if step * 32 + i not in present.injections: + assert a == b + else: + assert a.eid.startswith("canary") and b.eid == "background" + assert not any(r.eid.startswith("canary") for r in absent.sample(0)) + + +def test_sampling_reproducible_and_not_consumed_by_probes(): + a, b = study("associations"), study("associations") + expected = a.sample(1) + a.probes(final=True); a.sample(0) + assert a.sample(1) == expected == b.sample(1) + assert study("associations", seed=2027).sample(1) != expected + + +def test_rule_split_is_pair_disjoint_and_exhaustive(): + s = study("rule_half_noise") + train = {r.eid for r in s.train} + val = {r.eid for r in s.audit if r.group == "validation_rule"} + test = {r.eid for r in s.audit if r.group == "test_rule"} + assert not train & val and not train & test and not val & test + assert len(train | val | test) == CFG["modulus"] ** 2 + assert not any(r.group == "test_rule" for r in s.probes()) + assert any(r.group == "test_rule" for r in s.probes(final=True)) + + +def test_label_noise_is_frozen_and_holdouts_unchanged(): + clean, half, full = [study(c) for c in ("rule_clean", "rule_half_noise", "rule_random")] + assert [r.eid for r in clean.train] == [r.eid for r in full.train] + assert sum(r.group.endswith("randomized") for r in half.train) == round(len(half.train)/2) + assert all(r.group.endswith("randomized") for r in full.train) + assert [r for r in clean.audit if r.group.endswith("_rule")] == [r for r in full.audit if r.group.endswith("_rule")] + assert any(a.target != b.target for a, b in zip(clean.train, full.train)) + + +def test_association_controls_are_disjoint_and_template_is_unseen(): + s = study("associations") + seen = {r.prompt[2:-1] for r in s.train} + unseen = {r.prompt[2:-1] for r in s.audit if r.group == "unseen_key_control"} + assert len(seen) == len(unseen) == CFG["associations"] and not seen & unseen + assert {r.prompt[1] for r in s.sample(0)} == {4, 5} + assert all(r.prompt[1] == 6 for r in s.audit if r.group == "seen_key_new_template") + + +def test_forgetting_arms_have_identical_acquisition_but_different_interference(): + a, b = study("forgetting_disjoint"), study("forgetting_conflict") + assert a.train == b.train and a.sample(0) == b.sample(0) + assert {r.prompt for r in a.train}.isdisjoint({r.prompt for r in a.second}) + assert [r.prompt for r in b.train] == [r.prompt for r in b.second] + assert all(x.target != y.target for x, y in zip(b.train, b.second)) + assert all(r.eid.startswith("B_") for r in a.sample(a.boundary)) + + +def test_prefix_ablation_only_shortens_prompt_not_suffix(): + s = study("verbatim") + probes = s.probes(final=True) + for record in s.audit: + rows = [r for r in probes if r.eid == record.eid] + assert {len(r.prompt) for r in rows} == set(CFG["prefix_lengths_final"]) + assert all(r.target == record.target for r in rows) + + +def test_shared_aux_control_matches_decay_not_just_coefficient(): + src = {"optimizer_profiles": { + "adamw": {"learning_rate": 6e-4, "min_learning_rate": 6e-5, "weight_decay": .1, + "warmup_fraction": .01, "lr_schedule_epochs": 1, "beta1": .9, "beta2": .95, "epsilon": 1e-8}, + "muon": {"matrix_learning_rate": .02, "matrix_min_learning_rate": .002}}} + mu = m.resolve_profile(src, "muon", "shared_aux_decay") + ad = m.resolve_profile(src, "adamw", "shared_aux_decay") + assert mu["aux_learning_rate"] == ad["learning_rate"] + assert mu["warmup_fraction"] == ad["warmup_fraction"] + assert mu["matrix_learning_rate"] * mu["matrix_weight_decay"] == pytest.approx(ad["learning_rate"] * ad["weight_decay"]) + assert mu["matrix_min_learning_rate"] * mu["matrix_weight_decay"] == pytest.approx(ad["min_learning_rate"] * ad["weight_decay"]) + assert "aux_learning_rate" not in src["optimizer_profiles"]["muon"] + + +def test_interrupted_jsonl_only_allows_last_line(tmp_path): + path = tmp_path / "metrics.jsonl" + path.write_text('{"step": 1}\n{"step":') + assert m.read_rows(path) == [{"step": 1}] + path.write_text('{broken}\n{"step": 1}\n') + with pytest.raises(json.JSONDecodeError): + m.read_rows(path) + + +def test_evaluation_distinguishes_exact_and_token_recall(): + torch = pytest.importorskip("torch") + class Uniform: + cfg = SimpleNamespace(block_size=8) + training = True + def train(self, value=True): self.training = value + def eval(self): self.training = False + def __call__(self, x): return torch.zeros((*x.shape, 16)), None + def generate_greedy(self, prompt, n): + return torch.cat([prompt, torch.ones((len(prompt), n), dtype=torch.long)], dim=1) + model = Uniform() + out = m.evaluate(model, [m.Record("a", "g", (3,), (1, 1)), m.Record("b", "g", (3,), (2, 1))], 2, "cpu") + assert out["g"]["mean"]["nll"] == pytest.approx(np.log(16)) + assert out["g"]["mean"]["exact_match"] == .5 + assert out["g"]["mean"]["continuation_token_accuracy"] == .75 + assert model.training + + +def test_all_conditions_produce_valid_bounded_batches(): + for condition in CFG["conditions"]: + s = study(condition) + for step in range(s.steps): + x, y = m.pack(s.sample(step), 256) + assert x.shape == y.shape == (32, 256) + assert (x >= 0).all() and (x < 50257).all() + assert ((y == -100) | ((y >= 0) & (y < 50257))).all() + assert (y != -100).any(axis=1).all() + + +def test_weightwatcher_call_is_clipped_and_does_not_change_rng_or_weights(monkeypatch, tmp_path): + """Mocked API contract test, NOT a test of actual WW numerical fits.""" + torch = pytest.importorskip("torch") + pd = pytest.importorskip("pandas") + from types import ModuleType + model = torch.nn.Linear(32, 32, bias=False) + before = model.weight.detach().clone() + module = ModuleType("rg_nanogpt_one_head.model") + module.transformer_matrix_items = lambda model: [("L00_W_Q", "W_Q", 0, model.weight)] + monkeypatch.setitem(sys.modules, "rg_nanogpt_one_head.model", module) + fake = ModuleType("weightwatcher") + calls = [] + class Watcher: + def __init__(self, model): self.model = model + def analyze(self, **kwargs): + calls.append(kwargs) + np.random.rand(10); torch.rand(10) + return pd.DataFrame([dict(longname="L00_W_Q", alpha=2.1, raw_alpha=2.4, D=.1, + rand_distance=.2, ERG_gap=1, num_traps=0, num_pl_spikes=24)]) + fake.WeightWatcher = Watcher + monkeypatch.setitem(sys.modules, "weightwatcher", fake) + monkeypatch.setattr(m.importlib.metadata, "version", lambda name: "0.7.7") + np_state = np.random.get_state() + torch_state = torch.get_rng_state().clone() + m.weightwatch(model, tmp_path / "ww.csv", CFG["weightwatcher"], 1337, 0) + assert torch.equal(before, model.weight) + assert torch.equal(torch_state, torch.get_rng_state()) + assert np.array_equal(np_state[1], np.random.get_state()[1]) + assert calls == [{k:v for k,v in CFG["weightwatcher"].items() if k != "version"}] + table = pd.read_csv(tmp_path / "ww.csv") + assert table["alpha_clip_xmax"].iloc[0] == 2.1 + assert table["alpha_raw"].iloc[0] == 2.4 + assert table["matrix_name"].iloc[0] == "L00_W_Q"