From 2b890461c94ce1f45aabaa18c90cb3c5c4e5f7db Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Mon, 14 Sep 2026 15:51:24 -0700 Subject: [PATCH] Design source-pinned nanoGPT memorization study and tested measurement utilities Add a separate experiment folder for AdamW versus standard Muon, with clip_xmax WeightWatcher monitoring, controlled repetition/canary/association/retention protocols, and 32 passing component tests. Explicitly label the remaining training/data adapter and end-to-end checks as unimplemented; no training results or empirical optimality claimed. --- .../nanogpt_memorization/.gitignore | 7 + .../nanogpt_memorization/README.md | 314 ++++++++++++++++++ .../nanogpt_memorization/RESULTS.md | 30 ++ .../nanogpt_memorization/TRAINER_CONTRACT.md | 149 +++++++++ .../nanogpt_memorization/metrics.py | 127 +++++++ .../nanogpt_memorization/monitor.py | 22 ++ .../nanogpt_memorization/prepare_plan.py | 101 ++++++ .../nanogpt_memorization/probes.py | 59 ++++ .../nanogpt_memorization/study.json | 80 +++++ .../nanogpt_memorization/tests/test_study.py | 224 +++++++++++++ 10 files changed, 1113 insertions(+) create mode 100644 baseline/experiments/nanogpt_memorization/.gitignore create mode 100644 baseline/experiments/nanogpt_memorization/README.md create mode 100644 baseline/experiments/nanogpt_memorization/RESULTS.md create mode 100644 baseline/experiments/nanogpt_memorization/TRAINER_CONTRACT.md create mode 100644 baseline/experiments/nanogpt_memorization/metrics.py create mode 100644 baseline/experiments/nanogpt_memorization/monitor.py create mode 100644 baseline/experiments/nanogpt_memorization/prepare_plan.py create mode 100644 baseline/experiments/nanogpt_memorization/probes.py create mode 100644 baseline/experiments/nanogpt_memorization/study.json create mode 100644 baseline/experiments/nanogpt_memorization/tests/test_study.py diff --git a/baseline/experiments/nanogpt_memorization/.gitignore b/baseline/experiments/nanogpt_memorization/.gitignore new file mode 100644 index 0000000..447568c --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/.gitignore @@ -0,0 +1,7 @@ +__pycache__/ +.pytest_cache/ +*.pyc +runs/ +resolved/ +*.pt +*.bin diff --git a/baseline/experiments/nanogpt_memorization/README.md b/baseline/experiments/nanogpt_memorization/README.md new file mode 100644 index 0000000..be8b074 --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/README.md @@ -0,0 +1,314 @@ +# nanoGPT memorization: AdamW versus Muon + +**Status: experimental design and measurement components; no training results.** +This folder is independent of the existing baseline campaigns. It contains a +source-pinned plan resolver, behavioral metrics, synthetic-probe/schedule +utilities, and a hook into the existing WeightWatcher monitor. The end-to-end +training/data adapter is specified in `TRAINER_CONTRACT.md`, not implemented. +Do not run a generated YAML through the old trainer and describe its output as +this memorization study: that would omit the controlled data interventions. + +## Question and distinctions + +At matched architecture, initialization, sampled-token budget, and data order, +how do AdamW and Muon differ in acquisition, accessibility, and persistence of +specific training information? Do layer-resolved spectral changes accompany +those differences after accounting for training progress and validation loss? + +The measurements in the supplied figure are not five independent mechanisms. +Whole-sequence exact match is a criterion for verbatim extraction; token-level +match is a weaker criterion. Exposure measures rank within a specified candidate +universe. Prompt compression measures accessibility under a prompt class. +Training perplexity measures likelihood, not necessarily extractability. +Accordingly, vary the **information being learned**, and apply multiple +measurements to the same targets. + +The mechanisms to distinguish are: + +1. Rote recall of statistically uninformative, random token sequences. +2. Recall of particular natural-text passages, including repetition dependence. +3. Associative recall of synthetic key–value facts across surface templates. +4. Persistence versus forgetting after further presentations stop. + +A held-out compositional rule task is a positive control for generalization, +not a fifth claim of memorization. All canaries and key–value facts are synthetic; +no real personal information is needed. + +## Source settings: inherit, do not retune on memorization + +Source commit: `3749c36334382a20e48bfe2473c1dc4a1470a830`. +Source recipe: +`baseline/experiments/nanogpt_one_head_2026_08_21_baseline/configs/baseline.yaml`. +Git blob: `7fd3c592afc7fdfea952b1f1aa0f9b44cdc53a2b`. + +The repository calls these **source-backed center settings**, not empirically +frozen optimal settings. Its qualification protocol requires validation-only +selection and a winner lock before making the stronger claim. This experiment +does not invent such a lock. The changed data interventions also define a new +experimental protocol, even though model and optimizer settings are inherited. + +| Component | Inherited setting | +|---|---| +| Model | 1 block, 1 attention head, width 128, context 256 | +| Vocabulary | GPT-2 BPE, 50,257 tokens; tied embedding/output head | +| Dropout / biases | 0 / false | +| Background data | FineWeb-Edu `sample-10BT`, revision `593b3a867298afb8ce42625a270ef20ddcad28f9` | +| Splits | 80M train / 1M validation / 1M monitoring-only test; document disjoint | +| Effective batch | 4 sequences × 8 accumulation steps × 256 target tokens = 8,192 | +| Gradient clipping | Global norm 1.0 | +| AdamW | LR 6e-4 → 6e-5; betas (0.90, 0.95); epsilon 1e-8; matrix WD 0.10 | +| AdamW warmup | 1% of the inherited one-reference-epoch schedule | +| Muon hidden matrices | LR 0.02 → 0.002; momentum 0.95; Nesterov; 5 Newton–Schulz steps; WD 0.01 | +| Muon auxiliary AdamW | LR 3e-4 → 3e-5; betas (0.90, 0.95); epsilon 1e-8; matrix WD 0.01 | +| Muon warmup | 5% of the inherited one-reference-epoch schedule | +| Replication | Paired seeds 1337, 2027, 4099, 31415, 271828 | +| Runtime | Inherited deterministic FP32/highest-matmul policy; no TF32; MPS fallback recorded | + +Muon applies to Q, K, V, attention output, MLP input, and MLP output matrices. +Embeddings, the tied output head, normalization parameters, and other auxiliary +parameters retain the existing auxiliary-AdamW partition. Use the repository's +optimizer implementation; do not substitute a different package's Muon. + +The separate historical `muon_clip` profile uses 2e-4 → 2e-5, RMS scaling 0.20, +and QK threshold 100. It is **not standard Muon** and is not an active arm here. +A later MuonClip comparison must be explicitly labeled and separately configured. + +The LR schedule decays over one 80M-token reference epoch and remains at the +nonzero floor through reference epoch four, as in the selected campaign. Do not +stretch the cosine over four epochs, restart it at a retention fork, or equate +numerical learning rates across different optimizer update normalizations. + +## Matched campaign + +| Condition | Training intervention | Main comparison | +|---|---|---| +| `clean` | No experimental targets; matched ordinary-background replacement slots | Same-target counterfactual reference for the other arms | +| `sequence` | Reserved natural passages, uniform-random token sequences, finite-universe canaries | Whole-sequence recall, dose response, exact exposure, prefix accessibility | +| `association` | Synthetic key–value mappings in varied templates; separate rule-control examples | Fact recall versus literal template recall versus rule generalization | + +Each condition uses both optimizers and all five paired seeds: **30 root runs**. +At 39,063 updates, each run sees 320,004,096 prediction targets. The complete +root campaign is approximately 9.60 billion target tokens, before evaluation. +This is a staged design, not a recommendation to launch all jobs on a laptop. + +Start with a one-reference-epoch pilot, seed 1337, `clean` and `sequence`, both +optimizers: four runs. A pilot checks feasibility, learning signal, fit validity, +and runtime. It is not evidence for an optimizer advantage. Any pilot-driven +protocol changes must be locked before the five-seed confirmation; exclude the +pilot from confirmatory inference if it informed those changes. + +For the main campaign, retain fixed-budget final states, validation-best states, +and the complete trajectory. Do not early-stop when a memorization curve looks +favorable. The protected test split never selects a setting, checkpoint, horizon, +or spectral threshold. Muon-versus-AdamW estimates refer first to these inherited +**recipes**, not to an isolated mathematical update rule: their decay, auxiliary +learning rates, and warmups differ by design. A mechanism claim requires a +separately preregistered matched-decay/auxiliary ablation. + +### Sequence and repetition experiment + +Use 32 targets per dose per family, with doses **0, 1, 4, 16, 64 complete +presentations per reference epoch**. Each target has 128 context tokens followed +by a 64-token continuation. The natural and random families are separately +reported, never pooled into a single headline extraction rate. + +Natural passages come from a reserved document pool removed from ordinary +background sampling. Enforce document boundaries and reject exact target +contamination in the remaining train/validation/test streams. A passage described +as zero-dose must never be presented during training, including accidental +cross-window occurrences. Token-random strings exclude EOT and are not described +as ordinary natural language or random-character strings. + +Injection replaces complete 257-token records at predetermined training slots; +it never appends extra tokens to one optimizer's budget. Supply the full context +and continuation in one record. Random-window sampling through an injected +binary file is not sufficient: it can miss the target, split it, or expose it +more often than the stated duplication count. Log actual full-target visits, +partial-target visits, and supervised target-token counts. + +The two sequence families require 5,440 full-record presentations per reference +epoch. The canary cohort adds 680. This is approximately 1.96% of the 80M-target +budget when each occupies one 256-target record. Both optimizers receive exactly +the same slot schedule; record-level random filler is not a second repeated +memorization target. Match replacement-slot budgets across all conditions. + +### Canary exposure experiment + +Use eight independently keyed canaries per dose. A canary suffix is three tokens +from a fixed, unique 16-token alphabet. Its **entire declared universe** therefore +contains 16^3 = 4,096 candidates. Select the planted suffix uniformly and keep +candidate length and prefix fixed when ranking. Distinct nonces separate canary +contexts; a code may recur under another nonce, but the keyed mapping must not. + +Compute suffix-only summed NLL for every candidate. Report + +`exposure = log2(4096) - log2(rank)`. + +The maximum is 12 bits for this deliberately small universe. This is not a claim +of 12-bit privacy loss, a universal extraction probability, or the exposure of a +32-token unrestricted secret. Report rank bounds for tied scores; the primary +value is the conservative exposure bound. All-equal scores yield zero primary +exposure. Compare with zero-dose canaries and with the same target in the paired +clean model to control intrinsic token preferences. + +Exhaustive exposure is scheduled at reference epochs 0, 1, 2, and 4, not every +training step. Save all candidate scores and the universe hash for auditability. + +### Associative recall and generalization control + +Generate arbitrary, synthetic key–value mappings with no predictable relation +between key and value. For each entity, rotate one of three templates into a +held-out role; that template must occur for other entities during training. +Hold out **entity–template combinations**, not an entirely unfamiliar language +format. A one-presentation entity sees one training template. At larger doses, +distribute its total dose across the two available training templates. + +Measure value-only greedy exact match and suffix NLL on the original template, +held-out template, unseen keys, and deliberately permuted values. Correct recall +on a held-out template supports accessible associative storage; success on +unseen random keys is not expected and signals leakage or a construction error. + +Separately train a simple compositional rule, for example a small synthetic +attribute-to-token mapping with held-out attribute combinations. Improving on +held-out combinations is evidence of rule generalization. It must not be +counted as memorization merely because its target is predictable. + +### Retention fork (follow-on, not automatically launched) + +At reference epoch one of each `sequence` run, fork the full checkpoint into +continued-exposure and withdrawn-exposure children. The continued child is the +ordinary sequence trajectory; only the withdrawn child adds computation. It +replaces future target slots with matched background records. No reset of model, +optimizer momentum, auxiliary AdamW state, RNG, LR schedule, or sampled-token +counter is allowed. Parent and first child state hashes must agree. + +Record retention curves, loss of previously acquired exact matches, exposure +changes, and time since each target's last actual presentation. Report half-life +only when an identifiable decline crosses half the initial excess-over-control +signal; otherwise report a censored or undefined estimate, not an invented +half-life. Non-monotonic forgetting/relearning is allowed. + +## Measurements on the same fixed probes + +**Primary extraction endpoint:** whole-sequence greedy EM@32 with a 64-token +prefix, evaluated at the fixed final token budget. Report EM@16 and EM@64, +free-running token-match fraction, and longest exact prefix separately. +Teacher-forced token accuracy is not free-running extraction. + +**Prefix-constrained compression:** use prefix lengths 8, 16, 32, 64, 128, +all ending at the same target boundary. Report the shortest successful prefix +in that grid and continuation-length/prefix-length. Do not change the target +when changing prefix length, assume success is monotone, or call this optimized +adversarial compression. Failure is censored with respect to this prompt class. + +**Likelihood:** report mean suffix NLL and perplexity on matched seen/unseen +cohorts. Preserve token weighting and suffix masks. Low training perplexity by +itself is neither proof of verbatim storage nor proof of a privacy leak. + +Keep per-probe rows keyed by run, seed, optimizer, step, family, target ID, +intended dose, actual cumulative presentations, last-presentation step, template, +prefix length, and continuation length. Write counts/denominators with every +aggregate; a failed evaluation is missing, not zero memorization. + +## WeightWatcher monitoring + +Reuse `rg_nanogpt_one_head.spectral.run_weightwatcher` through `monitor.py`. +The required analysis is the repository's single-call configuration: + +```python +watcher.analyze( + ERG=True, + randomize=True, + fix_fingers="clip_xmax", + min_evals=20, + max_fingers=10, + plot=False, +) +``` + +This clips the **fitted spectral range**, not the model weights, and does not +activate MuonClip. Monitoring must never feed back into optimizer updates. + +Retain per-matrix `alpha`/`alpha_clip_xmax`, `raw_alpha`/`alpha_raw`, their +difference, finger count, `ERG_gap`, `num_traps`, `rand_distance`, fit distance +`D`, ranks, spectral norms, and whatever fit-support metadata is actually +returned. Preserve missing fields as missing. Save unmodified per-layer outputs +alongside the behavioral metrics and the exact checkpoint hash. + +Analyze initialization, updates 1/10/100/500, every quarter-reference-epoch, +each optimizer's actual inherited warmup endpoint, and final state. Behavioral +probes run every 500 updates; expensive exposure/prefix sweeps use their sparse +cadences in `study.json`. Save a restart checkpoint before expensive diagnostics. + +Do not treat a finite fitted alpha as proof that a random/MP-like layer has a +meaningful power-law tail. Inspect spectral support, tail sample count, fit +quality, and randomized-ESD separation. Predeclare a null/fit-validity procedure +before interpreting alpha as a heavy-tail signal. Flag random-like or inconclusive +layers; never average their alphas blindly across the network. The current +one-randomization distance is descriptive, not a calibrated significance test. + +Hidden matrices are the primary inherited monitor inventory. The tied embedding +is counted once in a separately labeled, lower-cadence follow-on audit if added; +do not pool it with hidden matrices or imply that the six hidden matrices alone +locate every stored association. + +## Analysis and interpretation + +Use a complete paired training seed as the replication unit. Layers, probes, +checkpoints, and candidate strings are repeated measurements, not independent +training replicates. Show every seed and paired optimizer difference; use +seed-level intervals, with the low precision of five seeds made explicit. +Predeclare multiplicity handling for the extraction, exposure, and associative +primary endpoint families; treat remaining grids as exploratory. + +First compare at equal sampled-token budgets. Secondarily compare at equal +validation NLL, but only within the actual overlapping loss range; never +extrapolate or choose levels from protected-test results. An optimizer that +learns faster can otherwise appear to memorize more at the same step simply +because it has progressed further. + +Relate spectral metrics to later behavioral changes using within-run differences +and models that account for step, dose, condition, and validation NLL. Check +whether results survive exclusion of invalid/random-like fits. A spectral +correlation is not a causal explanation, a stand-alone memorization detector, +or evidence that one optimizer is safer. Family-specific attribution requires +separate interventions, not a network-average alpha from a mixed corpus. + +## Files and commands + +- `study.json`: preregistered design, source identities, doses, seeds, and cadence. +- `prepare_plan.py`: verifies source blobs and emits inherited configs plus manifest. +- `metrics.py`: greedy recall, conditional NLL, exact exposure, prefix compression. +- `probes.py`: random-token probes, finite canary universe, exact slot schedules. +- `monitor.py`: strict adapter to the existing WeightWatcher implementation. +- `TRAINER_CONTRACT.md`: required data/training integration and acceptance gates. +- `tests/test_study.py`, `RESULTS.md`: bounded component tests and honest status. + +From the repository root, using its existing experiment environment: + +```bash +python -m pytest -q baseline/experiments/nanogpt_memorization/tests +python baseline/experiments/nanogpt_memorization/prepare_plan.py \ + --out /tmp/rg-nanogpt-memorization-plan-20260914 +``` + +The second command creates a plan, not jobs. It refuses an existing output +folder and refuses recipe/model/optimizer/spectral-source drift. Setup for the +existing source package, when needed, remains: + +```bash +python -m pip install -e './baseline/nanogpt_one_head[dev]' +``` + +Do not point results or mutated datasets at an older baseline's run/cache folder. + +## Primary methodological references + +- Carlini et al., *Quantifying Memorization Across Neural Language Models*, + arXiv:2202.07646 — extraction, context length, and repetition. +- Carlini et al., *The Secret Sharer*, arXiv:1802.08232 — canaries and exposure. +- Schwarzschild et al., *Rethinking LLM Memorization through the Lens of + Adversarial Compression*, arXiv:2404.15146 — prompt compression; the prefix + sweep here is deliberately narrower than that paper's optimized attack. +- Repository `baseline/FINAL_BASELINE_QUALIFICATION.md` — validation-only recipe + selection; source-backed defaults are not a demonstrated optimum. diff --git a/baseline/experiments/nanogpt_memorization/RESULTS.md b/baseline/experiments/nanogpt_memorization/RESULTS.md new file mode 100644 index 0000000..4ce1161 --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/RESULTS.md @@ -0,0 +1,30 @@ +# Execution ledger + +Date: 2026-09-14. + +## Completed in this preparation session + +- Read the repository's current baseline recipe, model, optimizer implementation, + spectral-monitor implementation, qualification document, and dependency file. +- Pinned source commit and Git blob identities in `study.json`. +- Created this separate experimental-design/measurement folder. +- Ran `python -m pytest -q tests`: **32 passed in 2.78 seconds**, CPU. +- Tests cover exact-versus-token recall, teacher-forced-versus-free-running + behavior, suffix alignment, exhaustive finite-universe exposure and ties, + prefix-search censoring, candidate batching, deterministic probes, exact + scheduled presentations, source-drift checks, plan counts, and WW guards. + +## Not completed / no results claimed + +- No end-to-end memorization trainer or natural/association data adapter has been + implemented; the integration contract is in `TRAINER_CONTRACT.md`. +- No FineWeb-Edu download, canary-injected training, or AdamW/Muon training + campaign was run in this session. +- No actual WeightWatcher call was executed here; WeightWatcher is not installed + in this execution environment. The hook reuses the inspected repository code. +- No pinned-nanoGPT optimizer-step, resume, or target-hardware integration test + was executed. The controlled test model is not nanoGPT. +- No empirically optimal configuration was established; the repository's + source-backed recipe is inherited without relabeling it as a frozen optimum. +- There are no measured optimizer advantages, memorization rates, spectral + trends, or significance claims yet. diff --git a/baseline/experiments/nanogpt_memorization/TRAINER_CONTRACT.md b/baseline/experiments/nanogpt_memorization/TRAINER_CONTRACT.md new file mode 100644 index 0000000..7f3b64a --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/TRAINER_CONTRACT.md @@ -0,0 +1,149 @@ +# Training adapter contract and acceptance gates + +This is a specification for the remaining integration, not a claim that the +trainer exists. The baseline model, optimizer, and spectral monitor are reused; +only the experimental data/sampling and behavioral-evaluation integration change. +The generated baseline YAML files alone do not implement these interventions. + +## 1. Load and verify + +Resolve `study.json` with `prepare_plan.py`. Read the source-pinned recipe rather +than copying hyperparameters by hand. Validate the prepared FineWeb cache with +`rg_nanogpt_one_head.data.validate_prepared_data(data_dir, baseline_config)`. +Require its dataset revision, exact token/byte counts, document-disjoint split +flag, and SHA-256 checks. Do not alter the original `train.bin` or its metadata. +Store intervention manifests and outputs under a new explicit experiment root. + +Use `GPT(GPTConfig(**cfg['model']))` from the pinned model module. Construct +optimizer handles using `make_optimizer_handles(model, profile)` from the pinned +optimizer module. Preserve the entire parameter partition and initialization. +A dedicated training wrapper must explicitly support `muon`: the old dated +campaign's CLI was configured around AdamW and MuonClip and is not an automatic +launcher for this study. + +## 2. Freeze information and sample controls + +Create target identities, canary universes, templates, rule examples, and the +reserved natural-document pool once with the study data seed. Save tokens, +source document identities, split membership, masks, and content hashes. + +Reject every natural target whose full 16/32/64-token scored suffix is already +present in the allowed background or another split at the corresponding score +length. Match held-out natural controls on token length and source domain; the +paired clean model controls the same target's intrinsic predictability. Disjoint +document IDs alone do not eliminate duplicate text across documents. + +For random-token targets, scan for collisions rather than assuming probability +zero. For canaries, audit context–suffix pairs and each template-equivalent +mapping. For associations, rotate held-out entity–template combinations while +ensuring every template is learned on other entities. `probes.py` currently +implements random sequences and finite canary enumeration only; the natural +reservation/decontamination and association/rule generators remain to be built. + +Do not infer empirical training frequency from the intended duplication factor. +A full experimental presentation is a complete context plus target in one +training record with all scored target tokens included in the loss. + +## 3. Paired training stream and exact dose + +Use distinct RNG streams for initialization, background batches, injection +slots, filler, evaluation, and spectral randomization. Pair initialization and +all training-data streams between optimizers. Hash initial model state and a +preflight trace of batch IDs to verify pairing. + +For each reference-epoch interval, create a `presentation_schedule` over the +actual available record slots. Its seed is derived deterministically from the +training seed and interval, with the same derivation for both optimizers. +Doses are full presentations per interval. The small rounding difference in the +last interval is handled by its actual slot count, not by changing the dose. + +The adapter replaces selected ordinary-background records with intervention +records. It must preserve the effective 8,192 supervised prediction targets per +optimizer update. Pack or fill a 257-token record so a 256-token input has 256 +next-token labels; do not leave padding targets unintentionally supervised or +silently mask different numbers of tokens between conditions. Record how much +loss comes from context, target, and filler. Do not train a scored target across +an attention-context boundary. + +Log each intervention visit with optimizer update, accumulation index, record +index, target ID, template, and number of supervised target tokens. At every +checkpoint compare actual cumulative visits with the planned schedule prefix. +Zero-dose targets must have zero visits. Equivalent base replacement slots are +used in the clean control, filled with ordinary background records. + +## 4. Optimizer step and checkpoint identity + +Reuse baseline learning-rate scheduling, gradient accumulation, global clipping, +and optimizer-step helpers. Resolve warmup steps using the actual baseline +trainer's rounding convention, then union these actual warmup endpoints into +`plan['spectral_steps_before_warmup_union']`. Do not guess that convention in a +new implementation or rescale the one-epoch schedule to the four-epoch horizon. + +An update's logged LR must be the value used for that update. Every behavioral +and spectral record is keyed to the same post-update model-state hash and +sampled-token counter. Initialization is update zero. Precision, device, compile +policy, and dependency versions belong in the run fingerprint. + +Save checkpoint state atomically to a temporary file, fsync, and rename. Keep a +rolling restart checkpoint plus immutable diagnostic states. Reuse the baseline +checkpoint/RNG implementation and retain model, optimizer handles, LR position, +CPU/accelerator RNG states, dedicated data RNG, dose counters, and current +schedule offset. Verify round-trip equality and the next resumed update before +starting long jobs. Never restart a partial result directory silently. + +## 5. Behavioral and spectral hooks + +Evaluate only fixed probes; never train on evaluation-generated continuations. +Use the functions in `metrics.py` for suffix-only NLL, teacher-forced accuracy, +free-running recall, canary ranking, and fixed-boundary prefix sweeps. Save the +returned candidate score vectors, not merely rounded ranks. + +After saving a checkpoint, call: + +```python +from monitor import monitor_training_state + +summary = monitor_training_state( + model, + run_dir, + step=completed_updates, + tokens_seen=sampled_target_tokens, + reference_tokens=80_000_000, + seed=training_seed, + fingerprint=run_fingerprint, + ww_config=baseline_config['weightwatcher'], +) +``` + +The inherited monitor uses one `fix_fingers='clip_xmax'` analysis on CPU clones, +with ERG and randomization enabled. It preserves RNG state and binds results to +the checkpoint hash. Test that enabling diagnostics does not change the next +training batch, next model update, or optimizer state. + +Keep `alpha_raw` and `alpha_clip_xmax` separately. Never substitute a raw fit +for a failed clipped fit, synthesize ERG values, label a failed fit zero, or use +WeightWatcher as the behavioral definition of memorization. If strict monitoring +fails, retain the restart checkpoint and write a failure status before stopping; +resume diagnostics only after resolving the cause. A finite alpha may still be +scientifically invalid for a random-like spectrum; the interpretation screen is +separate from successful software execution. + +## 6. Before a pilot may be labeled runnable + +The adapter must pass all of the following, beyond the component tests shipped +here: + +- Tiny synthetic overfit: increasing exposure increases exact recall on a + deliberately learnable toy cohort; zero-dose rows never enter optimization. +- Natural/association contamination audit and declared finite-universe audit. +- Paired initial hashes, batch/slot traces, and exact realized dose counts. +- Actual pinned nanoGPT AdamW and Muon forward/backward/update preflight. +- Actual WeightWatcher output schema and RNG/model invariance on target hardware. +- Full checkpoint round-trip, interruption/resume, and next-update equivalence. +- Matching masks/target-token counts, context bounds, and final token horizon. +- Same-state checkpoint agreement between online and offline behavioral metrics. + +The test suite in this folder currently checks measurement mathematics, suffix +alignment, a controlled transition model, scheduling counts, plan resolution, +source-drift refusal, and monitoring configuration guards. It does not establish +that these end-to-end integration gates have passed. diff --git a/baseline/experiments/nanogpt_memorization/metrics.py b/baseline/experiments/nanogpt_memorization/metrics.py new file mode 100644 index 0000000..d00210a --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/metrics.py @@ -0,0 +1,127 @@ +"""Behavioral measurements: no spectral metric is used to define memorization.""" +from __future__ import annotations +import math +from typing import Callable, Sequence +import numpy as np +import torch +import torch.nn.functional as F + + +def recall(generated: Sequence[int], target: Sequence[int]) -> dict: + g, y = np.asarray(generated), np.asarray(target) + if g.ndim != 1 or y.ndim != 1 or len(y) == 0 or g.shape != y.shape: + raise ValueError("Generated and target sequences must be nonempty, 1-D, and equal-length.") + match = g == y + failures = np.flatnonzero(~match) + return {"sequence_exact_match": bool(match.all()), + "free_running_token_match": float(match.mean()), + "longest_exact_prefix": int(failures[0]) if len(failures) else len(y), + "continuation_tokens": len(y)} + + +def exposure(candidate_nll: Sequence[float], target_index: int) -> dict: + """Exact finite-universe exposure; smaller NLL ranks first. + + Return tie bounds. Primary exposure is conservative: all-equal scores -> 0. + These are full-universe ranks ONLY if every possible candidate was scored. + """ + scores = np.asarray(candidate_nll, dtype=np.float64) + if scores.ndim != 1 or not len(scores) or not np.isfinite(scores).all(): + raise ValueError("All candidate NLL scores must be finite and present.") + if not 0 <= target_index < len(scores): + raise ValueError("Target index is outside the candidate universe.") + value = scores[target_index] + rank_min = 1 + int(np.count_nonzero(scores < value)) + rank_max = int(np.count_nonzero(scores <= value)) + bits = math.log2(len(scores)) + return {"candidate_count": len(scores), "rank_min": rank_min, "rank_max": rank_max, + "exposure_bits_lower": bits - math.log2(rank_max), + "exposure_bits_upper": bits - math.log2(rank_min), + "tie_count": rank_max - rank_min + 1} + + +def prefix_compression(context: Sequence[int], target: Sequence[int], + lengths: Sequence[int], generate: Callable[[list[int], int], list[int]]) -> dict: + """Sweep suffixes of ONE fixed context; never move the target boundary. + + This is prefix-constrained prompt compression, not optimized adversarial ACR. + Failure is censored, not proof that no shorter/adversarial prompt exists. + """ + if len(context) == 0 or len(target) == 0 or len(lengths) == 0: + raise ValueError("Context, target, and prefix grid must be nonempty.") + if any(int(n) != n for n in lengths): + raise ValueError("Prefix lengths must be integers.") + grid = sorted(set(int(n) for n in lengths)) + if grid[0] < 1 or grid[-1] > len(context): + raise ValueError("Prefix grid is outside the available context.") + rows = [] + for p in grid: + rows.append({"prefix_tokens": p, **recall(generate(list(context[-p:]), len(target)), target)}) + successes = [r["prefix_tokens"] for r in rows if r["sequence_exact_match"]] + shortest = min(successes) if successes else None + return {"rows": rows, "shortest_successful_prefix_in_grid": shortest, + "prefix_compression_ratio": len(target) / shortest if shortest is not None else None, + "search_censored": shortest is None, + "prompt_class": "fixed_context_suffixes_only"} + + +@torch.inference_mode() +def score_continuation(model, prefix: Sequence[int], target: Sequence[int]) -> dict: + """Teacher-forced suffix NLL/accuracy and separate free-running recall.""" + if len(prefix) == 0 or len(target) == 0: + raise ValueError("Prefix and continuation must both be nonempty.") + if len(prefix) + len(target) - 1 > model.cfg.block_size: + raise ValueError("Teacher-forced input would exceed model context.") + device = next(model.parameters()).device + full = torch.tensor([list(prefix) + list(target)], dtype=torch.long, device=device) + previous = model.training + model.eval() + try: + logits, _ = model(full[:, :-1]) + selected = logits[:, len(prefix) - 1:, :] + labels = full[:, len(prefix):] + losses = F.cross_entropy(selected.reshape(-1, selected.size(-1)), labels.reshape(-1), reduction="none") + nll = float(losses.mean().item()) + if not math.isfinite(nll): + raise RuntimeError("Nonfinite continuation loss; do not report a memorization score.") + greedy = model.generate_greedy(full[:, :len(prefix)], len(target))[0, len(prefix):].tolist() + return {**recall(greedy, target), "suffix_nll": nll, + "suffix_nll_sum": float(losses.sum().item()), + "suffix_perplexity": math.exp(nll) if nll < 709 else None, + "teacher_forced_token_accuracy": float((selected.argmax(-1) == labels).float().mean().item())} + finally: + model.train(previous) + + +@torch.inference_mode() +def rank_canary(model, prefix: Sequence[int], candidates: Sequence[Sequence[int]], + target_index: int, batch_size: int = 4) -> dict: + """Exhaustively score a small declared universe, not sampled rank estimation.""" + if len(prefix) == 0 or len(candidates) == 0 or batch_size < 1: + raise ValueError("Nonempty prefix/universe and positive batch size are required.") + if not 0 <= target_index < len(candidates): + raise ValueError("Target index is outside the candidate universe.") + length = len(candidates[0]) + if not length or any(len(c) != length for c in candidates): + raise ValueError("The declared universe must use equal nonzero token lengths.") + if len({tuple(c) for c in candidates}) != len(candidates): + raise ValueError("Candidate universe contains duplicates.") + if len(prefix) + length - 1 > model.cfg.block_size: + raise ValueError("Canary scoring would exceed context.") + device = next(model.parameters()).device + previous = model.training + model.eval() + scores = [] + try: + for start in range(0, len(candidates), batch_size): + chunk = candidates[start:start + batch_size] + full = torch.tensor([list(prefix) + list(c) for c in chunk], dtype=torch.long, device=device) + logits, _ = model(full[:, :-1]) + logits = logits[:, len(prefix) - 1:, :] + target = full[:, len(prefix):] + loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), target.reshape(-1), reduction="none") + scores.extend(loss.reshape(len(chunk), length).sum(-1).double().cpu().tolist()) + finally: + model.train(previous) + return {**exposure(scores, target_index), "candidate_nll_sum": scores, + "rank_scope": "entire_declared_finite_universe"} diff --git a/baseline/experiments/nanogpt_memorization/monitor.py b/baseline/experiments/nanogpt_memorization/monitor.py new file mode 100644 index 0000000..173335c --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/monitor.py @@ -0,0 +1,22 @@ +"""Live-model hook reusing the repository's audited spectral implementation.""" +from __future__ import annotations +from pathlib import Path + + +def monitor_training_state(model, run_dir: str | Path, *, step: int, tokens_seen: int, + reference_tokens: int, seed: int, fingerprint: str, ww_config: dict): + if ww_config.get("fix_fingers") != "clip_xmax": + raise ValueError("This study requires fix_fingers='clip_xmax'.") + for key in ("enabled", "ERG", "randomize", "strict", "require_raw_alpha"): + if ww_config.get(key) is not True: + raise ValueError(f"This study requires weightwatcher.{key}=true.") + if step < 0 or tokens_seen < 0 or reference_tokens < 1 or not fingerprint: + raise ValueError("Invalid checkpoint identity.") + from rg_nanogpt_one_head.spectral import run_weightwatcher + # Upstream creates CPU clones of Q/K/V/O/MLP-in/MLP-out, calls WW once, + # preserves CPU/accelerator RNG, binds output to the model-state hash, + # and stores alpha, raw_alpha, num_fingers, ERG_gap, traps and rand_distance. + # No mutation of weights, clipping of weights, or feedback to optimization. + return run_weightwatcher(model, run_dir, step=step, tokens_seen=tokens_seen, + train_tokens=reference_tokens, config=dict(ww_config), + seed=seed, fingerprint=fingerprint) diff --git a/baseline/experiments/nanogpt_memorization/prepare_plan.py b/baseline/experiments/nanogpt_memorization/prepare_plan.py new file mode 100644 index 0000000..a835e8c --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/prepare_plan.py @@ -0,0 +1,101 @@ +"""Resolve a source-pinned study plan. This does not launch training.""" +from __future__ import annotations +import argparse +import copy +import hashlib +import json +import math +from pathlib import Path +import yaml + +HERE = Path(__file__).resolve().parent + + +def git_blob_sha(data: bytes) -> str: + return hashlib.sha1(b"blob " + str(len(data)).encode() + b"\0" + data).hexdigest() + + +def fingerprint(value: object) -> str: + return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + +def resolve(repo: Path, study: dict, *, verify_modules: bool = True) -> dict: + source = study["source"] + raw = (repo / source["recipe"]).read_bytes() + if git_blob_sha(raw) != source["recipe_git_blob_sha"]: + raise RuntimeError("Baseline recipe changed: review and repin the study; do not silently reuse it.") + if verify_modules: + root = repo / "baseline/nanogpt_one_head/src/rg_nanogpt_one_head" + for name, expected in source["modules"].items(): + if git_blob_sha((root / name).read_bytes()) != expected: + raise RuntimeError(f"Baseline {name} changed: review and repin the study.") + base = yaml.safe_load(raw) + if study["optimizers"] != ["adamw", "muon"]: + raise ValueError("Primary comparison is AdamW versus Muon, not silently substituted MuonClip.") + for key in study["optimizers"]: + if base["optimizer_profiles"][key]["family"] != key: + raise ValueError(f"Optimizer family mismatch for {key}.") + if study["seeds"] != base["training"]["seeds"]: + raise ValueError("Paired seeds differ from the selected repository campaign.") + if base["dataset"]["train_tokens"] != study["training"]["reference_tokens"]: + raise ValueError("Reference token horizon does not match the inherited corpus.") + for key in ("fix_fingers", "ERG", "randomize", "min_evals", "max_fingers", "require_raw_alpha"): + if base["weightwatcher"][key] != study["monitoring"][key]: + raise ValueError(f"WeightWatcher setting mismatch: {key}.") + t = base["training"] + tokens_per_step = t["batch_size"] * t["grad_accum_steps"] * base["model"]["block_size"] + ref = study["training"]["reference_tokens"] + epochs = study["training"]["target_reference_epochs"] + steps_at = lambda epoch: math.ceil(epoch * ref / tokens_per_step) + spectral = set(study["monitoring"]["spectral_early_steps"]) + interval = study["monitoring"]["spectral_every_reference_epochs"] + spectral.update(steps_at(i * interval) for i in range(math.floor(epochs / interval) + 1)) + # Warmup end is resolved by the baseline trainer, not guessed by this planner. + runs = [] + for condition in study["conditions"]: + for seed in study["seeds"]: + for optimizer in study["optimizers"]: + cfg = copy.deepcopy(base) + cfg["protocol"]["name"] = study["name"] + cfg["protocol"]["description"] = "Memorization study derived from source-backed settings; not a qualification lock." + cfg["training"]["seeds"] = [seed] + cfg["training"]["target_epochs"] = epochs + cfg["optimizer_profiles"] = {optimizer: cfg["optimizer_profiles"][optimizer]} + # Historical display text says Muon was not an arm of the older campaign. + cfg["optimizer_profiles"][optimizer]["display_name"] = { + "adamw": "AdamW", "muon": "Muon + auxiliary AdamW" + }[optimizer] + runs.append({"run_id": f"{condition}/{optimizer}/seed_{seed}", + "condition": condition, "optimizer": optimizer, "seed": seed, + "baseline_config": cfg}) + result = { + "study": study, "tokens_per_update": tokens_per_step, + "schedule_steps": steps_at(1), "total_steps": steps_at(epochs), + "actual_target_tokens_per_run": steps_at(epochs) * tokens_per_step, + "spectral_steps_before_warmup_union": sorted(s for s in spectral | {steps_at(epochs)} if s <= steps_at(epochs)), + "full_exposure_steps": [steps_at(e) for e in study["monitoring"]["full_exposure_reference_epochs"]], + "prefix_sweep_steps": [steps_at(e) for e in study["monitoring"]["prefix_sweep_reference_epochs"]], + "runs": runs, + "execution_status": "plan_only_training_adapter_not_implemented", + } + result["fingerprint"] = fingerprint(result) + return result + + +def write_plan(plan: dict, out: Path) -> None: + # Refuse reuse instead of silently overwriting another experiment. + out.mkdir(parents=True, exist_ok=False) + (out / "resolved_plan.json").write_text(json.dumps(plan, indent=2) + "\n") + for run in plan["runs"]: + path = out / "configs" / (run["run_id"].replace("/", "__") + ".yaml") + path.parent.mkdir(exist_ok=True) + path.write_text(yaml.safe_dump(run["baseline_config"], sort_keys=False)) + print(f"Wrote {len(plan['runs'])} planned runs to {out}. No training was started.") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo", type=Path, default=HERE.parents[2]) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + write_plan(resolve(args.repo, json.loads((HERE / "study.json").read_text())), args.out) diff --git a/baseline/experiments/nanogpt_memorization/probes.py b/baseline/experiments/nanogpt_memorization/probes.py new file mode 100644 index 0000000..69102b3 --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/probes.py @@ -0,0 +1,59 @@ +"""Deterministic synthetic probes and exact full-record presentation schedules. + +Natural-span decontamination and the baseline training adapter are deliberately +not implemented here; see TRAINER_CONTRACT.md before running a study. +""" +from __future__ import annotations +from itertools import product +import numpy as np + + +def canary_universe(alphabet: list[int], length: int) -> list[list[int]]: + if length < 1 or len(alphabet) < 2 or len(set(alphabet)) != len(alphabet): + raise ValueError("Use a unique alphabet with at least two tokens and positive length.") + if len(alphabet) ** length > 65536: + raise ValueError("Use at most 65,536 candidates for exhaustive exposure.") + return [list(x) for x in product(alphabet, repeat=length)] + + +def random_sequence_probes(seed: int, doses=(0, 1, 4, 16, 64), per_dose=32, + vocab_size=50257, context_tokens=128, target_tokens=64) -> list[dict]: + if per_dose < 1 or vocab_size < 3 or min(context_tokens, target_tokens) < 1: + raise ValueError("Invalid probe dimensions.") + if len(set(doses)) != len(doses) or any(d < 0 or int(d) != d for d in doses): + raise ValueError("Doses must be unique nonnegative integers.") + rng = np.random.default_rng(seed) + probes = [] + seen = set() + for dose in doses: + for index in range(per_dose): + while True: + # Exclude GPT-2 EOT; sequences are token-random, not character-random. + context = rng.integers(0, vocab_size - 1, context_tokens).tolist() + target = rng.integers(0, vocab_size - 1, target_tokens).tolist() + if tuple(target) not in seen: + seen.add(tuple(target)) + break + probes.append({"id": f"random_d{dose}_{index:03d}", "family": "random_tokens", + "dose": int(dose), "context": context, "target": target}) + return probes + + +def presentation_schedule(probes: list[dict], slots: int, seed: int) -> list[dict]: + """One slot is one COMPLETE context+target training record, not a bin offset. + + Allocate every declared presentation without replacement. The adapter must + log realized visits and verify these counts; it must never sample windows + through the injected records and call intended copies actual exposures. + """ + if slots < 1 or len({p['id'] for p in probes}) != len(probes): + raise ValueError("Require positive slots and unique probe IDs.") + for p in probes: + if int(p["dose"]) != p["dose"] or p["dose"] < 0: + raise ValueError("Doses must be nonnegative integers.") + ids = [p["id"] for p in probes for _ in range(int(p["dose"]))] + if len(ids) > slots: + raise ValueError("Not enough training slots for the declared presentation counts.") + rng = np.random.default_rng(seed) + chosen = rng.choice(slots, size=len(ids), replace=False) + return sorted(({'slot': int(s), 'probe_id': p} for s, p in zip(chosen, ids)), key=lambda r: r['slot']) diff --git a/baseline/experiments/nanogpt_memorization/study.json b/baseline/experiments/nanogpt_memorization/study.json new file mode 100644 index 0000000..622ed78 --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/study.json @@ -0,0 +1,80 @@ +{ + "schema_version": 1, + "name": "nanogpt_memorization_2026_09_14", + "status": "designed_not_trained", + "source": { + "repository": "CalculatedContent/rg_optimizers", + "commit": "3749c36334382a20e48bfe2473c1dc4a1470a830", + "recipe": "baseline/experiments/nanogpt_one_head_2026_08_21_baseline/configs/baseline.yaml", + "recipe_git_blob_sha": "7fd3c592afc7fdfea952b1f1aa0f9b44cdc53a2b", + "qualification_status": "source_backed_center_not_frozen_optimum", + "modules": { + "model.py": "029b42675ce230c6f1a142afaa8730b9f4d4ced3", + "optimizers.py": "b1a97ad684afd882ffcff4ae0fe5d70140e7d22c", + "spectral.py": "856a12ecad29bc85582a10601e5708b29e2cbb8a" + } + }, + "optimizers": ["adamw", "muon"], + "historical_comparator_not_active": "muon_clip", + "seeds": [1337, 2027, 4099, 31415, 271828], + "data_seed": 20260914, + "conditions": ["clean", "sequence", "association"], + "training": { + "reference_tokens": 80000000, + "target_reference_epochs": 4, + "schedule_reference_epochs": 1, + "early_stopping": false, + "protected_test_used_for_selection": false, + "checkpoint_policy": "atomic_latest_plus_permanent_diagnostics", + "preserve_baseline_model_optimizer_precision_and_batch": true + }, + "probes": { + "doses_per_reference_epoch": [0, 1, 4, 16, 64], + "sequence_targets_per_dose_and_family": 32, + "sequence_families": ["natural_reserved", "random_tokens"], + "prefix_tokens": [8, 16, 32, 64, 128], + "continuation_tokens": [16, 32, 64], + "primary_prefix_tokens": 64, + "primary_continuation_tokens": 32, + "canaries_per_dose": 8, + "canary_alphabet_size": 16, + "canary_suffix_tokens": 3, + "canary_universe_size": 4096, + "canary_exposure_bits_max": 12, + "association_entities_per_dose": 32, + "association_training_templates_per_entity": 2, + "association_heldout_templates_per_entity": 1, + "association_controls": ["unseen_keys", "permuted_values", "heldout_rule_combinations"], + "membership_controls": "same_target_in_paired_clean_run" + }, + "monitoring": { + "behavior_every_steps": 500, + "spectral_early_steps": [0, 1, 10, 100, 500], + "spectral_every_reference_epochs": 0.25, + "include_warmup_end": true, + "full_exposure_reference_epochs": [0, 1, 2, 4], + "prefix_sweep_reference_epochs": [0, 0.25, 1, 2, 4], + "fix_fingers": "clip_xmax", + "ERG": true, + "randomize": true, + "min_evals": 20, + "max_fingers": 10, + "require_raw_alpha": true, + "missing_or_failed_fit": "record_failure_never_impute" + }, + "retention": { + "enabled_in_initial_campaign": false, + "parent_condition": "sequence", + "fork_reference_epoch": 1, + "children": ["continued_exposure", "withdrawn_exposure"], + "restore_full_optimizer_and_rng_state": true, + "never_restart_schedule": true + }, + "analysis": { + "replication_unit": "paired_complete_training_seed", + "primary_endpoint": "fixed_budget_EM32_P64_sequence_and_canary_exposure_and_association_recall", + "secondary_comparison": "common_validation_NLL_range_only", + "spectral_relation": "exploratory_not_a_causal_memorization_detector", + "alpha_interpretation": "requires_tail_fit_and_correlation_evidence_not_layer_average_alone" + } +} diff --git a/baseline/experiments/nanogpt_memorization/tests/test_study.py b/baseline/experiments/nanogpt_memorization/tests/test_study.py new file mode 100644 index 0000000..e067cf7 --- /dev/null +++ b/baseline/experiments/nanogpt_memorization/tests/test_study.py @@ -0,0 +1,224 @@ +from __future__ import annotations +from collections import Counter +import copy +import json +from pathlib import Path +from types import SimpleNamespace +import sys +import numpy as np +import pytest +import torch +import yaml + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) +from metrics import recall, exposure, prefix_compression, score_continuation, rank_canary +from probes import canary_universe, random_sequence_probes, presentation_schedule +from prepare_plan import resolve, git_blob_sha, write_plan +from monitor import monitor_training_state + + +class TinyModel(torch.nn.Module): + """Controlled transition model, not nanoGPT performance evidence.""" + def __init__(self): + super().__init__() + self.anchor = torch.nn.Parameter(torch.zeros(())) + self.cfg = SimpleNamespace(block_size=16) + + def forward(self, idx): + logits = torch.nn.functional.one_hot((idx + 1) % 8, 8).float() * 12 + return logits, None + + def generate_greedy(self, prompts, max_new_tokens): + ids = prompts.clone() + for _ in range(max_new_tokens): + ids = torch.cat((ids, (ids[:, -1:] + 1) % 8), dim=1) + return ids + + +def test_exact_and_partial_recall_are_distinct(): + row = recall([1, 8, 3, 4], [1, 2, 3, 4]) + assert not row['sequence_exact_match'] + assert row['free_running_token_match'] == .75 + assert row['longest_exact_prefix'] == 1 + assert recall([1, 2], [1, 2])['sequence_exact_match'] + + +@pytest.mark.parametrize('g,t', [([], []), ([1], [1, 2]), ([[1]], [[1]])]) +def test_recall_rejects_malformed(g, t): + with pytest.raises(ValueError): + recall(g, t) + + +def test_exact_exposure(): + row = exposure([8, 3, 5, 9], 1) + assert row['rank_min'] == row['rank_max'] == 1 + assert row['exposure_bits_lower'] == 2 + assert exposure([8, 3, 5, 9], 3)['exposure_bits_lower'] == 0 + + +def test_exposure_ties_do_not_invent_memorization(): + row = exposure([2, 2, 2, 2], 0) + assert row['exposure_bits_lower'] == 0 + assert row['exposure_bits_upper'] == 2 + assert row['tie_count'] == 4 + + +@pytest.mark.parametrize('scores,index', [([], 0), ([1, float('nan')], 0), ([1], 1), ([float('inf')], 0)]) +def test_exposure_rejects_missing_scores(scores, index): + with pytest.raises(ValueError): + exposure(scores, index) + + +def test_prefix_sweep_fixed_boundary_and_nonmonotone_success(): + observed = [] + def generate(p, n): + observed.append(p) + return [9, 9] if len(p) in (1, 3) else [8, 8] + row = prefix_compression([1, 2, 3], [9, 9], [1, 2, 3], generate) + assert observed == [[3], [2, 3], [1, 2, 3]] + assert row['shortest_successful_prefix_in_grid'] == 1 + assert row['prefix_compression_ratio'] == 2 + + +def test_prefix_failure_is_censored(): + row = prefix_compression([1, 2], [5, 6], [1, 2], lambda p, n: [0, 0]) + assert row['search_censored'] + assert row['prefix_compression_ratio'] is None + + +@pytest.mark.parametrize('grid', [[], [0], [3], [1.5]]) +def test_bad_prefix_grid(grid): + with pytest.raises(ValueError): + prefix_compression([1, 2], [5], grid, lambda p, n: [5]) + + +def test_suffix_alignment_and_teacher_forcing_distinction(): + model = TinyModel().train() + good = score_continuation(model, [0, 1], [2, 3]) + assert good['sequence_exact_match'] and good['suffix_nll'] < .01 + bad = score_continuation(model, [0, 1], [2, 5, 6]) + assert bad['teacher_forced_token_accuracy'] == pytest.approx(2 / 3) + assert bad['free_running_token_match'] == pytest.approx(1 / 3) + assert bad['longest_exact_prefix'] == 1 + assert model.training + + +def test_score_preserves_eval_mode_and_rejects_overflow(): + model = TinyModel().eval() + score_continuation(model, [1], [2]) + assert not model.training + with pytest.raises(ValueError): + score_continuation(model, [1] * 16, [2, 3]) + + +def test_candidate_ranking_and_batching(): + model = TinyModel().train() + candidates = [[2, 3], [4, 3], [2, 5], [4, 5]] + a = rank_canary(model, [0, 1], candidates, 0, batch_size=1) + b = rank_canary(model, [0, 1], candidates, 0, batch_size=3) + assert a['rank_min'] == 1 and a['exposure_bits_lower'] == 2 + assert a['candidate_nll_sum'] == pytest.approx(b['candidate_nll_sum']) + assert model.training + + +def test_candidate_duplicates_rejected(): + with pytest.raises(ValueError): + rank_canary(TinyModel(), [1], [[2], [2]], 0) + + +def test_canary_universe(): + values = canary_universe(list(range(16)), 3) + assert len(values) == len({tuple(x) for x in values}) == 4096 + assert values[0] == [0, 0, 0] and values[-1] == [15, 15, 15] + + +def test_random_probes_are_paired_and_do_not_touch_global_rng(): + np.random.seed(92) + before = np.random.get_state() + a = random_sequence_probes(11, per_dose=2) + b = random_sequence_probes(11, per_dose=2) + after = np.random.get_state() + assert a == b + assert np.array_equal(before[1], after[1]) + assert before[2:] == after[2:] + assert len({tuple(x['target']) for x in a}) == len(a) + + +def test_schedule_counts_full_presentations_and_excludes_zero_dose(): + probes = random_sequence_probes(2, per_dose=2) + rows = presentation_schedule(probes, 1000, 3) + counts = Counter(row['probe_id'] for row in rows) + assert all(counts[p['id']] == p['dose'] for p in probes) + assert len({row['slot'] for row in rows}) == len(rows) + assert rows == presentation_schedule(probes, 1000, 3) + + +def test_schedule_rejects_insufficient_slots(): + with pytest.raises(ValueError): + presentation_schedule([{'id': 'a', 'dose': 2}], 1, 0) + + +def config_fixture(tmp_path): + study = json.loads((ROOT / 'study.json').read_text()) + base = {'protocol': {}, 'dataset': {'train_tokens': 80000000}, + 'model': {'block_size': 256}, + 'training': {'batch_size': 4, 'grad_accum_steps': 8, 'seeds': study['seeds']}, + 'optimizer_profiles': {'adamw': {'family': 'adamw', 'learning_rate': .0006}, + 'muon': {'family': 'muon', 'matrix_learning_rate': .02}}, + 'weightwatcher': {k: study['monitoring'][k] for k in + ('fix_fingers', 'ERG', 'randomize', 'min_evals', 'max_fingers', 'require_raw_alpha')}} + path = tmp_path / study['source']['recipe'] + path.parent.mkdir(parents=True) + raw = yaml.safe_dump(base).encode() + path.write_bytes(raw) + study['source']['recipe_git_blob_sha'] = git_blob_sha(raw) + study['source']['modules'] = {} + return study, path + + +def test_plan_counts_horizon_and_optimizer_inheritance(tmp_path): + study, _ = config_fixture(tmp_path) + plan = resolve(tmp_path, study) + assert len(plan['runs']) == 30 + assert plan['tokens_per_update'] == 8192 + assert plan['schedule_steps'] == 9766 + assert plan['total_steps'] == 39063 + assert plan['actual_target_tokens_per_run'] == 320004096 + assert plan['runs'][0]['baseline_config']['optimizer_profiles']['adamw']['learning_rate'] == .0006 + assert plan['runs'][1]['baseline_config']['optimizer_profiles']['muon']['matrix_learning_rate'] == .02 + assert plan['execution_status'].startswith('plan_only') + assert plan['fingerprint'] == resolve(tmp_path, study)['fingerprint'] + + +def test_source_drift_fails_closed(tmp_path): + study, path = config_fixture(tmp_path) + path.write_text(path.read_text() + '\n# drift\n') + with pytest.raises(RuntimeError, match='changed'): + resolve(tmp_path, study) + + +def test_plan_refuses_implicit_muonclip_substitution(tmp_path): + study, _ = config_fixture(tmp_path) + study['optimizers'] = ['adamw', 'muon_clip'] + with pytest.raises(ValueError, match='MuonClip'): + resolve(tmp_path, study) + + +def test_plan_refuses_existing_output_directory(tmp_path): + study, _ = config_fixture(tmp_path) + plan = resolve(tmp_path, study) + out = tmp_path / 'output' + write_plan(plan, out) + assert len(list((out / 'configs').glob('*.yaml'))) == 30 + with pytest.raises(FileExistsError): + write_plan(plan, out) + + +@pytest.mark.parametrize('key,value', [('fix_fingers', False), ('ERG', False), ('randomize', False), ('strict', False)]) +def test_monitor_refuses_invalid_settings_before_import(key, value): + cfg = dict(fix_fingers='clip_xmax', ERG=True, randomize=True, strict=True, enabled=True, require_raw_alpha=True) + cfg[key] = value + with pytest.raises(ValueError): + monitor_training_state(None, '/tmp/not-created', step=0, tokens_seen=0, + reference_tokens=80000000, seed=1337, fingerprint='x', ww_config=cfg)