diff options
| -rw-r--r-- | design/preregistration.md | 52 | ||||
| -rw-r--r-- | design/reviews/codex-review.md | 68 | ||||
| -rw-r--r-- | src/config.py | 18 | ||||
| -rw-r--r-- | src/data.py | 28 | ||||
| -rw-r--r-- | src/eval.py | 171 | ||||
| -rw-r--r-- | src/model_api.py | 7 | ||||
| -rw-r--r-- | src/models/rnn.py | 8 | ||||
| -rw-r--r-- | src/train.py | 20 | ||||
| -rw-r--r-- | tests/test_codex_fixes.py | 104 | ||||
| -rw-r--r-- | tests/test_data.py | 9 |
10 files changed, 402 insertions, 83 deletions
diff --git a/design/preregistration.md b/design/preregistration.md index f800461..30f731f 100644 --- a/design/preregistration.md +++ b/design/preregistration.md @@ -92,3 +92,55 @@ The interpretation matrix (O/H/P codes) above is NOT amended; only setup details - Giannou et al. (2023), "Looped Transformers as Programmable Computers", arXiv:2301.13196 - Xu et al. (ICLR 2020), "What Can Neural Networks Reason About?", arXiv:1905.13211 - Xu et al. (2021), "How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks", arXiv:2009.11848 + +--- + +## Addendum 2 (2026-08-14, pre-launch — operationalization + correctness fixes) + +Trigger: OpenAI Codex code review (experiment repo `design/reviews/codex-review.md`). +The interpretation matrix is unchanged; this addendum (a) locks the code thresholds that the +prose left unquantified, and (b) records setup correctness fixes. Everything below is +pre-launch and pre-run. + +### Operationalization of O/H/P codes (thresholds now locked) + +- **O1:** train EM ≥ 0.95 for ≥ 10 CONSECUTIVE evals (window start `sat_start`), THEN val EM + reaches ≥ 0.9 at eval `hi` with `hi ≥ sat_start + 10` (strictly after the window), where the + last eval with val ≤ 0.2 (`lo`) satisfies `hi - lo ≤ 5`. +- **O2:** train EM ≥ 0.95 sustained AND val EM never reaches 0.9 AND final val EM ≤ 0.3. +- **O3:** train EM ≥ 0.95 sustained AND val EM reaches 0.9 but not via the O1 pattern. +- **O4:** train EM never sustained at ≥ 0.95 for 10 evals (setup/optimization failure). +- **O-PARTIAL** (new — the original matrix left this region unspecified): train saturated, + val never reaches 0.9, final val EM in (0.3, 0.9). Reported as measurements, interpreted + cautiously as partial in-range generalization; NOT retrofitted into O1/O2/O3. +- **H1:** mean steps at run end ≤ min_steps + 0.5 (collapse to the floor — see fix 2: with + min_steps = 2 the floor is 2, so literal "collapse to 1" is impossible by construction). +- **H2:** mean steps ≥ K − 0.5 (never learned to halt). +- **H3:** min_steps + 0.5 < mean < K − 0.5 AND Pearson ρ(gap-to-next-prime, steps) ≥ +0.3 + (positive correlation = larger gap consumes more compute steps). +- **H4:** intermediate but ρ < 0.3 (halt signal noisy/unused). +- **P1:** ≥ 3 of {121, 143, 169, 187} wrong AND ≤ 6 total errors AND all errors within the flagged set. +- **P2:** scattered errors not matching P1/P3/P4. +- **P3:** probe accuracy ≥ 85% (checked after P4). +- **P4:** > 50% of trivial-composite inputs (even or multiple of 5) wrong. + +### Correctness fixes (codex review) + +1. **Layout invariance (BLOCKER).** Inputs are now LEFT-padded to a fixed global length + (3 digits for [2,100]) in EVERY context — training batches, eval, greedy decoding; outputs + right-padded to the global length (4 incl. EOS). Previously batch-max padding made an + example's representation depend on its batchmates (RNN pad steps transformed the state; + transformer logit positions misaligned and absolute position embeddings shifted). + RNN pad steps are now exact no-ops. Regression tests: mixed-batch vs singleton logit + invariance for both models. +2. **min_steps off-by-one.** Halting probs forced to 0 only for the first `min_steps − 1` + steps (was: first `min_steps`), so the earliest halt is step `min_steps` — matching + "execute at least min_steps steps". +3. **integers vocab EOS alias.** Integers-mode vocab is now `next_prime(range_end) + 2`, so + the value token 101 (= next_prime(100)) is NOT aliased with EOS. Boundary test added. +4. **Rerun protection.** train.py refuses to run if `metrics.csv` exists (append + header + would corrupt provenance). Fresh `--out_dir` required per run. +5. **Checkpoint-selection honesty.** eval reports BOTH `best.pt` (val-selected — flagged as + selection-holed) and `last.pt` (unselected); probe + halting analyses use `last.pt`. +6. **Dead config removed:** `halt_eps` (was never read). +7. **`run_meta.json`:** python/torch/numpy versions, device, thread count recorded per run. diff --git a/design/reviews/codex-review.md b/design/reviews/codex-review.md new file mode 100644 index 0000000..55ca55d --- /dev/null +++ b/design/reviews/codex-review.md @@ -0,0 +1,68 @@ +# Hostile reproducibility review — 2026-08-14 + +Scope reviewed: `src/`, `tests/`, and `design/preregistration.md`. This is a code/design audit only; no training artifacts were assumed valid. `pytest` could not be run because it is not installed in the environment. + +## Findings + +### BLOCKER — Variable-length batches change the input seen by both models + +`make_batch` right-pads inputs to the longest input in the batch. In the tied RNN, `_encode` iterates across *all* padded positions and calls `_cell_step` even when `mask` zeroes the embedding/position contribution ([`src/models/rnn.py`](../../src/models/rnn.py), `_encode`). Thus a one-digit input receives additional tied transformations whenever it is batched with a two- or three-digit input. At evaluation and greedy decoding, every example is passed as a singleton, so it receives no such transformations. The same numerical input therefore has a batch-composition-dependent representation during training and a different representation at evaluation. + +The transformer has the corresponding, and more severe, alignment error ([`src/models/transformer.py`](../../src/models/transformer.py), `forward`): logits always start at `T_in - 1`. For a shorter example in a mixed-length batch, this is a **padded input position**, not that example's final real input digit. Output positions are also shifted by the batch maximum input length, so their learned absolute positional embeddings differ from singleton evaluation. `key_padding_mask` prevents padding from being a key, but does not move the prediction position or remove its residual/query representation. + +This silently makes both reported training behavior and validation/probe behavior depend on batching, and invalidates comparisons. Mask recurrent state updates by each sequence length and gather each transformer's first prediction at `length(x)-1` (or use a packing/layout that places each output immediately after its own input); train and evaluate with the identical layout. Add mixed-vs-singleton invariance tests for logits. + +### BLOCKER — The implementation assigns preregistered O/H/P labels using unregistered and incorrect rules + +`src/eval.py` claims to implement the locked codes but materially changes them. + +* **O1:** `grokking_signature` permits the ten-evaluation train-saturation window to occur anywhere, including after the validation transition. The preregistration requires train EM to be sustained first, then the transition. It also classifies an early `val >= .9` result as O1 based on an unrelated earlier low validation point. +* **O2:** its stated condition is final validation EM `<= .3`; the code labels *any* run with `max(train) >= .95` and no `val >= .9` as O2, including a final validation EM of .31--.89. That directly contradicts O2 and hides an unclassified/ambiguous outcome. +* **H1--H4:** `halting_report` evaluates the selected checkpoint rather than the state “by run end,” declares H1 at mean `<= min_steps + .5` (default `<=2.5`) rather than collapse to 1, and declares H3 based solely on `abs(correlation) >= .3`. It neither checks the registered intermediate 2--18 range nor evolution during training; a negative correlation is treated as desired evidence. H2 uses an unregistered `K-.5` threshold. The logged scalar batch mean cannot establish the required per-input evolution either. +* **P1--P4:** the code invents thresholds absent from the locked matrix (`acc >= .85`, `easy_misses >=5`, and `<=6` errors). In particular P4 requires failure across the probe including easy cases, but five easy misses can be called P4 even at 95% accuracy; P3 is tested first, so it can also override that condition. The P2 catch-all does not test scattered/uniform composite errors. These labels are post-hoc classifications, not the preregistered outcomes. + +Do not publish any O/H/P code from this evaluator. Either make each condition faithfully operational with thresholds added in a preregistration amendment before runs, or emit measurements plus `unclassified` where the prose does not supply a testable boundary. Preserve checkpoint/step provenance in every measurement. + +### MAJOR — ACT `min_steps` is off by one, making the declared H1 condition impossible + +In `_run_compute`, loop index `t=0` represents compute step 1, but halting probability is zeroed when `t < cfg.min_steps`. With the default `min_steps=2`, probabilities for steps 1 *and 2* are forced to zero, so the earliest non-remainder halt is step 3. The returned expected steps can never collapse to 1 (or even below 3), while the preregistration calls H1 “collapse to 1.” The test only asserts `steps >= min_steps`, which does not expose this error. + +If `min_steps` means “execute at least two steps,” zero only the first `min_steps - 1` probabilities. Then align H1/H3 code definitions and tests with the chosen indexing convention. + +### MAJOR — Validation is repeatedly used for checkpoint selection and then reported as the primary held-out result + +There is no train/validation input overlap: `get_splits` uses one seeded shuffle and disjoint slices. Repeated next-prime *labels* across distinct inputs are inherent to the task, not input leakage. However, training evaluates the validation set every 200 updates, stops on it, and saves `best.pt` whenever its validation EM improves ([`src/train.py`](../../src/train.py), `_eval_pass`). `src/eval.py` defaults to that val-selected checkpoint and reports its validation EM as the primary result and its probe outcome. The same holdout therefore controls stopping/checkpoint selection and supplies the headline estimate; it is no longer an unbiased final estimate after hundreds of looks. + +The early-stopping rule is preregistered, but the report must call this a validation-selected result, not an untouched held-out test result. For confirmatory performance, freeze a predeclared checkpoint rule (for example, the first qualifying stop) and use an untouched test split or independent seeds/replications; do not select the maximum observed validation EM. + +### MAJOR — Rerunning a seed in place corrupts `metrics.csv` and breaks analysis provenance + +The output path is deterministic (`runs/<model>/seed<seed>`), but `train.py` opens `metrics.csv` in append mode and unconditionally writes a header on each invocation. A second invocation interleaves another CSV header among old rows; `grokking_signature` then attempts `float("train_em")` and fails. If the new run has no improvement, a prior `best.pt` can also remain in place. Config JSON, metrics, and checkpoints can consequently originate from different invocations. + +Require an empty/new run directory, or create a unique run ID and write a manifest atomically. Refuse to resume without explicit resume semantics that restore model, optimizer, step, and all RNG states. + +### MAJOR — Determinism is only partial and not documented as CPU-only + +`set_seed` seeds Python, NumPy, and PyTorch CPU RNGs, and the data shuffle has its own seeded `Random`; this is a good start. But it neither records/enforces the PyTorch version and threading environment nor enables deterministic algorithms or CUDA/cuDNN deterministic settings. Future device placement, GPU use, or nondeterministic kernels will therefore make same-seed runs diverge without warning. There is also no resume checkpoint for optimizer/RNG state. + +Record environment/version/device metadata, explicitly select a device, and enforce deterministic settings appropriate to it. Test two fresh, isolated same-seed runs for byte-equivalent metric rows (or document a bounded tolerance). Keep each run in a fresh directory as above. + +### MAJOR — `integers` vocabulary aliases EOS with a valid target and is not range-safe + +Although the preregistered arm uses digit tokens, the supported `vocab_mode="integers"` is broken. At the default range end, token 101 is both the valid target `next_prime(100)` and `eos_id` (`vocab-1`); larger ranges can have a next prime beyond `range_end+1`, causing an embedding/target index failure. `decode_tokens` in integers mode ignores EOS entirely, which hides the alias. The existing test checks only output shape. + +Reserve separate BOS/EOS/PAD IDs outside the full input/target integer domain, derive that domain from the actual maximum target, and add boundary tests. Until then, reject integer mode rather than silently producing invalid semantics. + +### MINOR — ACT aggregation itself is normalized, but `halt_eps` is dead configuration + +The ACT aggregation is mathematically normalized: `w_t = remaining * p_t`, `remaining *= 1-p_t`, followed by the final remainder on `h_list[-1]`; returned expected steps and its mean penalty are differentiable. Penalizing `mean(halt_steps)` with the scheduled `lambda` is coherent with that distributional definition and correctly disabled for fixed-K runs. There is no normalization bug here. + +However, `Config.halt_eps` is documented as an ACT cumulative threshold but is never read. The code always computes all K states and uses a soft remainder, contrary to that configuration's description. Remove the unused setting or implement and test the specified behavior; do not imply a threshold is active in run configs. + +### MINOR — Pad-target clamping is safe only because masking follows it, but lacks a guard/test + +`y_safe = y.clamp(max=vocab-1)` maps pad ID to EOS solely to make cross-entropy indexable, then multiplication by `y_mask` removes those terms. With the current `make_batch`, this produces correct token loss and avoids pad embeddings in the output head. It is fragile: any future mask mismatch would silently train padded positions toward EOS. Assert that every masked-in target is `< vocab` and every masked-out target is exactly `pad_id`, or use an ignored-index loss. Add a test that perturbing logits at padded target positions cannot change loss or gradients. + +### NIT — Tests do not exercise the reviewed failure modes + +The suite checks shapes, a broad halting range, and finite backward passes, but not ACT mass summing to one, penalty gradients, mixed-length versus singleton behavior, evaluator-code boundary cases, fresh-run handling, or deterministic repeatability. These omissions allowed the issues above to pass. Also, `log_examples` orders with `len(str(x))` where `x` is a token list; it is deterministic but obscures the intended numeric ordering. diff --git a/src/config.py b/src/config.py index c0532bc..0b161e7 100644 --- a/src/config.py +++ b/src/config.py @@ -17,11 +17,10 @@ class Config: d_model: int = 128 max_steps: int = 20 # K tied iterations (RNN) halting: bool = True # False -> fixed-K ablation - halt_eps: float = 0.05 # ACT cumulative threshold (documented; K small so no early break) halt_penalty: float = 0.01 # lambda on mean steps halt_warmup_steps: int = 1000 # penalty = 0 before this (anti-collapse) halt_ramp_end_steps: int = 5000 # penalty ramps linearly 0 -> halt_penalty between warmup and here - min_steps: int = 2 # ACT: halt prob forced to 0 before this step + min_steps: int = 2 # ACT: halt prob forced to 0 for the first min_steps-1 steps n_layers: int = 2 n_heads: int = 4 # training @@ -38,10 +37,19 @@ class Config: @property def vocab(self) -> int: - """Token count. Integers mode: 0..range_end+1 (next prime can exceed range_end).""" + """Token count. Integers mode: value tokens 0..next_prime(range_end), plus EOS (+1) and pad (+1).""" if self.vocab_mode == "integers": - return self.range_end + 2 - return 11 # digits 0-9 + EOS + # inline mini-sieve: next prime above range_end bounds the target domain + limit = self.range_end + 100 + is_prime = [True] * (limit + 1) + is_prime[0] = is_prime[1] = False + for p in range(2, int(limit ** 0.5) + 1): + if is_prime[p]: + for m in range(p * p, limit + 1, p): + is_prime[m] = False + np_ = next(i for i in range(self.range_end + 1, limit + 1) if is_prime[i]) + return np_ + 2 # values 0..np_, EOS=np_+1, pad=np_+2 + return 11 # digits 0-9 + EOS @property def eos_id(self) -> int: diff --git a/src/data.py b/src/data.py index 35682d3..e6ea57b 100644 --- a/src/data.py +++ b/src/data.py @@ -63,15 +63,33 @@ def build_examples(inputs: list[int], cfg: Config) -> list[tuple[list[int], list return out +def _global_lengths(cfg: Config) -> tuple[int, int]: + """(in_max, out_max): fixed global lengths so batch layout == singleton layout (codex BLOCKER fix).""" + primes = sieve_primes(cfg.range_end + 100) + max_target = next_prime(cfg.range_end, primes) + if cfg.vocab_mode == "integers": + return 1, 2 # [value], [value, EOS] + return len(str(cfg.range_end)), len(str(max_target)) + 1 # digits + EOS + + +def pad_inputs(x: torch.Tensor, cfg: Config) -> torch.Tensor: + """LEFT-pad inputs to the global in_max so absolute positions are layout-invariant.""" + in_max, _ = _global_lengths(cfg) + if x.shape[1] < in_max: + pad = torch.full((x.shape[0], in_max - x.shape[1]), cfg.pad_id, dtype=x.dtype, device=x.device) + x = torch.cat([pad, x], dim=1) + return x + + def make_batch(examples, cfg: Config) -> dict[str, torch.Tensor]: + """Fixed global layout (not batch-max): x left-padded to in_max, y right-padded to out_max.""" xs, ys = zip(*examples) - T_in = max(len(x) for x in xs) - T_out = max(len(y) for y in ys) + in_max, out_max = _global_lengths(cfg) B = len(examples) - x = torch.full((B, T_in), cfg.pad_id, dtype=torch.long) - y = torch.full((B, T_out), cfg.pad_id, dtype=torch.long) + x = torch.full((B, in_max), cfg.pad_id, dtype=torch.long) + y = torch.full((B, out_max), cfg.pad_id, dtype=torch.long) for i, (xi, yi) in enumerate(examples): - x[i, : len(xi)] = torch.tensor(xi, dtype=torch.long) + x[i, in_max - len(xi):] = torch.tensor(xi, dtype=torch.long) y[i, : len(yi)] = torch.tensor(yi, dtype=torch.long) # teacher-forced decoder input: BOS(=EOS reuse) then shifted y y_in = torch.cat([torch.full((B, 1), cfg.eos_id, dtype=torch.long), y[:, :-1]], dim=1) diff --git a/src/eval.py b/src/eval.py index 4485523..fdc2dec 100644 --- a/src/eval.py +++ b/src/eval.py @@ -1,7 +1,9 @@ """Post-run analysis: final metrics, [101,200] probe with sieve diagnostic, grokking signature. -Interpretation codes are locked in design/preregistration.md — this module only MEASURES -and classifies against those definitions (O1-O4, H1-H4, P1-P4). +Interpretation codes are locked in design/preregistration.md (with operationalization +thresholds in Addendum 2). This module MEASURES and classifies against those definitions. +Where the prereg prose supplies no testable boundary, the code emits the measurements +plus "unclassified" rather than inventing a label. """ import argparse import csv @@ -23,38 +25,44 @@ def probe_report(model, cfg: Config, lo: int = 101, hi: int = 200) -> dict: primes = sieve_primes(hi + 200) correct = 0 errors = [] - easy_misses = 0 + easy_total = 0 + easy_wrong = 0 for n in range(lo, hi + 1): x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0) gen = greedy_decode(model, x, cfg)[0].tolist() pred = decode_tokens(gen, cfg) target = next_prime(n, primes) + is_easy = (n % 2 == 0) or (n % 5 == 0) # trivial composites (skip-evens / skip-5s) + if is_easy: + easy_total += 1 if pred == target: correct += 1 else: errors.append({"n": n, "target": target, "pred": pred}) - if n % 2 == 0 or n % 5 == 0: - easy_misses += 1 + if is_easy: + easy_wrong += 1 total = hi - lo + 1 acc = correct / total flagged = [e for e in errors if e["n"] in FLAGGED_COMPOSITES] - # P-code classification (see preregistration.md) - if acc >= 0.85: - code = "P3" - elif easy_misses >= 5: - code = "P4" - elif errors and all(e["n"] in FLAGGED_COMPOSITES for e in errors) and len(errors) <= 6: - code = "P1" + # classification per prereg + Addendum 2 operationalization; P4 checked first + if easy_total and easy_wrong / easy_total > 0.5: + code = "P4" # fails trivial evens/5-multiples -> pure memorization + elif acc >= 0.85: + code = "P3" # surprising success beyond expectation + elif len(flagged) >= 3 and len(errors) <= 6 and all(e["n"] in FLAGGED_COMPOSITES for e in errors): + code = "P1" # errors concentrated on composites needing divisors 11,13 -> learned sieve else: - code = "P2" + code = "P2" # scattered errors -> memorization / non-transferable heuristics return { "code": code, "acc": acc, "correct": correct, "total": total, - "errors": errors, "flagged_errors": flagged, "easy_misses": easy_misses, + "errors": errors, "flagged_errors": flagged, + "easy_total": easy_total, "easy_wrong": easy_wrong, + "easy_err_rate": (easy_wrong / easy_total) if easy_total else None, } def grokking_signature(metrics_path: str) -> dict: - """Classify the training curve against preregistered codes O1-O4.""" + """Classify the training curve against preregistered codes O1-O4 (Addendum 2 operationalization).""" with open(metrics_path) as fh: rows = list(csv.DictReader(fh)) if not rows: @@ -62,31 +70,41 @@ def grokking_signature(metrics_path: str) -> dict: train = [float(r["train_em"]) for r in rows] val = [float(r["val_em"]) for r in rows] n = len(rows) - saturated = any(all(t >= 0.95 for t in train[i:i + 10]) for i in range(n - 9)) if n >= 10 else False + # first index where train EM >= 0.95 for 10 CONSECUTIVE evals + sat_start = next((i for i in range(n - 9) if all(t >= 0.95 for t in train[i:i + 10])), None) hi = next((i for i, v in enumerate(val) if v >= 0.9), None) - trans = None - if hi is not None: - lo_cands = [i for i in range(hi) if val[i] <= 0.2] + # transition: last eval with val <= 0.2 strictly before hi, and after saturation window + lo = None + if hi is not None and sat_start is not None: + lo_cands = [i for i in range(sat_start + 10, hi) if val[i] <= 0.2] if lo_cands: - trans = hi - max(lo_cands) - if saturated and hi is not None and trans is not None and trans <= 5: - code = "O1" - elif max(train) >= 0.95 and hi is not None and (trans is None or trans > 5): - code = "O3" - elif max(train) >= 0.95 and hi is None: - code = "O2" + lo = max(lo_cands) + width = (hi - lo) if (hi is not None and lo is not None) else None + + max_train = max(train) + if max_train < 0.95: + code = "O4" # train never reached 0.95 -> setup/optimization failure + elif sat_start is None: + code = "O4" # reached 0.95 but never sustained 10 evals within budget + elif hi is not None and width is not None and width <= 5: + code = "O1" # sharp transition AFTER sustained train saturation + elif hi is not None: + code = "O3" # reached 0.9+ but not via the sharp O1 pattern (gradual) + elif val[-1] <= 0.3: + code = "O2" # train memorized, val stayed low else: - code = "O4" + code = "O-PARTIAL" # val ended in (0.3, 0.9) with no 0.9 reach — prereg has no boundary return { - "code": code, "train_saturated": saturated, "val_hi_eval_idx": hi, - "transition_width_evals": trans, "n_evals": n, + "code": code, "sat_start_eval": sat_start, "val_hi_eval_idx": hi, + "transition_width_evals": width, "n_evals": n, "final_train_em": train[-1], "final_val_em": val[-1], + "max_train_em": max_train, } @torch.no_grad() def halting_report(model, cfg: Config) -> dict: - """RNN halting structure: mean steps + correlation with gap-to-next-prime (H1-H4).""" + """RNN halting structure: mean steps at run end + correlation with gap-to-next-prime (H1-H4).""" primes = sieve_primes(300) gaps, steps = [], [] for n in range(cfg.range_start, cfg.range_end + 1): @@ -97,54 +115,79 @@ def halting_report(model, cfg: Config) -> dict: steps.append(float(s.mean())) mean = float(np.mean(steps)) rho = float(np.corrcoef(gaps, steps)[0, 1]) if len(set(gaps)) > 1 else 0.0 - if mean <= cfg.min_steps + 0.5: - code = "H1" - elif mean >= cfg.max_steps - 0.5: - code = "H2" - elif abs(rho) >= 0.3: - code = "H3" + lo, hi = cfg.min_steps + 0.5, cfg.max_steps - 0.5 + if mean <= lo: + code = "H1" # collapse to the min_steps floor + elif mean >= hi: + code = "H2" # pinned at K — never learned to halt + elif rho >= 0.3: + code = "H3" # intermediate + positive gap correlation -> computation budget else: - code = "H4" - return {"code": code, "mean_steps": mean, "corr_gap": rho, "min_steps": cfg.min_steps, "max_steps": cfg.max_steps} + code = "H4" # intermediate but no positive gap correlation — noisy/unused + return {"code": code, "mean_steps": mean, "corr_gap": rho, + "min_steps": cfg.min_steps, "max_steps": cfg.max_steps} + + +def _load(out_dir: str, ckpt: str, cfg: Config): + model = build_model(cfg) + model.load_state_dict(torch.load(os.path.join(out_dir, ckpt), map_location="cpu")) + model.eval() + return model def main() -> None: ap = argparse.ArgumentParser(description="prime-grokking eval") ap.add_argument("model") ap.add_argument("seed") - ap.add_argument("--ckpt", default="best.pt") a = ap.parse_args() out_dir = os.path.join("runs", a.model, f"seed{a.seed}") cfg = Config.load(os.path.join(out_dir, "config.json")) - model = build_model(cfg) - model.load_state_dict(torch.load(os.path.join(out_dir, a.ckpt), map_location="cpu")) - model.eval() _, val_in = get_splits(cfg) val_ex = build_examples(val_in, cfg) - tok, em, per = evaluate(model, val_ex, cfg) - - probe = probe_report(model, cfg) - sig = grokking_signature(os.path.join(out_dir, "metrics.csv")) - hlt = halting_report(model, cfg) if cfg.model == "rnn" else None - - results = { - "model": cfg.model, "seed": cfg.seed, "ckpt": a.ckpt, - "params": model.param_count(), - "val_token_acc": tok, "val_exact_match": em, - "per_example": [{"n": "".join(map(str, x)), "target": t, "pred": p, "ok": ok} - for x, t, p, ok in per], - "probe": probe, - "signature": sig, - "halting": hlt, - } + + report = {"model": cfg.model, "seed": cfg.seed, "params": None, "val_selected": {}, "final": {}} + for ckpt in ("best.pt", "last.pt"): + path = os.path.join(out_dir, ckpt) + if not os.path.exists(path): + continue + model = _load(out_dir, ckpt, cfg) + tok, em, per = evaluate(model, val_ex, cfg) + entry = { + "ckpt": ckpt, + "params": model.param_count(), + "val_token_acc": tok, + "val_exact_match": em, + "note": ("val-selected checkpoint (early stop / best-on-val per prereg — " + "val EM here is selection-holed, not an untouched test estimate" if ckpt == "best.pt" + else "final checkpoint, no val selection"), + } + if ckpt == "last.pt": + entry["probe"] = probe_report(model, cfg) + if cfg.model == "rnn": + entry["halting"] = halting_report(model, cfg) + entry["per_example"] = [{"n": "".join(map(str, x)), "target": t, "pred": p, "ok": ok} + for x, t, p, ok in per] + if ckpt == "best.pt": + report["val_selected"] = entry + else: + report["final"] = entry + + report["signature"] = grokking_signature(os.path.join(out_dir, "metrics.csv")) with open(os.path.join(out_dir, "results.json"), "w") as fh: - json.dump(results, fh, indent=2) - print(json.dumps({"val_token_acc": round(tok, 4), "val_em": round(em, 4), - "probe": probe["code"], "probe_acc": round(probe["acc"], 3), - "flagged_errors": [e["n"] for e in probe["flagged_errors"]], - "signature": sig["code"], "halting": hlt["code"] if hlt else None, - "results": os.path.join(out_dir, "results.json")}, indent=2)) + json.dump(report, fh, indent=2) + + f = report["final"] + p = f.get("probe", {}) + print(json.dumps({ + "val_em_best": round(report["val_selected"].get("val_exact_match", float("nan")), 4), + "val_em_last": round(f.get("val_exact_match", float("nan")), 4), + "probe_code": p.get("code"), "probe_acc": round(p.get("acc", float("nan")), 3), + "flagged_errors": [e["n"] for e in p.get("flagged_errors", [])], + "signature": report["signature"]["code"], + "halting": (f.get("halting") or {}).get("code"), + "results": os.path.join(out_dir, "results.json"), + }, indent=2)) if __name__ == "__main__": diff --git a/src/model_api.py b/src/model_api.py index 66c2ace..31ffe31 100644 --- a/src/model_api.py +++ b/src/model_api.py @@ -3,6 +3,7 @@ import torch import torch.nn as nn from src.config import Config +from src.data import pad_inputs class PrimeModel(nn.Module): @@ -27,7 +28,11 @@ def build_model(cfg: Config) -> PrimeModel: @torch.no_grad() def greedy_decode(model: PrimeModel, x: torch.Tensor, cfg: Config, max_len: int | None = None) -> torch.Tensor: - """Autoregressive greedy decode of output digits. Returns (B, max_len) tokens (BOS stripped).""" + """Autoregressive greedy decode of output digits. Returns (B, max_len) tokens (BOS stripped). + + Inputs are LEFT-padded to the global layout so positions match training exactly + (batch/singleton invariance — codex BLOCKER fix).""" + x = pad_inputs(x, cfg) max_len = max_len or cfg.max_out_len B = x.shape[0] y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=x.device) diff --git a/src/models/rnn.py b/src/models/rnn.py index 19123b7..44f0a0d 100644 --- a/src/models/rnn.py +++ b/src/models/rnn.py @@ -45,7 +45,8 @@ class TiedRNN(PrimeModel): return pe def _encode(self, x: torch.Tensor) -> torch.Tensor: - """Read input digits through the tied cell (pad positions inject nothing).""" + """Read input digits through the tied cell. Pad positions are exact no-ops + (state update masked) so batch composition cannot change an example's state.""" B, T = x.shape d = self.cfg.d_model pos = self._sinusoidal(T, d).to(x.device) # (T,d) @@ -53,7 +54,8 @@ class TiedRNN(PrimeModel): mask = (x != self.cfg.pad_id).float().unsqueeze(-1) # (B,T,1) h = torch.zeros(B, d, device=x.device) for t in range(T): - h = self._cell_step(h + (e[:, t] + pos[t]) * mask[:, t]) + h_new = self._cell_step(h + (e[:, t] + pos[t]) * mask[:, t]) + h = mask[:, t] * h_new + (1 - mask[:, t]) * h # pad step = no-op return h def _run_compute(self, h0: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -72,7 +74,7 @@ class TiedRNN(PrimeModel): for t in range(cfg.max_steps): h = self._cell_step(h) p = torch.sigmoid(self.halt_head(self.ln1(h))).squeeze(-1) # (B,) - if t < cfg.min_steps: + if t < cfg.min_steps - 1: p = p * 0.0 h_list.append(h) p_list.append(p) diff --git a/src/train.py b/src/train.py index 975598d..46bd8e7 100644 --- a/src/train.py +++ b/src/train.py @@ -1,8 +1,10 @@ """Training loop. Usage: python -m src.train [model] [seed] [--flag ...]""" import csv +import json import math import os import random +import sys import numpy as np import torch @@ -62,8 +64,21 @@ def main() -> None: cfg = parse_args() set_seed(cfg.seed) out_dir = os.path.join(cfg.out_dir, cfg.model, f"seed{cfg.seed}") + csv_path = os.path.join(out_dir, "metrics.csv") + if os.path.exists(csv_path): + raise SystemExit(f"REFUSING to rerun in place: {csv_path} exists. Use a fresh --out_dir " + f"(reruns would corrupt the CSV and checkpoint provenance).") os.makedirs(out_dir, exist_ok=True) cfg.save(os.path.join(out_dir, "config.json")) + with open(os.path.join(out_dir, "run_meta.json"), "w") as fh: + json.dump({ + "python": sys.version.split()[0], + "torch": torch.__version__, + "numpy": np.__version__, + "device": "cpu", + "torch_threads": torch.get_num_threads(), + "cmd": sys.argv, + }, fh, indent=2) train_in, val_in = get_splits(cfg) train_ex = build_examples(train_in, cfg) @@ -75,9 +90,8 @@ def main() -> None: opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay) ce = nn.CrossEntropyLoss(reduction="none") - log_examples = sorted(val_ex, key=lambda x: (len(str(x)), x))[: cfg.log_n_examples] + log_examples = sorted(val_ex, key=lambda ex: (len(ex[0]), ex[0]))[: cfg.log_n_examples] - csv_path = os.path.join(out_dir, "metrics.csv") fieldnames = ["step", "train_loss", "train_token_acc", "train_em", "val_token_acc", "val_em", "mean_halt_steps", "log_examples", "param_count"] best_val_em = -1.0 @@ -152,6 +166,8 @@ def main() -> None: _eval_pass(loss.detach(), halt.detach() if halt is not None else None) if done: break + if step >= cfg.max_train_steps: + break torch.save(model.state_dict(), os.path.join(out_dir, "last.pt")) print(f"DONE steps={step} best_val_em={best_val_em:.4f}") diff --git a/tests/test_codex_fixes.py b/tests/test_codex_fixes.py new file mode 100644 index 0000000..9e14113 --- /dev/null +++ b/tests/test_codex_fixes.py @@ -0,0 +1,104 @@ +"""Regression tests for codex-review fixes (design/reviews/codex-review.md).""" +import torch + +from src.config import Config +from src.data import build_examples, decode_tokens, make_batch +from src.model_api import build_model, greedy_decode + + +def test_mixed_batch_vs_singleton_invariance_rnn(): + cfg = Config(model="rnn") + m = build_model(cfg) + ex = build_examples([2, 42, 99], cfg) + batch = make_batch(ex, cfg) + out_batch = m(batch["x"], batch["y_in"]) + for i in range(3): + single = make_batch([ex[i]], cfg) + out_single = m(single["x"], single["y_in"]) + assert torch.allclose(out_batch["logits"][i], out_single["logits"][0], atol=1e-6), f"row {i}" + assert torch.allclose(out_batch["halt_steps"][i], out_single["halt_steps"][0], atol=1e-6) + + +def test_mixed_batch_vs_singleton_invariance_transformer(): + cfg = Config(model="transformer") + m = build_model(cfg) + ex = build_examples([2, 42, 99], cfg) + batch = make_batch(ex, cfg) + out_batch = m(batch["x"], batch["y_in"]) + for i in range(3): + single = make_batch([ex[i]], cfg) + out_single = m(single["x"], single["y_in"]) + assert torch.allclose(out_batch["logits"][i], out_single["logits"][0], atol=1e-6), f"row {i}" + + +def test_greedy_decode_layout_invariant(): + cfg = Config(model="transformer") + m = build_model(cfg) + ex = build_examples([42], cfg) + single = make_batch(ex, cfg) + out_a = greedy_decode(m, single["x"], cfg) + raw = torch.tensor([[4, 2]], dtype=torch.long) # unpadded — greedy must left-pad internally + out_b = greedy_decode(m, raw, cfg) + assert torch.equal(out_a, out_b) + + +def test_min_steps_floor_reachable(): + """min_steps off-by-one fix: earliest halt is step min_steps (forced-1 halt prob -> exactly floor).""" + cfg = Config(model="rnn", min_steps=2) + m = build_model(cfg) + with torch.no_grad(): + m.halt_head.bias.fill_(50.0) # p -> 1 at the first free step + ex = build_examples([2, 7, 42], cfg) + b = make_batch(ex, cfg) + out = m(b["x"], b["y_in"]) + assert torch.allclose(out["halt_steps"], torch.full_like(out["halt_steps"], float(cfg.min_steps)), atol=1e-3) + + +def test_pad_logits_cannot_change_loss(): + """Loss and grads must be insensitive to logits at padded target positions.""" + cfg = Config(model="rnn") + m = build_model(cfg) + b = make_batch(build_examples([2, 7], cfg), cfg) + out = m(b["x"], b["y_in"]) + y_safe = b["y"].clamp(max=cfg.vocab - 1) + loss = torch.nn.functional.cross_entropy(out["logits"].reshape(-1, cfg.vocab), + y_safe.reshape(-1), reduction="none") + mask = b["y_mask"].float().reshape(-1) + l1 = (loss * mask).sum() / mask.sum() + l1.backward() + grads = {k: v.grad.clone() for k, v in m.named_parameters() if v.grad is not None} + m.zero_grad() + out2 = m(b["x"], b["y_in"]) + logits2 = out2["logits"].clone() + pad_positions = ~b["y_mask"] # (B,T_out) + logits2[pad_positions] += 100.0 # huge perturbation on pad positions only + loss2 = torch.nn.functional.cross_entropy(logits2.reshape(-1, cfg.vocab), + y_safe.reshape(-1), reduction="none") + l2 = (loss2 * mask).sum() / mask.sum() + assert torch.allclose(l1.detach(), l2.detach()) + l2.backward() + for k, v in m.named_parameters(): + if v.grad is not None: + assert torch.allclose(grads[k], v.grad, atol=1e-6), k + + +def test_integers_vocab_eos_not_aliased(): + """Value 101 (= next_prime(100)) must not collide with EOS in integers mode.""" + cfg = Config(vocab_mode="integers", range_end=100) + ex = build_examples([100], cfg) + _, target = ex[0] + assert target[0] == 101 and target[1] == cfg.eos_id + assert cfg.eos_id != 101 + assert cfg.vocab == 103 + assert decode_tokens(target, cfg) == 101 + + +def test_act_weights_sum_to_one(): + """ACT mass conservation: steps must always lie in [min_steps, K] for random states.""" + cfg = Config(model="rnn", max_steps=20, min_steps=2) + m = build_model(cfg) + b = make_batch(build_examples([2, 7, 42, 99, 3, 15, 60, 97], cfg), cfg) + out = m(b["x"], b["y_in"]) + s = out["halt_steps"] + assert bool((s >= cfg.min_steps).all()) + assert bool((s <= cfg.max_steps).all()) diff --git a/tests/test_data.py b/tests/test_data.py index fe84e94..cd4e620 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -30,7 +30,7 @@ def test_encode_decode_roundtrip_integers(): cfg = Config(vocab_mode="integers", range_end=100) for n in [2, 42, 100]: assert decode_tokens(encode_int(n, cfg) + [cfg.eos_id], cfg) == n - assert cfg.vocab == 102 # 0..101 + EOS + assert cfg.vocab == 103 # values 0..101, EOS=102, pad=103 (no alias with target 101) def test_splits_sizes_and_no_overlap(): @@ -60,8 +60,8 @@ def test_make_batch_shapes_and_padding(): cfg = Config() ex = build_examples([2, 7, 42, 99], cfg) b = make_batch(ex, cfg) - assert b["x"].shape == (4, 2) # max input len 2 (42, 99) - assert b["y"].shape == (4, 4) # max output len 4 ("101"+EOS) + assert b["x"].shape == (4, 3) # GLOBAL layout: left-padded to 3 digits (max of [2,100]) + assert b["y"].shape == (4, 4) # GLOBAL layout: right-padded to 4 ("101"+EOS) assert b["y_in"].shape == (4, 4) assert b["y_mask"].shape == (4, 4) # pad positions masked out @@ -69,6 +69,9 @@ def test_make_batch_shapes_and_padding(): assert b["y_mask"][0, 1] # BOS position of y_in is eos_id assert (b["y_in"][:, 0] == cfg.eos_id).all() + # LEFT-padding: input 2 sits at the last column, pads at the front + assert b["x"][0, 2] == 2 and b["x"][0, 0] == cfg.pad_id + assert b["x"][2, 1] == 4 and b["x"][2, 2] == 2 def test_flagged_composites_are_composite(): |