From d453f64f7f4951f45d5944d98c8771e52eb3389d Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 22 Aug 2026 20:52:48 -0700 Subject: [PATCH] =?UTF-8?q?Add=20reproducible=20AdamW=E2=80=93MuonClip=20n?= =?UTF-8?q?anoGPT=20WeightWatcher=20campaign?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 37 +- baseline/.gitignore | 14 + baseline/README.md | 38 +- .../README.md | 319 ++ .../RESULTS.md | 15 + .../campaign.yaml | 13 + .../configs/baseline.yaml | 145 + .../environment/README.md | 27 + .../01_Performance_and_Spectra.ipynb | 176 + .../runs/README.md | 11 + .../scripts/build_report.py | 3159 ++++++++++++++++ .../scripts/run_experiment.py | 3269 +++++++++++++++++ baseline/nanogpt_one_head/README.md | 25 +- baseline/nanogpt_one_head/TPU.md | 5 +- .../notebooks/01_sgd_momentum_baseline.ipynb | 2 +- .../notebooks/02_adamw_baseline.ipynb | 2 +- .../notebooks/03_muon_baseline.ipynb | 2 +- .../notebooks/04_compare_baselines.ipynb | 2 +- .../notebooks/05_muonclip_esd_clip_xmax.ipynb | 108 +- baseline/nanogpt_one_head/pyproject.toml | 3 +- baseline/nanogpt_one_head/requirements.txt | 4 +- .../src/rg_nanogpt_one_head/analysis.py | 39 +- .../src/rg_nanogpt_one_head/checkpoints.py | 253 +- .../src/rg_nanogpt_one_head/completion.py | 689 +++- .../src/rg_nanogpt_one_head/config.py | 48 +- .../src/rg_nanogpt_one_head/data.py | 19 +- .../src/rg_nanogpt_one_head/doctor_smoke.py | 554 +++ .../src/rg_nanogpt_one_head/engine.py | 79 +- .../src/rg_nanogpt_one_head/evaluation.py | 56 +- .../src/rg_nanogpt_one_head/monitor.py | 42 +- .../src/rg_nanogpt_one_head/muonclip.py | 41 +- .../src/rg_nanogpt_one_head/optimizers.py | 7 +- .../src/rg_nanogpt_one_head/provenance.py | 183 + .../src/rg_nanogpt_one_head/run_utils.py | 257 +- .../src/rg_nanogpt_one_head/runtime.py | 176 +- .../src/rg_nanogpt_one_head/spectral.py | 395 +- .../src/rg_nanogpt_one_head/train_loop.py | 128 +- .../tests/test_20260821_campaign.py | 909 +++++ .../nanogpt_one_head/tests/test_completion.py | 431 ++- .../nanogpt_one_head/tests/test_muonclip.py | 51 +- .../tests/test_muonclip_integration.py | 2 + .../tests/test_muonclip_walk.py | 2 + .../nanogpt_one_head/tests/test_one_head.py | 91 +- 43 files changed, 11494 insertions(+), 334 deletions(-) create mode 100644 baseline/experiments/nanogpt_one_head_2026_08_21_baseline/README.md create mode 100644 baseline/experiments/nanogpt_one_head_2026_08_21_baseline/RESULTS.md create mode 100644 baseline/experiments/nanogpt_one_head_2026_08_21_baseline/campaign.yaml create mode 100644 baseline/experiments/nanogpt_one_head_2026_08_21_baseline/configs/baseline.yaml create mode 100644 baseline/experiments/nanogpt_one_head_2026_08_21_baseline/environment/README.md create mode 100644 baseline/experiments/nanogpt_one_head_2026_08_21_baseline/notebooks/01_Performance_and_Spectra.ipynb create mode 100644 baseline/experiments/nanogpt_one_head_2026_08_21_baseline/runs/README.md create mode 100644 baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/build_report.py create mode 100644 baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/run_experiment.py create mode 100644 baseline/nanogpt_one_head/src/rg_nanogpt_one_head/doctor_smoke.py create mode 100644 baseline/nanogpt_one_head/src/rg_nanogpt_one_head/provenance.py create mode 100644 baseline/nanogpt_one_head/tests/test_20260821_campaign.py diff --git a/README.md b/README.md index b9b46979..bcab5cdb 100644 --- a/README.md +++ b/README.md @@ -49,11 +49,14 @@ python -m pip install -e './baseline[experiment]' jupyter lab baseline/notebooks ``` -Long-running outputs should live under `$HOME`, not `/tmp`: +Put generated data and runs beneath an explicit experiment root. The dated +one-head nanoGPT campaign enforces a `/tmp` root and redirects third-party +caches there as well: ```bash -export RG_BASELINE_DATA_DIR="$HOME/rg-optimizer-data" -export RG_BASELINE_RUN_ROOT="$HOME/rg-optimizer-runs" +export RG_BASELINE_EXPERIMENT_ROOT="/tmp/rg-optimizer-baselines" +export RG_BASELINE_DATA_DIR="$RG_BASELINE_EXPERIMENT_ROOT/data" +export RG_BASELINE_RUN_ROOT="$RG_BASELINE_EXPERIMENT_ROOT/runs" ``` ### Recommended notebook order @@ -76,16 +79,26 @@ baseline/notebooks/CIFAR10_ViT_Optimizer_Baselines.ipynb One-head nanoGPT: ```bash -cd baseline/nanogpt_one_head -bash scripts/setup_mac.sh -bash scripts/prepare_data.sh -bash scripts/smoke_test.sh - -export RG_NANOGPT_ONE_HEAD_ROOT="$HOME/rg-nanogpt-one-head" -caffeinate -dimsu bash scripts/run_all_baselines.sh \ - 2>&1 | tee "$RG_NANOGPT_ONE_HEAD_ROOT/run_all.log" +cd baseline/experiments/nanogpt_one_head_2026_08_21_baseline +export RG_NANOGPT_EXPERIMENT_ROOT="/tmp/rg-nanogpt-one-head-20260821" +mkdir -p "$RG_NANOGPT_EXPERIMENT_ROOT"/{cache/{home,pip,xdg/{cache,config,data,state},matplotlib},tmp} +export HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/home" +export PIP_CACHE_DIR="$RG_NANOGPT_EXPERIMENT_ROOT/cache/pip" +export XDG_CACHE_HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/cache" +export XDG_CONFIG_HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/config" +export XDG_DATA_HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/data" +export XDG_STATE_HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/state" +export MPLCONFIGDIR="$RG_NANOGPT_EXPERIMENT_ROOT/cache/matplotlib" +export TMPDIR="$RG_NANOGPT_EXPERIMENT_ROOT/tmp" +python -m pip install -e ../../nanogpt_one_head +python scripts/run_experiment.py doctor --device mps +python scripts/run_experiment.py prepare +caffeinate -dimsu python scripts/run_experiment.py run --device mps ``` +The full protocol, exact commands, reporting contract, and honest results ledger +are in the [dated experiment folder](baseline/experiments/nanogpt_one_head_2026_08_21_baseline/README.md). + nanochat: ```text @@ -200,7 +213,7 @@ The committed values above are strong source-backed centers. They become frozen best baselines only after the bounded qualification protocol: 1. screen the preregistered neighborhood using validation data only; -2. run the finalists with the complete three-seed protocol; +2. run the finalists with the complete five-seed protocol; 3. select the lowest mean best-validation loss; 4. write the full winning configuration and evidence to a lock file; 5. inspect protected-test comparisons only after the lock exists. diff --git a/baseline/.gitignore b/baseline/.gitignore index 7f46a370..46630584 100644 --- a/baseline/.gitignore +++ b/baseline/.gitignore @@ -3,3 +3,17 @@ __pycache__/ .ipynb_checkpoints/ data/ runs/ + +# The dated nanoGPT campaign deliberately archives only small, reviewed run +# records in the repository. Keep ordinary/raw run directories ignored while +# allowing this one experiment's README and timestamped archive directories. +!experiments/nanogpt_one_head_2026_08_21_baseline/runs/ +!experiments/nanogpt_one_head_2026_08_21_baseline/runs/README.md +!experiments/nanogpt_one_head_2026_08_21_baseline/runs/*/ +!experiments/nanogpt_one_head_2026_08_21_baseline/runs/*/** +experiments/nanogpt_one_head_2026_08_21_baseline/runs/**/*.pt +experiments/nanogpt_one_head_2026_08_21_baseline/runs/**/*.bin +experiments/nanogpt_one_head_2026_08_21_baseline/runs/**/*.partial* +experiments/nanogpt_one_head_2026_08_21_baseline/runs/.*.partial-*/ +experiments/nanogpt_one_head_2026_08_21_baseline/runs/**/cache/ +experiments/nanogpt_one_head_2026_08_21_baseline/runs/**/logs/ diff --git a/baseline/README.md b/baseline/README.md index 0daa1a90..0dcfec5c 100644 --- a/baseline/README.md +++ b/baseline/README.md @@ -71,21 +71,20 @@ When installing from the repository root, use: python -m pip install -e './baseline[experiment]' ``` -Set persistent locations before running long jobs: +Set one explicit experiment location before running long jobs: ```bash -export RG_BASELINE_DATA_DIR="$HOME/rg-optimizer-data" -export RG_BASELINE_RUN_ROOT="$HOME/rg-optimizer-runs" +export RG_BASELINE_EXPERIMENT_ROOT="/tmp/rg-optimizer-baselines" +export RG_BASELINE_DATA_DIR="$RG_BASELINE_EXPERIMENT_ROOT/data" +export RG_BASELINE_RUN_ROOT="$RG_BASELINE_EXPERIMENT_ROOT/runs" ``` -The isolated one-head nanoGPT suite uses: +The dated one-head nanoGPT campaign uses and validates: ```bash -export RG_NANOGPT_ONE_HEAD_ROOT="$HOME/rg-nanogpt-one-head" +export RG_NANOGPT_EXPERIMENT_ROOT="/tmp/rg-nanogpt-one-head-20260821" ``` -Do not put long-running results in `/tmp`. - ## 1. MNIST / MLP3 Run these notebooks in order: @@ -243,15 +242,26 @@ that produced that checkpoint. Mac workflow: ```bash -cd nanogpt_one_head -bash scripts/setup_mac.sh -bash scripts/prepare_data.sh -bash scripts/smoke_test.sh - -caffeinate -dimsu bash scripts/run_all_baselines.sh \ - 2>&1 | tee "$RG_NANOGPT_ONE_HEAD_ROOT/run_all.log" +cd experiments/nanogpt_one_head_2026_08_21_baseline +export RG_NANOGPT_EXPERIMENT_ROOT="/tmp/rg-nanogpt-one-head-20260821" +mkdir -p "$RG_NANOGPT_EXPERIMENT_ROOT"/{cache/{home,pip,xdg/{cache,config,data,state},matplotlib},tmp} +export HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/home" +export PIP_CACHE_DIR="$RG_NANOGPT_EXPERIMENT_ROOT/cache/pip" +export XDG_CACHE_HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/cache" +export XDG_CONFIG_HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/config" +export XDG_DATA_HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/data" +export XDG_STATE_HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/state" +export MPLCONFIGDIR="$RG_NANOGPT_EXPERIMENT_ROOT/cache/matplotlib" +export TMPDIR="$RG_NANOGPT_EXPERIMENT_ROOT/tmp" +python -m pip install -e ../../nanogpt_one_head +python scripts/run_experiment.py doctor --device mps +python scripts/run_experiment.py prepare +caffeinate -dimsu python scripts/run_experiment.py run --device mps ``` +See the [dated experiment README](experiments/nanogpt_one_head_2026_08_21_baseline/README.md) +for the AdamW/MuonClip five-seed protocol, reporting, and archive steps. + ## 4. nanochat Run: diff --git a/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/README.md b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/README.md new file mode 100644 index 00000000..c451c47b --- /dev/null +++ b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/README.md @@ -0,0 +1,319 @@ +# One-head nanoGPT AdamW/MuonClip baseline — prepared 2026-08-22 + +This is the reproducible AdamW / MuonClip campaign for the +one-block, one-head nanoGPT model. It is deliberately separate from the +historical SGD/AdamW/Muon notebook suite. + +The campaign runs five paired seeds (`1337`, `2027`, `4099`, `31415`, +`271828`) for each of two optimizers. Every seed uses the same initialization +convention, sampled training windows, and fixed evaluation probes. The unit of +replication is one complete seeded run; layers and checkpoints are repeated +measurements, not extra samples. Before analysis or archive, the launcher hashes +the exact step-zero model tensor inventory and requires equality across both +optimizer arms for each seed. + +## What is being run + +| Item | Frozen value | +|---|---:| +| FineWeb-Edu revision | `593b3a867298afb8ce42625a270ef20ddcad28f9` | +| Train / validation / test | 80M / 1M / 1M tokens | +| Model | 1 block, 1 head, width 128, context 256 | +| Unique parameters | 6,662,656 | +| Effective batch | 8,192 tokens/update | +| Horizon | 4 corpus-equivalent epochs, about 320M sampled tokens | +| Optimizer steps | 39,063 | +| Permanent states | 17: epoch 0, 0.25, ..., 4.0 | +| Seeds | 1337, 2027, 4099, 31415, 271828 | +| Arms | AdamW, MuonClip | + +The learning-rate schedule completes the established one-epoch warm-up/cosine +recipe and then holds the LR floor for three further epochs. That makes this a +combined optimization and spectral-relaxation baseline. It is not presented as +a performance-optimal 320M-token schedule. + +The committed hyperparameters are source-backed centers already used by the +repository. There is no checked-in nanoGPT qualification lock, so the results +must be called **baseline results**, not globally optimized winners. A future +validation-only bounded search can qualify replacements without changing this +record. + +## WeightWatcher: one call, two alphas + +At every permanent state, one CPU copy containing all six transformer matrices +is analyzed once: + +```python +watcher.analyze( + ERG=True, + randomize=True, + plot=False, + min_evals=20, + fix_fingers="clip_xmax", + max_fingers=10, +) +``` + +With pinned WeightWatcher 0.7.7, `alpha` is the finger-corrected exponent and +`raw_alpha` is the exponent before finger removal. The persisted canonical +columns are `alpha_clip_xmax` and `alpha_raw`; `alpha_clip_xmax` is primary and +`alpha_raw` is the required sensitivity curve. WeightWatcher is not run twice. +This interpretation follows the Calculated Content +[clip-Xmax/raw-alpha description](https://calculatedcontent.com/2024/01/29/evaluating-llms-with-weightwatcher-part-iii-the-magic-of-mistral-a-story-of-dragon-kings/). + +The six matrices are `W_Q`, `W_K`, `W_V`, `W_O`, `W_MLP_IN`, and +`W_MLP_OUT`. The much larger token embedding / tied language head is reported +separately in the model parameter count and is not silently mixed into the +six-matrix alpha summary. + +The raw per-matrix tables also retain WeightWatcher's `ERG_gap`, `num_traps`, +`detX_num`, `detX_val`, and `rand_distance` fields. The report produces a +separate ERG-gap/correlation-trap trajectory plot for every optimizer and +matrix; these values are never reconstructed from a proxy statistic. + +## Metrics and their exact meaning + +- Loss is mean next-token cross-entropy in nats/token on a fixed probe. +- Perplexity is `exp(loss)` with no hidden clipping. +- “Accuracy” is next-token top-1 token accuracy, not classification accuracy. +- Top-5 next-token accuracy and bits/token are recorded as diagnostics. +- Train and validation probes each use 64 fixed batches, or 65,536 target + tokens, during training. They are fixed-probe estimates, not exhaustive split + scans. +- The test probe remains untouched during optimization. After the fixed horizon, + it is evaluated once for `checkpoint_final.pt` and once for the + validation-selected `checkpoint_best.pt`, using the same 65,536-target-token + probe. +- BLEU is a secondary post-training lexical-overlap diagnostic on 64 fixed + greedy continuations. Continuation token accuracy and exact match are also + recorded for those two checkpoints. None of these test/generation diagnostics + select a checkpoint, optimizer, horizon, or hyperparameter. + +The report computes a descriptive plateau flag from validation only. A run is +called plateau-like when its validation NLL changes by at most 0.01 nats/token +over each of the last two complete one-epoch intervals. Every arm still runs to +the same fixed four-epoch budget. + +## Never HOME: required environment + +Use one explicit `/tmp` root. The launcher rejects a missing, relative, home, or +non-`/tmp` root and redirects Hugging Face, tiktoken, Matplotlib, Jupyter, +IPython, Torch, XDG, and the child-process `HOME` beneath it. + +```bash +cd baseline/experiments/nanogpt_one_head_2026_08_21_baseline +export RG_NANOGPT_EXPERIMENT_ROOT="/tmp/rg-nanogpt-one-head-20260821" +mkdir -p "$RG_NANOGPT_EXPERIMENT_ROOT"/{cache/{home,pip,xdg/{cache,config,data,state},matplotlib},tmp} +export HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/home" +export PIP_CACHE_DIR="$RG_NANOGPT_EXPERIMENT_ROOT/cache/pip" +export XDG_CACHE_HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/cache" +export XDG_CONFIG_HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/config" +export XDG_DATA_HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/data" +export XDG_STATE_HOME="$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/state" +export MPLCONFIGDIR="$RG_NANOGPT_EXPERIMENT_ROOT/cache/matplotlib" +export TMPDIR="$RG_NANOGPT_EXPERIMENT_ROOT/tmp" +export PYTORCH_ENABLE_MPS_FALLBACK=1 +``` + +Budget at least **20 GB free** beneath that `/tmp` root, and more if pip/model +caches are retained. The 17 permanent model-only checkpoints are about 26.7 MB +each, or roughly 9.1 GB across twenty runs, before latest/best/final +optimizer-state checkpoints, the approximately 164 MB token corpus, logs, +caches, and reports. + +Use the currently activated conda environment. Install the package dependencies +once into that environment: + +```bash +python -m pip install -e ../../nanogpt_one_head +``` + +Commit the experiment setup and core changes before the production preflight. +The launcher deliberately refuses an untracked config or dirty source tree, so +every real run has an immutable source commit rather than a workspace-only +configuration. + +## Preflight and corpus preparation + +```bash +python scripts/run_experiment.py doctor --device mps +python scripts/run_experiment.py prepare +``` + +`doctor` is a real backend gate. On the requested accelerator it runs a tiny +one-head forward/backward/update for AdamW and MuonClip, round-trips each +optimizer through a schema-v5 restart checkpoint, and performs one +`fix_fingers="clip_xmax"` WeightWatcher call over all six matrices. The +production run is not started if any numerical, device, checkpoint, raw-alpha, +or clipped-alpha check fails. `run` revalidates that the successful doctor +artifact belongs to the same source commit, config, complete dependency +closure, campaign root, and hardware block; rerun `doctor` after any of those +change. + +`prepare` reuses a cache only after verifying the pinned dataset identity, +split, tokenizer, vocabulary, EOT token, exact byte counts, and SHA-256 for all +three token files. In the sandbox used to prepare this protocol there was no +existing nanoGPT corpus under `/tmp`, so no full training run was fabricated. + +## Run commands + +Run one overnight replicate on the Mac: + +```bash +caffeinate -dimsu python scripts/run_experiment.py run \ + --optimizers adamw,muon_clip \ + --seeds 1337 \ + --device mps +``` + +Then build a clearly marked provisional report to inspect seed 1337 before +committing the Mac to the other four seeds: + +```bash +python scripts/run_experiment.py analyze --allow-incomplete +``` + +This provisional report has no across-seed uncertainty claim. Use its loss, +raw-alpha, `clip_xmax`-alpha, ERG-gap, and trap trajectories to decide whether +the four-epoch horizon has reached a sufficiently stable late regime. The run +budget remains frozen; do not select a shorter per-optimizer stopping point. + +Run or resume the complete 2 × 5 campaign: + +```bash +caffeinate -dimsu python scripts/run_experiment.py run --device mps +``` + +The runner streams every training row and each checkpoint's clipped and raw +median alpha to the terminal while retaining a per-replicate log. From another +terminal, a specific overnight run can be followed without creating files: + +```bash +tail -f "$RG_NANOGPT_EXPERIMENT_ROOT/logs/runs/adamw/seed_1337.log" +``` + +For a live table that explicitly shows both raw and `clip_xmax` alpha for every +matrix at the latest permanent state: + +```bash +python scripts/run_experiment.py monitor \ + --optimizer adamw \ + --seed 1337 \ + --interval 30 +``` + +Use `--once --no-clear` for a single snapshot. The underlying direct command is +`python -m rg_nanogpt_one_head.monitor`; the launcher supplies the campaign's +strict `/tmp` results and cache roots. + +Run a single arm on an H100: + +```bash +python scripts/run_experiment.py doctor --device cuda +python scripts/run_experiment.py prepare +python scripts/run_experiment.py run \ + --optimizers muon_clip \ + --seeds 1337,2027,4099,31415,271828 \ + --device cuda +``` + +Run on a single TPU/XLA device using explicitly ephemeral `/tmp` storage: + +```bash +python -m pip install -e '../../nanogpt_one_head[tpu]' +export RG_NANOGPT_ALLOW_EPHEMERAL_TPU_STORAGE=1 +export RG_NANOGPT_HARDWARE_BLOCK_ID='tpu-homogeneous-block-a' +python scripts/run_experiment.py doctor --device tpu +python scripts/run_experiment.py prepare +python scripts/run_experiment.py run --device tpu +``` + +Replace the example block ID with a stable description of the actual TPU +pool. It is required when the provider does not expose +`TPU_ACCELERATOR_TYPE`. This dated campaign intentionally obeys the strict +`/tmp` output rule even on TPU, so the opt-in above is genuinely ephemeral: +loss of the VM can destroy the corpus and checkpoints. Preserve the campaign +root outside the VM during long runs and run `archive` immediately after the +complete 2 × 5 analysis; the Git run record intentionally excludes the large +raw checkpoints and corpus. + +Use `--device auto` only when an intentional CPU fallback is acceptable. The +advertised Mac workflow names `mps` explicitly so a CPU-only or incompatible +Torch build fails at preflight instead of silently starting a multi-day CPU run. + +Mac/MPS, H100/CUDA, and TPU/XLA runs are separate hardware blocks. Do not pool +seeds from different accelerator types into one confidence interval. +Runtime provenance includes the CUDA driver, UUID, memory and device +properties; Mac model/SoC/memory; or TPU accelerator type. If a platform cannot +report those fields, `doctor` fails and asks for +`RG_NANOGPT_HARDWARE_BLOCK_ID`. Use one stable, descriptive value only for a +genuinely homogeneous device block; changing it intentionally starts a separate +campaign block. + +For parallel homogeneous H100/TPU hosts, set the same explicit block ID before +`doctor` and assign disjoint optimizer/seed subsets. Complete run directories +may then be copied into one aggregation root and revalidated with `status`. +The comparison still requires identical accelerator model/capability, driver, +Torch/dependency closure, config, source, and corpus hashes; it ignores only +host/install-path and physical-device instance fields. Never move a partial +run to another accelerator instance for resume—finish it on the originating +device or start that replicate in a new root. + +The launcher fails nonzero if any requested replicate fails. It resumes from +the last finite atomic checkpoint; if a process dies before the first periodic +restart state, it falls back to the immutable step-zero checkpoint and +truncates partial CSV/spectral/QK rows before replay. It never treats partial +success as a complete campaign. + +Every full-state checkpoint embeds exact model and optimizer-state digests; +every permanent model checkpoint embeds its model digest. WeightWatcher raw +CSVs are bound to the run fingerprint, seed, diagnostic seed, and exact model +state, with a separate raw-file SHA-256 status record. These bindings are +recomputed before reuse, reporting, or archive. + +Supported launcher jobs use adjacent nonblocking file locks. Separate +optimizer/seed jobs can run concurrently on workers that share a campaign root, +but a second process targeting the same replicate fails before it can append +metrics or replace checkpoints. `prepare`, `doctor`, and `analyze` are likewise +single-writer operations. + +MuonClip QK diagnostics must cover steps `500, 1000, ..., 39000, 39063` +exactly. Interval counts must sum to 39,063 and, for this one-block/one-head +model, head observations must equal optimizer steps in each interval. + +## Status, report, and executed notebook + +```bash +python scripts/run_experiment.py status +python scripts/run_experiment.py analyze +``` + +`analyze` requires all ten runs by default and writes aggregate CSVs, +separate AdamW/MuonClip plots, a documented HTML report, a Markdown summary, a +provenance manifest, and an executed notebook below the same `/tmp` root. +`analyze --allow-incomplete` is diagnostic only; `archive` continues to reject +anything other than the exact complete 2 × 5 campaign. + +After inspection, create a small, check-in-ready archive containing manifests, +aggregate tables, figures, HTML, and the executed notebook—but no corpus or +model checkpoints: + +```bash +python scripts/run_experiment.py archive +``` + +The archived run record includes the exact Git commit / describe string, dirty +state, config hash, command, dependency freeze, hardware backend, data hashes, +UTC timestamps, and instructions to reproduce by checking out that commit. +Its requirements lock covers the complete installed campaign dependency +closure and rejects opaque direct/VCS/file origins instead of silently +rewriting them. `verify-lock` must match the recreated Python and package +inventory before training. Large binary packages are not vendored into Git; +preserve the public package-channel configuration or an external wheelhouse +needed to resolve the pinned builds, and replay fails closed if it cannot. + +## Repository state and actual results + +See [RESULTS.md](RESULTS.md). It starts as an honest “not yet run” ledger. The +archive command creates a dated run folder only after a complete report exists; +executed notebooks are never mistaken for source notebooks. diff --git a/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/RESULTS.md b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/RESULTS.md new file mode 100644 index 00000000..1b399163 --- /dev/null +++ b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/RESULTS.md @@ -0,0 +1,15 @@ +# Results ledger + +## 2026-08-22 — protocol revised + +- Status: not run in the preparation sandbox. +- Reason: no full FineWeb-Edu corpus preparation or production training was + performed in the audit sandbox. +- Scientific result: none claimed. +- Next action: run `doctor`, `prepare`, then seed 1337 across AdamW and MuonClip; + inspect loss and alpha convergence before resuming the paired 2 × 5 campaign from a + clean repository commit. + +Actual archived executions belong in `runs/_/` and +must include the generated `RUN_RECORD.md`, provenance JSON, aggregate tables, +plots, report, and executed notebook. diff --git a/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/campaign.yaml b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/campaign.yaml new file mode 100644 index 00000000..05d1a1ad --- /dev/null +++ b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/campaign.yaml @@ -0,0 +1,13 @@ +campaign: + id: nanogpt_one_head_2026_08_21_baseline_v3 + prepared_date: 2026-08-22 + config: configs/baseline.yaml + optimizers: [adamw, muon_clip] + seeds: [1337, 2027, 4099, 31415, 271828] + require_clean_git: true + require_tmp_root: true + require_complete_replicates: 10 + primary_checkpoint_policy: minimum_validation_probe_nll + protected_test_policy: held_out_posthoc_never_selects + primary_weightwatcher_alpha: alpha_clip_xmax + sensitivity_weightwatcher_alpha: alpha_raw diff --git a/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/configs/baseline.yaml b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/configs/baseline.yaml new file mode 100644 index 00000000..7fd3c592 --- /dev/null +++ b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/configs/baseline.yaml @@ -0,0 +1,145 @@ +protocol: + name: nanogpt_one_head_2026_08_21_ww_baseline + version: 3 + description: Four-corpus-equivalent-epoch AdamW and MuonClip source-backed baselines with five paired seeds, one-call WeightWatcher clip_xmax/raw alpha monitoring, and seventeen permanent analysis states. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + +training: + seeds: [1337, 2027, 4099, 31415, 271828] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 4.0 + epoch_interval: 0.25 + eval_interval_steps: 500 + eval_batches: 64 + checkpoint_interval_steps: 500 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum (not a campaign arm) + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.10 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 0.00006 + warmup_fraction: 0.01 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + adam: + display_name: Adam (not a campaign arm) + family: adam + learning_rate: 0.0006 + min_learning_rate: 0.00006 + warmup_fraction: 0.01 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.0 + + muon: + display_name: Muon + auxiliary AdamW (not a campaign arm) + family: muon + matrix_learning_rate: 0.02 + matrix_min_learning_rate: 0.002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00003 + warmup_fraction: 0.05 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.01 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.01 + + muon_clip: + display_name: MuonClip + auxiliary AdamW + family: muon_clip + learning_rate: 0.0002 + min_learning_rate: 0.00002 + warmup_fraction: 0.0512 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + momentum: 0.95 + nesterov: false + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + weight_decay: 0.10 + update_rms_scale: 0.20 + qk_clip_threshold: 100.0 + qk_clip_balance: 0.50 + qk_diagnostics_interval: 500 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + fix_fingers: clip_xmax + max_fingers: 10 + require_raw_alpha: true + +runtime: + matmul_precision: highest + allow_tf32: false + cudnn_benchmark: false + mps_fallback: true + deterministic_algorithms: true + deterministic_warn_only: false + empty_mps_cache_after_weightwatcher: true diff --git a/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/environment/README.md b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/environment/README.md new file mode 100644 index 00000000..3ed1eff8 --- /dev/null +++ b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/environment/README.md @@ -0,0 +1,27 @@ +# Environment locks + +The campaign launcher writes the actual `pip freeze`, Python executable and +version, Torch / accelerator metadata, CUDA/cuDNN or XLA versions, and cache +paths into the `/tmp` provenance directory at run time. The archive command +copies those observed locks here only as part of a dated completed run. + +Launcher child processes also receive an ephemeral `HOME` below the required +campaign root. This catches libraries that ignore their dedicated cache or XDG +variables; production commands do not write experiment data or caches to the +user's real home directory. + +Each archive retains the raw `pip freeze`, creates a version-pinned +`requirements_replay.txt` for the complete installed campaign dependency +closure, and records `dependency_lock.json`. Archive fails on a dependency +installed from an opaque direct/VCS/file origin rather than pretending that a +name/version pin is equivalent. The archived reproduction sequence runs +`verify-lock` after installing the checked-out project and before touching the +corpus or starting training. + +The Git archive does not vendor large wheels or conda packages. Preserve the +public package-channel configuration (or an external wheelhouse) needed to +resolve the recorded builds. Replay fails closed if those artifacts are no +longer resolvable or the recreated dependency closure differs. + +No synthetic Darwin, CUDA, or TPU lock file is checked in before that platform +has actually run the experiment. diff --git a/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/notebooks/01_Performance_and_Spectra.ipynb b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/notebooks/01_Performance_and_Spectra.ipynb new file mode 100644 index 00000000..b9074100 --- /dev/null +++ b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/notebooks/01_Performance_and_Spectra.ipynb @@ -0,0 +1,176 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "campaign-overview", + "metadata": {}, + "source": [ + "# One-head nanoGPT baseline: performance and spectra\n", + "\n", + "This is the source report notebook for the frozen AdamW and MuonClip campaign prepared on 2026-08-22. It does **not** train models or select checkpoints. It invokes the campaign report builder against already-completed run artifacts, then displays its Markdown summary, aggregate tables, and plots.\n", + "\n", + "Execute it through Papermill with explicit `RESULTS_ROOT` and `OUTPUT_ROOT` parameters. Set `REQUIRE_COMPLETE=True` for the preregistered 2 x 5 report." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "papermill-parameters", + "metadata": { + "tags": [ + "parameters" + ] + }, + "outputs": [], + "source": [ + "RESULTS_ROOT = \"\"\n", + "OUTPUT_ROOT = \"\"\n", + "REQUIRE_COMPLETE = True" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "resolve-inputs", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import subprocess\n", + "import sys\n", + "\n", + "import pandas as pd\n", + "from IPython.display import Image, Markdown, display\n", + "\n", + "EXPERIMENT_ID = \"nanogpt_one_head_2026_08_21_baseline\"\n", + "\n", + "\n", + "def locate_experiment_root() -> Path:\n", + " for anchor in (Path.cwd(), *Path.cwd().parents):\n", + " if (\n", + " anchor.name == EXPERIMENT_ID\n", + " and (anchor / \"scripts\" / \"build_report.py\").is_file()\n", + " ):\n", + " return anchor.resolve()\n", + " candidate = (\n", + " anchor / \"baseline\" / \"experiments\" / EXPERIMENT_ID\n", + " )\n", + " if (candidate / \"scripts\" / \"build_report.py\").is_file():\n", + " return candidate.resolve()\n", + " raise FileNotFoundError(\n", + " f\"Could not locate {EXPERIMENT_ID}; execute from the repository \"\n", + " \"or the dated experiment directory.\"\n", + " )\n", + "\n", + "\n", + "def required_absolute_path(value: str, parameter: str) -> Path:\n", + " if not str(value).strip():\n", + " raise ValueError(f\"Papermill parameter {parameter} is required\")\n", + " path = Path(value)\n", + " if not path.is_absolute():\n", + " raise ValueError(f\"{parameter} must be an absolute path: {path}\")\n", + " return path.resolve()\n", + "\n", + "\n", + "EXPERIMENT_PATH = locate_experiment_root()\n", + "RESULTS_PATH = required_absolute_path(RESULTS_ROOT, \"RESULTS_ROOT\")\n", + "OUTPUT_PATH = required_absolute_path(OUTPUT_ROOT, \"OUTPUT_ROOT\")\n", + "if not RESULTS_PATH.is_dir():\n", + " raise FileNotFoundError(f\"Results directory does not exist: {RESULTS_PATH}\")\n", + "print(f\"Experiment: {EXPERIMENT_PATH}\")\n", + "print(f\"Results: {RESULTS_PATH}\")\n", + "print(f\"Report: {OUTPUT_PATH}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "build-report", + "metadata": {}, + "outputs": [], + "source": [ + "report_command = [\n", + " sys.executable,\n", + " str(EXPERIMENT_PATH / \"scripts\" / \"build_report.py\"),\n", + " \"--results-root\",\n", + " str(RESULTS_PATH),\n", + " \"--output-root\",\n", + " str(OUTPUT_PATH),\n", + "]\n", + "if REQUIRE_COMPLETE:\n", + " report_command.append(\"--require-complete\")\n", + "else:\n", + " report_command.append(\"--allow-incomplete\")\n", + "completed = subprocess.run(\n", + " report_command,\n", + " check=True,\n", + " capture_output=True,\n", + " text=True,\n", + ")\n", + "if completed.stdout:\n", + " print(completed.stdout.rstrip())\n", + "if completed.stderr:\n", + " print(completed.stderr.rstrip(), file=sys.stderr)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "display-summary", + "metadata": {}, + "outputs": [], + "source": [ + "summary_paths = sorted(OUTPUT_PATH.rglob(\"SUMMARY.md\"))\n", + "if not summary_paths:\n", + " raise FileNotFoundError(f\"Report builder wrote no SUMMARY.md below {OUTPUT_PATH}\")\n", + "for summary_path in summary_paths:\n", + " display(Markdown(summary_path.read_text(encoding=\"utf-8\")))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "display-tables", + "metadata": {}, + "outputs": [], + "source": [ + "csv_paths = sorted(OUTPUT_PATH.rglob(\"*.csv\"))\n", + "if not csv_paths:\n", + " raise FileNotFoundError(f\"Report builder wrote no CSV tables below {OUTPUT_PATH}\")\n", + "for csv_path in csv_paths:\n", + " label = csv_path.relative_to(OUTPUT_PATH)\n", + " display(Markdown(f\"### `{label}`\"))\n", + " display(pd.read_csv(csv_path).head(20))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "display-plots", + "metadata": {}, + "outputs": [], + "source": [ + "png_paths = sorted(OUTPUT_PATH.rglob(\"*.png\"))\n", + "if not png_paths:\n", + " raise FileNotFoundError(f\"Report builder wrote no PNG plots below {OUTPUT_PATH}\")\n", + "for png_path in png_paths:\n", + " label = png_path.relative_to(OUTPUT_PATH)\n", + " display(Markdown(f\"### `{label}`\"))\n", + " display(Image(filename=str(png_path)))" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/runs/README.md b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/runs/README.md new file mode 100644 index 00000000..e4999643 --- /dev/null +++ b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/runs/README.md @@ -0,0 +1,11 @@ +# Completed run records + +`run_experiment.py archive` creates one immutable, review-ready directory here +after all ten replicates and the report notebook validate successfully. The +directory name is `_`. + +Each record contains the exact protocol and environment provenance, aggregate +tables and figures, the HTML/Markdown report, the executed notebook, a +SHA-256 file manifest, and lightweight per-replicate manifests/test outcomes. +Tokenized corpus files, caches, training logs, and model checkpoints remain +beneath `RG_NANOGPT_EXPERIMENT_ROOT` and are never copied here. diff --git a/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/build_report.py b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/build_report.py new file mode 100644 index 00000000..16a8df68 --- /dev/null +++ b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/build_report.py @@ -0,0 +1,3159 @@ +#!/usr/bin/env python3 +"""Build the audited report for the dated one-head nanoGPT campaign. + +The report builder is intentionally stricter than a generic plotting script. +By default it requires exactly the preregistered 2 x 5 campaign: + + optimizers: adamw, muon_clip + seeds: 1337, 2027, 4099, 31415, 271828 + +It validates completion and matched campaign invariants, preserves the seeded +run as the unit of replication, computes paired seed differences, aggregates +WeightWatcher alphas only after taking the six-matrix median within each run, +and uses validation loss alone for the saturation diagnostic. + +All generated artifacts are written below an explicitly supplied /tmp output +directory. The source results are read-only inputs. +""" + +from __future__ import annotations + +import argparse +from collections import defaultdict +from datetime import datetime, timezone +import hashlib +import html +from itertools import combinations +import json +import math +import os +from pathlib import Path +import re +import subprocess +import sys +from typing import Any, Iterable, Mapping, Sequence + +import numpy as np +import pandas as pd + + +# Importing Matplotlib can create a configuration directory immediately. Keep +# even ``--help`` read-only, then bind its cache to the validated /tmp report +# tree before the first plotting import. +plt: Any | None = None + + +def _initialize_matplotlib(output_root: Path) -> None: + global plt + if plt is not None: + return + cache = output_root.parent / "cache" / "matplotlib-report" + cache.mkdir(parents=True, exist_ok=True) + os.environ["MPLCONFIGDIR"] = str(cache) + os.environ.setdefault("MPLBACKEND", "Agg") + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as pyplot + + plt = pyplot + + +OPTIMIZERS: tuple[str, ...] = ("adamw", "muon_clip") +SEEDS: tuple[int, ...] = (1337, 2027, 4099, 31415, 271828) +MATRIX_TYPES: tuple[str, ...] = ( + "W_Q", + "W_K", + "W_V", + "W_O", + "W_MLP_IN", + "W_MLP_OUT", +) +OPTIMIZER_LABELS: Mapping[str, str] = { + "adamw": "AdamW", + "muon_clip": "MuonClip + auxiliary AdamW", +} +OPTIMIZER_COLORS: Mapping[str, str] = { + "adamw": "#D55E00", + "muon_clip": "#CC79A7", +} +SPLIT_COLORS: Mapping[str, str] = { + "train": "#0072B2", + "val": "#E69F00", + "test": "#009E73", + "other": "#6B7280", +} +T_975_DF1 = 12.7062047364 +T_975_DF2 = 4.3026527297 +T_975_DF3 = 3.1824463053 +T_975_DF4 = 2.7764451052 +SATURATION_DELTA_NATS = 0.01 +MIN_PERMANENT_CHECKPOINTS = 10 +EXPECTED_TOTAL_STEPS = 39_063 +QK_DIAGNOSTIC_INTERVAL = 500 +EXPECTED_QK_STEPS: tuple[int, ...] = ( + *range(QK_DIAGNOSTIC_INTERVAL, EXPECTED_TOTAL_STEPS + 1, QK_DIAGNOSTIC_INTERVAL), + EXPECTED_TOTAL_STEPS, +) +EXPECTED_PERMANENT_STEPS: tuple[int, ...] = ( + 0, + 2_441, + 4_883, + 7_324, + 9_766, + 12_207, + 14_648, + 17_090, + 19_531, + 21_973, + 24_414, + 26_855, + 29_297, + 31_738, + 34_180, + 36_621, + 39_063, +) +FROZEN_CONFIG_SHA256 = ( + "ebbbdfa30efe96b0b0c1c68ae4fc81909361502d89ad336d1181d00fcb85876a" +) +PINNED_WEIGHTWATCHER = "0.7.7" +PINNED_PACKAGE_VERSION = "0.5.1" +SCRIPT_PATH = Path(__file__).resolve() +EXPERIMENT_DIR = SCRIPT_PATH.parents[1] +REPOSITORY_ROOT = SCRIPT_PATH.parents[4] +FROZEN_CONFIG = EXPERIMENT_DIR / "configs" / "baseline.yaml" +MANIFEST_PACKAGES: tuple[str, ...] = ( + "python", + "rg-nanogpt-one-head", + "torch", + "torch-xla", + "numpy", + "pandas", + "scipy", + "PyYAML", + "datasets", + "tiktoken", + "sacrebleu", + "weightwatcher", + "powerlaw", + "papermill", + "packaging", +) + +TEST_RESULT_METRICS: tuple[str, ...] = ( + "test_loss", + "test_perplexity", + "test_bits_per_token", + "test_accuracy", + "test_top5_accuracy", + "test_bleu", + "test_continuation_token_accuracy", + "test_continuation_exact_match", +) + +FINAL_EPOCH_METRICS: tuple[str, ...] = ( + "train_loss", + "val_loss", + "test_loss", + "train_perplexity", + "val_perplexity", + "test_perplexity", + "train_bits_per_token", + "val_bits_per_token", + "test_bits_per_token", + "train_accuracy", + "val_accuracy", + "test_accuracy", + "train_top5_accuracy", + "val_top5_accuracy", + "test_top5_accuracy", + "test_bleu", + "test_continuation_token_accuracy", + "test_continuation_exact_match", +) + +FINITE_EPOCH_METRICS: tuple[str, ...] = ( + "train_loss", + "val_loss", + "train_perplexity", + "val_perplexity", + "train_bits_per_token", + "val_bits_per_token", + "train_accuracy", + "val_accuracy", + "train_top5_accuracy", + "val_top5_accuracy", +) + +HELD_OUT_CURVE_COLUMNS: tuple[str, ...] = ( + "test_loss", + "test_perplexity", + "test_bits_per_token", + "test_accuracy", + "test_top5_accuracy", + "test_bleu", + "test_continuation_token_accuracy", + "test_continuation_exact_match", + "test_generalization_gap", +) + +REQUIRED_RUN_FILES: tuple[str, ...] = ( + "manifest.json", + "run_complete.json", + "metrics.csv", + "epoch_metrics.csv", + "spectral/layers.csv", + "spectral/summary.csv", + "test_results.json", +) + +CSV_ARTIFACTS: tuple[str, ...] = ( + "campaign_runs.csv", + "metrics_all.csv", + "epoch_metrics_all.csv", + "spectral_layers_all.csv", + "spectral_summary_all.csv", + "test_results_all.csv", + "qk_diagnostics_all.csv", + "qk_summary.csv", + "performance_summary.csv", + "paired_seed_differences.csv", + "alpha_run_medians.csv", + "alpha_across_seed_summary.csv", + "saturation_diagnostics.csv", + "saturation_integer_epoch_validation.csv", + "saturation_across_seed_summary.csv", + "checkpoint_sha256.csv", +) + + +class CampaignValidationError(RuntimeError): + """Raised when the input cannot support the requested scientific report.""" + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _requires_complete(args: argparse.Namespace) -> bool: + """Resolve strictness for the CLI and programmatic legacy callers.""" + + if hasattr(args, "require_complete"): + return bool(args.require_complete) + return not bool(getattr(args, "allow_incomplete", False)) + + +def _json_ready(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _json_ready(item) for key, item in value.items()} + if isinstance(value, (list, tuple, set)): + return [_json_ready(item) for item in value] + if isinstance(value, np.generic): + return _json_ready(value.item()) + if isinstance(value, float) and not math.isfinite(value): + return None + if isinstance(value, Path): + return str(value) + return value + + +def _canonical_json(value: Any) -> str: + return json.dumps( + _json_ready(value), + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + default=str, + ) + + +def _sha256(path: Path, chunk_size: int = 4 * 1024 * 1024) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(chunk_size): + digest.update(chunk) + return digest.hexdigest() + + +def _canonical_sha256(value: Any) -> str: + return hashlib.sha256(_canonical_json(value).encode("utf-8")).hexdigest() + + +def _frozen_campaign_config() -> dict[str, Any]: + try: + import yaml + except ImportError as exc: # pragma: no cover - launcher dependency gate + raise CampaignValidationError("PyYAML is required to verify the frozen config") from exc + try: + payload = yaml.safe_load(FROZEN_CONFIG.read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError) as exc: + raise CampaignValidationError( + f"could not read the frozen campaign config {FROZEN_CONFIG}: {exc}" + ) from exc + if not isinstance(payload, dict): + raise CampaignValidationError("the frozen campaign config is not a mapping") + observed = _canonical_sha256(payload) + if observed != FROZEN_CONFIG_SHA256: + raise CampaignValidationError( + "the checked-out campaign config differs from its frozen contract: " + f"observed={observed}, expected={FROZEN_CONFIG_SHA256}" + ) + return payload + + +def _repository_head() -> str: + try: + completed = subprocess.run( + ["git", "-C", str(REPOSITORY_ROOT), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + timeout=30, + ) + except (OSError, subprocess.SubprocessError) as exc: + raise CampaignValidationError(f"could not resolve the source Git commit: {exc}") from exc + commit = completed.stdout.strip() + if not commit: + raise CampaignValidationError("Git returned an empty source commit") + return commit + + +def _atomic_write_text(path: Path, text: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(text, encoding="utf-8") + temporary.replace(path) + + +def _atomic_json(path: Path, payload: Mapping[str, Any]) -> None: + _atomic_write_text( + path, + json.dumps( + _json_ready(payload), + indent=2, + sort_keys=True, + allow_nan=False, + default=str, + ) + + "\n", + ) + + +def _atomic_csv(path: Path, frame: pd.DataFrame) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + frame.to_csv(temporary, index=False) + temporary.replace(path) + + +def _atomic_figure(path: Path, figure: plt.Figure) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(path.stem + ".tmp" + path.suffix) + figure.savefig( + temporary, + dpi=180, + bbox_inches="tight", + format=path.suffix.lstrip("."), + ) + temporary.replace(path) + + +def _is_within(path: Path, parent: Path) -> bool: + try: + path.resolve(strict=False).relative_to(parent.resolve(strict=False)) + return True + except ValueError: + return False + + +def _is_strictly_within(path: Path, parent: Path) -> bool: + resolved_path = path.resolve(strict=False) + resolved_parent = parent.resolve(strict=False) + return resolved_path != resolved_parent and _is_within( + resolved_path, resolved_parent + ) + + +def _validate_paths(results_root: Path, output_root: Path) -> None: + allowed_tmp_roots = { + Path("/tmp").resolve(strict=False), + Path("/private/tmp").resolve(strict=False), + } + for label, path in ( + ("results root", results_root), + ("output root", output_root), + ): + if not any( + _is_strictly_within(path, tmp_root) + for tmp_root in allowed_tmp_roots + ): + raise CampaignValidationError( + f"{label} must be strictly below resolved /tmp or " + f"/private/tmp; observed {path}" + ) + experiment_root_value = os.environ.get("RG_NANOGPT_EXPERIMENT_ROOT", "") + if not experiment_root_value.strip(): + raise CampaignValidationError( + "RG_NANOGPT_EXPERIMENT_ROOT is required for report generation" + ) + experiment_root = Path(experiment_root_value) + if not experiment_root.is_absolute() or "~" in experiment_root_value: + raise CampaignValidationError( + "RG_NANOGPT_EXPERIMENT_ROOT must be an absolute non-tilde path" + ) + experiment_root = experiment_root.resolve(strict=False) + if not any( + _is_strictly_within(experiment_root, tmp_root) + for tmp_root in allowed_tmp_roots + ): + raise CampaignValidationError( + "RG_NANOGPT_EXPERIMENT_ROOT must resolve strictly below /tmp or " + "/private/tmp" + ) + for label, path in ( + ("results root", results_root), + ("output root", output_root), + ): + if not _is_strictly_within(path, experiment_root): + raise CampaignValidationError( + f"{label} must be strictly below RG_NANOGPT_EXPERIMENT_ROOT: {path}" + ) + home_value = os.environ.get("HOME") + if home_value: + home = Path(home_value).resolve(strict=False) + for label, path in ( + ("results root", results_root), + ("output root", output_root), + ): + if _is_within(path, home): + raise CampaignValidationError( + f"{label} must never be HOME or below HOME: {path}" + ) + if not results_root.is_dir(): + raise CampaignValidationError( + f"results root does not exist or is not a directory: {results_root}" + ) + if _is_within(output_root, results_root) or _is_within(results_root, output_root): + raise CampaignValidationError( + "results root and report output root must not contain one another" + ) + + +def _read_json(path: Path) -> dict[str, Any]: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise CampaignValidationError(f"could not read valid JSON {path}: {exc}") from exc + if not isinstance(payload, dict): + raise CampaignValidationError(f"JSON artifact is not an object: {path}") + return payload + + +def _read_csv(path: Path, *, allow_empty: bool = False) -> pd.DataFrame: + try: + frame = pd.read_csv(path) + except Exception as exc: + raise CampaignValidationError(f"could not read CSV {path}: {exc}") from exc + if frame.empty and not allow_empty: + raise CampaignValidationError(f"required CSV is empty: {path}") + return frame + + +def _require_columns(frame: pd.DataFrame, columns: Iterable[str], label: str) -> None: + missing = sorted(set(columns) - set(frame.columns)) + if missing: + raise CampaignValidationError(f"{label} is missing columns: {missing}") + + +def _numeric(frame: pd.DataFrame, columns: Iterable[str]) -> pd.DataFrame: + result = frame.copy() + for column in columns: + if column in result.columns: + result[column] = pd.to_numeric(result[column], errors="coerce") + return result + + +def _with_identity( + frame: pd.DataFrame, + *, + optimizer: str, + seed: int, + source_path: Path, +) -> pd.DataFrame: + result = frame.copy() + if "optimizer" in result.columns: + observed = set(result["optimizer"].dropna().astype(str)) + if observed and observed != {optimizer}: + raise CampaignValidationError( + f"optimizer identity mismatch in {source_path}: {sorted(observed)}" + ) + result = result.drop(columns=["optimizer"]) + if "seed" in result.columns: + observed_seed = set(pd.to_numeric(result["seed"], errors="coerce").dropna().astype(int)) + if observed_seed and observed_seed != {int(seed)}: + raise CampaignValidationError( + f"seed identity mismatch in {source_path}: {sorted(observed_seed)}" + ) + result = result.drop(columns=["seed"]) + result.insert(0, "optimizer", optimizer) + result.insert(1, "optimizer_label", OPTIMIZER_LABELS[optimizer]) + result.insert(2, "seed", int(seed)) + result.insert(3, "source_file", str(source_path)) + return result + + +def _matrix_type(value: Any) -> str: + text = str(value).upper() + for matrix in sorted(MATRIX_TYPES, key=len, reverse=True): + if text == matrix or text.endswith("_" + matrix) or matrix in text: + return matrix + return text + + +def _first_existing_column(frame: pd.DataFrame, names: Sequence[str]) -> str | None: + return next((name for name in names if name in frame.columns), None) + + +def _normalize_spectral_layers(frame: pd.DataFrame, label: str) -> pd.DataFrame: + _require_columns(frame, ("step", "epoch"), label) + result = _numeric(frame, ("step", "tokens_seen", "epoch")) + if "matrix_type" not in result.columns: + if "matrix_name" not in result.columns: + raise CampaignValidationError( + f"{label} contains neither matrix_type nor matrix_name" + ) + result["matrix_type"] = result["matrix_name"].map(_matrix_type) + else: + result["matrix_type"] = result["matrix_type"].map(_matrix_type) + + if "matrix_name" not in result.columns: + result["matrix_name"] = result["matrix_type"] + + clipped_column = _first_existing_column( + result, + ("alpha_clip_xmax", "clip_xmax_alpha", "fixed_alpha", "alpha"), + ) + raw_column = _first_existing_column( + result, + ("alpha_raw", "raw_alpha", "alpha_before_finger_clip"), + ) + if clipped_column is None or raw_column is None: + raise CampaignValidationError( + f"{label} must contain both clipped and raw WeightWatcher alpha; " + f"observed columns={sorted(result.columns)}" + ) + result["alpha_clip_xmax"] = pd.to_numeric( + result[clipped_column], errors="coerce" + ) + result["alpha_raw"] = pd.to_numeric(result[raw_column], errors="coerce") + result["alpha_clip_minus_raw"] = ( + result["alpha_clip_xmax"] - result["alpha_raw"] + ) + return result + + +def _expected_run_directory(results_root: Path, optimizer: str, seed: int) -> Path: + return results_root / optimizer / f"seed_{int(seed)}" + + +def _unexpected_seed_directories(results_root: Path) -> list[str]: + unexpected: list[str] = [] + pattern = re.compile(r"^seed_(-?\d+)$") + expected = set(SEEDS) + for optimizer in OPTIMIZERS: + optimizer_root = results_root / optimizer + if not optimizer_root.is_dir(): + continue + for child in optimizer_root.iterdir(): + if not child.is_dir(): + continue + match = pattern.match(child.name) + if match and int(match.group(1)) not in expected: + unexpected.append(str(child)) + for child in results_root.iterdir(): + if not child.is_dir() or child.name in OPTIMIZERS: + continue + if any( + grandchild.is_dir() and pattern.match(grandchild.name) + for grandchild in child.iterdir() + ): + unexpected.append(str(child)) + return sorted(unexpected) + + +def _validate_step_span( + frame: pd.DataFrame, + *, + total_steps: int, + label: str, +) -> None: + _require_columns(frame, ("step",), label) + steps = pd.to_numeric(frame["step"], errors="coerce") + if steps.isna().any(): + raise CampaignValidationError(f"{label} contains nonnumeric steps") + integer_steps = steps.astype(int) + if not np.allclose(steps.to_numpy(dtype=float), integer_steps.to_numpy(dtype=float)): + raise CampaignValidationError(f"{label} contains noninteger steps") + if 0 not in set(integer_steps) or total_steps not in set(integer_steps): + raise CampaignValidationError( + f"{label} does not span step zero through {total_steps}" + ) + if int(integer_steps.max()) != int(total_steps): + raise CampaignValidationError( + f"{label} extends to {integer_steps.max()}, expected {total_steps}" + ) + + +def _validate_spectral_inventory( + layers: pd.DataFrame, + epoch_metrics: pd.DataFrame, + summary: pd.DataFrame, + *, + label: str, +) -> None: + _require_columns( + layers, + ( + "step", + "matrix_type", + "alpha_raw", + "alpha_clip_xmax", + "alpha", + "raw_alpha", + "alpha_delta", + "num_fingers", + "finger_policy", + "primary_alpha_variant", + "weightwatcher_analysis_calls", + ), + f"{label} spectral/layers.csv", + ) + alpha_values = layers[["alpha_raw", "alpha_clip_xmax"]].apply( + pd.to_numeric, + errors="coerce", + ) + if not np.isfinite(alpha_values.to_numpy(dtype=float)).all(): + raise CampaignValidationError( + f"{label} has non-finite raw or clip_xmax alpha values" + ) + aliases = layers[ + ["alpha", "raw_alpha", "alpha_delta", "num_fingers"] + ].apply(pd.to_numeric, errors="coerce") + if not np.isfinite(aliases.to_numpy(dtype=float)).all(): + raise CampaignValidationError(f"{label} has non-finite alpha alias values") + clipped = pd.to_numeric(layers["alpha_clip_xmax"], errors="coerce") + raw = pd.to_numeric(layers["alpha_raw"], errors="coerce") + if not np.allclose(aliases["alpha"], clipped, rtol=0.0, atol=0.0): + raise CampaignValidationError(f"{label} alpha alias differs from alpha_clip_xmax") + if not np.allclose(aliases["raw_alpha"], raw, rtol=0.0, atol=0.0): + raise CampaignValidationError(f"{label} raw_alpha alias differs from alpha_raw") + if not np.allclose( + aliases["alpha_delta"], raw - clipped, rtol=1e-12, atol=1e-12 + ) or (aliases["num_fingers"] < 0).any(): + raise CampaignValidationError( + f"{label} has an invalid alpha_delta or negative num_fingers" + ) + analysis_calls = pd.to_numeric( + layers["weightwatcher_analysis_calls"], errors="coerce" + ) + if analysis_calls.isna().any() or not analysis_calls.eq(1).all(): + raise CampaignValidationError( + f"{label} does not record exactly one WeightWatcher analysis call " + "per checkpoint row" + ) + if not layers["finger_policy"].astype(str).eq( + "fix_fingers=clip_xmax" + ).all(): + raise CampaignValidationError( + f"{label} does not consistently declare fix_fingers=clip_xmax" + ) + if not layers["primary_alpha_variant"].astype(str).eq( + "clip_xmax" + ).all(): + raise CampaignValidationError( + f"{label} does not consistently declare clip_xmax as primary alpha" + ) + + expected_matrices = set(MATRIX_TYPES) + epoch_steps = set(pd.to_numeric(epoch_metrics["step"], errors="coerce").astype(int)) + layer_steps = set(pd.to_numeric(layers["step"], errors="coerce").astype(int)) + summary_steps = set(pd.to_numeric(summary["step"], errors="coerce").astype(int)) + if layer_steps != epoch_steps or summary_steps != epoch_steps: + raise CampaignValidationError( + f"{label} spectral steps do not exactly match permanent checkpoint steps" + ) + if summary["step"].duplicated().any(): + raise CampaignValidationError( + f"{label} spectral/summary.csv contains duplicate steps" + ) + _require_columns(summary, ("n_matrices",), f"{label} spectral/summary.csv") + matrix_counts = pd.to_numeric(summary["n_matrices"], errors="coerce") + if matrix_counts.isna().any() or not matrix_counts.eq(len(MATRIX_TYPES)).all(): + raise CampaignValidationError( + f"{label} spectral summaries do not report exactly six matrices" + ) + duplicated = layers.duplicated(["step", "matrix_type"], keep=False) + if duplicated.any(): + rows = layers.loc[duplicated, ["step", "matrix_type"]].to_dict("records") + raise CampaignValidationError( + f"{label} contains duplicate step/matrix spectral rows: {rows[:12]}" + ) + for step, group in layers.groupby("step", sort=True): + observed = set(group["matrix_type"].astype(str)) + if observed != expected_matrices or len(group) != len(MATRIX_TYPES): + raise CampaignValidationError( + f"{label} step={int(step)} has matrices={sorted(observed)}, " + f"expected={sorted(expected_matrices)}" + ) + + +def _normalize_test_results( + payload: Mapping[str, Any], + *, + optimizer: str, + seed: int, + source_path: Path, +) -> pd.DataFrame: + policy = str(payload.get("policy", "")).lower() + if "held out" not in policy or "validation" not in policy or "never" not in policy: + raise CampaignValidationError( + f"{source_path} does not declare the held-out post-training test policy" + ) + rows: list[dict[str, Any]] = [] + aliases = { + "final": ("final",), + "validation_selected": ("validation_selected", "best", "selected"), + } + for canonical, candidates in aliases.items(): + key = next((name for name in candidates if name in payload), None) + if key is None or not isinstance(payload[key], Mapping): + raise CampaignValidationError( + f"{source_path} lacks required test result {canonical!r}" + ) + values = payload[key] + row = { + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "seed": int(seed), + "source_file": str(source_path), + "checkpoint": canonical, + } + for output, source_names in { + "step": ("step",), + "test_loss": ("loss", "test_loss"), + "test_perplexity": ("perplexity", "test_perplexity"), + "test_bits_per_token": ( + "bits_per_token", + "test_bits_per_token", + ), + "test_accuracy": ("accuracy", "test_accuracy"), + "test_top5_accuracy": ( + "top5_accuracy", + "test_top5_accuracy", + ), + "test_bleu": ("bleu", "test_bleu"), + "test_continuation_token_accuracy": ( + "continuation_token_accuracy", + "test_continuation_token_accuracy", + ), + "test_continuation_exact_match": ( + "continuation_exact_match", + "test_continuation_exact_match", + ), + }.items(): + source = next((name for name in source_names if name in values), None) + if source is None: + raise CampaignValidationError( + f"{source_path} {canonical} lacks {output}" + ) + row[output] = values[source] + rows.append(row) + result = _numeric(pd.DataFrame(rows), ("step", *TEST_RESULT_METRICS)) + values = result[list(TEST_RESULT_METRICS)] + if not np.isfinite(values.to_numpy(dtype=float)).all(): + raise CampaignValidationError(f"{source_path} has non-finite test metrics") + for _, row in result.iterrows(): + label = str(row["checkpoint"]) + loss = float(row["test_loss"]) + perplexity = float(row["test_perplexity"]) + bits = float(row["test_bits_per_token"]) + bounded = ( + float(row["test_accuracy"]), + float(row["test_top5_accuracy"]), + float(row["test_continuation_token_accuracy"]), + float(row["test_continuation_exact_match"]), + ) + if ( + loss < 0.0 + or bits < 0.0 + or perplexity <= 0.0 + or not math.isclose(math.log(perplexity), loss, rel_tol=1e-10, abs_tol=1e-10) + or not math.isclose(bits * math.log(2.0), loss, rel_tol=1e-10, abs_tol=1e-10) + or any(not 0.0 <= value <= 1.0 for value in bounded) + or float(row["test_top5_accuracy"]) < float(row["test_accuracy"]) + or float(row["test_continuation_exact_match"]) + > float(row["test_continuation_token_accuracy"]) + 1e-12 + or not 0.0 <= float(row["test_bleu"]) <= 100.0 + ): + raise CampaignValidationError( + f"{source_path} {label} has inconsistent or out-of-range test metrics" + ) + return result + + +def _campaign_invariants(manifest: Mapping[str, Any]) -> dict[str, Any]: + runtime = manifest.get("runtime_environment", {}) + runtime_mapping = runtime if isinstance(runtime, Mapping) else {} + source = manifest.get( + "source_repository", + manifest.get("source", {}), + ) + source_mapping = source if isinstance(source, Mapping) else {} + return { + "protocol": manifest.get("protocol"), + "config_sha256": manifest.get("config_sha256"), + "model": manifest.get("model"), + "training": manifest.get("training"), + "evaluation": manifest.get("evaluation"), + "weightwatcher": manifest.get("weightwatcher"), + "data_metadata": manifest.get("data_metadata"), + "tokens_per_step": manifest.get("tokens_per_step"), + "max_steps": manifest.get("max_steps"), + "package_versions": manifest.get("package_versions"), + "runtime_environment": _runtime_block_identity(runtime_mapping), + "accelerator": runtime_mapping.get("accelerator"), + "torch_version": manifest.get("torch_version", runtime_mapping.get("torch_version")), + "git_available": source_mapping.get("available"), + "git_commit": manifest.get( + "git_commit", + manifest.get("source_commit", source_mapping.get("commit")), + ), + "git_dirty": manifest.get( + "git_dirty", + source_mapping.get("dirty"), + ), + } + + +def _runtime_block_identity(runtime: Mapping[str, Any]) -> dict[str, Any]: + identity = dict(runtime) + if str(identity.get("hardware_block_id_source", "")) == "user": + for field in ( + "python_executable", + "processor", + "cuda_device_uuid", + "cuda_device_count", + "xla_process_index", + ): + identity.pop(field, None) + return identity + + +def _validate_matched_campaign( + manifests: Sequence[Mapping[str, Any]], + *, + allow_mixed_runtime: bool, +) -> None: + if not manifests: + raise CampaignValidationError("no complete run manifests were loaded") + frozen = _frozen_campaign_config() + current_commit = _repository_head() + reference = _campaign_invariants(manifests[0]) + scientific_keys = ( + "protocol", + "config_sha256", + "model", + "training", + "evaluation", + "weightwatcher", + "data_metadata", + "tokens_per_step", + "max_steps", + "package_versions", + ) + for index, manifest in enumerate(manifests): + observed = _campaign_invariants(manifest) + if index: + for key in scientific_keys: + if _canonical_json(observed[key]) != _canonical_json(reference[key]): + raise CampaignValidationError( + f"campaign invariant {key!r} differs in manifest index {index}" + ) + for key in ("git_available", "git_commit", "git_dirty"): + if _canonical_json(observed[key]) != _canonical_json(reference[key]): + raise CampaignValidationError( + f"source invariant {key!r} differs across runs" + ) + if not allow_mixed_runtime and _canonical_json( + observed["runtime_environment"] + ) != _canonical_json(reference["runtime_environment"]): + raise CampaignValidationError( + "runtime/hardware-block identity differs across runs; " + "create separate reports for separate hardware blocks" + ) + if observed.get("git_available") is not True: + raise CampaignValidationError( + f"manifest index {index} has no readable Git source identity" + ) + if observed.get("git_dirty") is not False: + raise CampaignValidationError( + "run manifests must identify a clean source tree; " + f"manifest index {index} records dirty={observed.get('git_dirty')!r}" + ) + commit = str(observed.get("git_commit", "")) + if not commit or commit == "unknown": + raise CampaignValidationError( + f"manifest index {index} has no exact Git commit" + ) + if commit != current_commit: + raise CampaignValidationError( + f"manifest index {index} was produced by {commit}, but the " + f"checked-out report source is {current_commit}" + ) + + if str(manifest.get("config_sha256", "")) != FROZEN_CONFIG_SHA256: + raise CampaignValidationError( + f"manifest index {index} does not use the frozen campaign config" + ) + initial_model_hash = str(manifest.get("initial_model_sha256", "")) + if ( + len(initial_model_hash) != 64 + or any( + character not in "0123456789abcdef" + for character in initial_model_hash.lower() + ) + ): + raise CampaignValidationError( + f"manifest index {index} has no initial-model tensor hash" + ) + for key in ("protocol", "model", "training", "evaluation", "weightwatcher"): + if _canonical_json(manifest.get(key)) != _canonical_json(frozen[key]): + raise CampaignValidationError( + f"manifest index {index} {key} differs from the frozen config" + ) + optimizer = str(manifest.get("optimizer", "")) + expected_profile = dict(frozen["optimizer_profiles"].get(optimizer, {})) + expected_profile["name"] = optimizer + if _canonical_json(manifest.get("optimizer_profile")) != _canonical_json( + expected_profile + ): + raise CampaignValidationError( + f"manifest index {index} optimizer profile differs from the frozen config" + ) + if int(manifest.get("max_steps", -1)) != EXPECTED_TOTAL_STEPS: + raise CampaignValidationError( + f"manifest index {index} has the wrong optimizer-step horizon" + ) + expected_tokens_per_step = ( + int(frozen["training"]["batch_size"]) + * int(frozen["training"]["grad_accum_steps"]) + * int(frozen["model"]["block_size"]) + ) + if int(manifest.get("tokens_per_step", -1)) != expected_tokens_per_step: + raise CampaignValidationError( + f"manifest index {index} has the wrong tokens_per_step" + ) + packages = manifest.get("package_versions") + if not isinstance(packages, Mapping): + raise CampaignValidationError( + f"manifest index {index} has no package_versions mapping" + ) + missing_packages = [name for name in MANIFEST_PACKAGES if not packages.get(name)] + if missing_packages: + raise CampaignValidationError( + f"manifest index {index} lacks dependency versions: {missing_packages}" + ) + if str(packages.get("weightwatcher")) != PINNED_WEIGHTWATCHER: + raise CampaignValidationError( + f"manifest index {index} did not use WeightWatcher {PINNED_WEIGHTWATCHER}" + ) + if str(packages.get("rg-nanogpt-one-head")) != PINNED_PACKAGE_VERSION: + raise CampaignValidationError( + f"manifest index {index} did not use campaign package " + f"{PINNED_PACKAGE_VERSION}" + ) + data = manifest.get("data_metadata") + expected_dataset_fields = { + "schema_version": 2, + "dataset_name": frozen["dataset"]["name"], + "dataset_config": frozen["dataset"]["config"], + "dataset_split": frozen["dataset"]["split"], + "dataset_revision": frozen["dataset"]["revision"], + "tokenizer": frozen["dataset"]["tokenizer"], + "vocab_size": frozen["model"]["vocab_size"], + "eot_token": 50_256, + "dtype": "uint16", + "document_disjoint_splits": True, + "splits": { + "train": frozen["dataset"]["train_tokens"], + "val": frozen["dataset"]["val_tokens"], + "test": frozen["dataset"]["test_tokens"], + }, + } + if not isinstance(data, Mapping) or any( + data.get(key) != value for key, value in expected_dataset_fields.items() + ): + raise CampaignValidationError( + f"manifest index {index} data metadata differs from the frozen corpus" + ) + files = data.get("files") + if not isinstance(files, Mapping) or any( + not isinstance(files.get(split), Mapping) + or files[split].get("path") != f"{split}.bin" + or len(str(files[split].get("sha256", ""))) != 64 + or any( + character not in "0123456789abcdef" + for character in str(files[split].get("sha256", "")).lower() + ) + or int(files[split].get("bytes", -1)) != int(tokens) * 2 + for split, tokens in expected_dataset_fields["splits"].items() + ): + raise CampaignValidationError( + f"manifest index {index} has no complete corpus hash inventory" + ) + + runtime = manifest.get("runtime_environment", {}) + if not isinstance(runtime, Mapping): + raise CampaignValidationError( + f"manifest index {index} has no runtime_environment mapping" + ) + if runtime.get("float32_matmul_precision") != "highest": + raise CampaignValidationError( + f"manifest index {index} did not use matmul_precision=highest" + ) + if runtime.get("deterministic_algorithms") is not True: + raise CampaignValidationError( + f"manifest index {index} did not enable deterministic algorithms" + ) + if runtime.get("deterministic_warn_only") is not False: + raise CampaignValidationError( + f"manifest index {index} used deterministic warn-only mode" + ) + if runtime.get("accelerator") not in {"cpu", "mps", "cuda", "tpu"}: + raise CampaignValidationError( + f"manifest index {index} records an unsupported accelerator: " + f"{runtime.get('accelerator')!r}" + ) + if not str(runtime.get("hardware_block_id", "")).strip() or not str( + runtime.get("hardware_block_id_source", "") + ).strip(): + raise CampaignValidationError( + f"manifest index {index} has no complete hardware-block identity" + ) + if runtime.get("accelerator") == "cuda" and ( + runtime.get("cuda_matmul_allow_tf32") is not False + or runtime.get("cudnn_allow_tf32") is not False + ): + raise CampaignValidationError( + f"manifest index {index} enabled CUDA TF32" + ) + + if len(manifests) == len(OPTIMIZERS) * len(SEEDS): + initial_hashes_by_seed: dict[int, dict[str, str]] = defaultdict(dict) + for manifest in manifests: + initial_hashes_by_seed[int(manifest["seed"])][ + str(manifest["optimizer"]) + ] = str(manifest["initial_model_sha256"]) + for seed in SEEDS: + hashes = initial_hashes_by_seed.get(seed, {}) + if set(hashes) != set(OPTIMIZERS) or len(set(hashes.values())) != 1: + raise CampaignValidationError( + "optimizer arms do not share identical step-zero tensors for " + f"seed {seed}: {hashes}" + ) + + model = reference.get("model") + training = reference.get("training") + weightwatcher = reference.get("weightwatcher") + protocol = reference.get("protocol") + if not isinstance(model, Mapping) or ( + int(model.get("n_layer", -1)) != 1 + or int(model.get("n_head", -1)) != 1 + ): + raise CampaignValidationError( + "the dated baseline requires a one-layer, one-head model" + ) + if not isinstance(training, Mapping) or ( + tuple(int(value) for value in training.get("seeds", ())) != SEEDS + or not math.isclose(float(training.get("target_epochs", -1.0)), 4.0) + or not math.isclose(float(training.get("epoch_interval", -1.0)), 0.25) + ): + raise CampaignValidationError( + "the dated baseline requires seeds " + "1337/2027/4099/31415/271828, four epochs, " + "and quarter-epoch permanent states" + ) + if not isinstance(weightwatcher, Mapping) or ( + weightwatcher.get("enabled") is not True + or weightwatcher.get("strict") is not True + or weightwatcher.get("fix_fingers") != "clip_xmax" + or int(weightwatcher.get("max_fingers", -1)) != 10 + or weightwatcher.get("require_raw_alpha") is not True + ): + raise CampaignValidationError( + "the dated baseline requires strict one-call WeightWatcher with " + "fix_fingers=clip_xmax, max_fingers=10, and raw alpha" + ) + if not isinstance(protocol, Mapping) or protocol.get("name") != ( + "nanogpt_one_head_2026_08_21_ww_baseline" + ): + raise CampaignValidationError( + "run manifests do not identify the dated baseline protocol" + ) + + +def _load_campaign( + results_root: Path, + *, + require_complete: bool, + allow_extra_runs: bool, + allow_mixed_runtime: bool, +) -> dict[str, Any]: + unexpected = _unexpected_seed_directories(results_root) + if unexpected and not allow_extra_runs: + raise CampaignValidationError( + "unexpected seed directories make this more than the exact 2 x 5 campaign: " + + ", ".join(unexpected) + ) + + run_rows: list[dict[str, Any]] = [] + metrics_frames: list[pd.DataFrame] = [] + epoch_frames: list[pd.DataFrame] = [] + layer_frames: list[pd.DataFrame] = [] + spectral_summary_frames: list[pd.DataFrame] = [] + test_frames: list[pd.DataFrame] = [] + qk_frames: list[pd.DataFrame] = [] + manifests: list[dict[str, Any]] = [] + completions: dict[tuple[str, int], dict[str, Any]] = {} + run_dirs: dict[tuple[str, int], Path] = {} + errors: list[str] = [] + + for optimizer in OPTIMIZERS: + for seed in SEEDS: + run_dir = _expected_run_directory(results_root, optimizer, seed) + run_dirs[(optimizer, seed)] = run_dir + missing = [ + relative + for relative in REQUIRED_RUN_FILES + if not (run_dir / relative).is_file() + or (run_dir / relative).stat().st_size == 0 + ] + if missing: + message = ( + f"optimizer={optimizer} seed={seed} missing required artifacts: {missing}" + ) + errors.append(message) + run_rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "seed": seed, + "run_dir": str(run_dir), + "complete": False, + "validation_error": message, + }) + continue + + try: + manifest = _read_json(run_dir / "manifest.json") + completion = _read_json(run_dir / "run_complete.json") + manifest_optimizer = str(manifest.get("optimizer", "")) + completion_optimizer = str(completion.get("optimizer", "")) + manifest_seed = int(manifest.get("seed")) + completion_seed = int(completion.get("seed")) + if manifest_optimizer != optimizer or completion_optimizer != optimizer: + raise CampaignValidationError( + f"optimizer identity mismatch in {run_dir}" + ) + if manifest_seed != seed or completion_seed != seed: + raise CampaignValidationError(f"seed identity mismatch in {run_dir}") + if completion.get("completed") is not True: + raise CampaignValidationError( + f"run_complete.json does not declare completed=true: {run_dir}" + ) + if str(completion.get("fingerprint", "")) != str( + manifest.get("protocol_fingerprint", "") + ): + raise CampaignValidationError( + f"manifest/completion protocol fingerprint mismatch in {run_dir}" + ) + optimizer_profile = manifest.get("optimizer_profile", {}) + if not isinstance(optimizer_profile, Mapping) or str( + optimizer_profile.get("family", "") + ) != optimizer: + raise CampaignValidationError( + f"optimizer profile family mismatch in {run_dir}" + ) + + total_steps = int(completion["optimizer_steps"]) + if int(manifest.get("max_steps", total_steps)) != total_steps: + raise CampaignValidationError( + f"manifest/completion total-step mismatch in {run_dir}" + ) + if total_steps != EXPECTED_TOTAL_STEPS: + raise CampaignValidationError( + f"optimizer-step horizon is {total_steps}, expected " + f"{EXPECTED_TOTAL_STEPS}: {run_dir}" + ) + + metrics_path = run_dir / "metrics.csv" + metrics = _numeric( + _read_csv(metrics_path), + ( + "step", + "tokens_seen", + "epoch", + "train_loss", + "val_loss", + "test_loss", + "train_perplexity", + "val_perplexity", + "test_perplexity", + "train_bits_per_token", + "val_bits_per_token", + "test_bits_per_token", + "train_accuracy", + "val_accuracy", + "test_accuracy", + "train_top5_accuracy", + "val_top5_accuracy", + "test_top5_accuracy", + "test_bleu", + "test_continuation_token_accuracy", + "test_continuation_exact_match", + "val_generalization_gap", + "test_generalization_gap", + "tokens_per_sec", + ), + ) + _require_columns( + metrics, + ("step", "epoch", *FINAL_EPOCH_METRICS), + str(metrics_path), + ) + _validate_step_span(metrics, total_steps=total_steps, label=str(metrics_path)) + always_finite = metrics[[ + "train_loss", + "val_loss", + "train_perplexity", + "val_perplexity", + "train_bits_per_token", + "val_bits_per_token", + "train_accuracy", + "val_accuracy", + "train_top5_accuracy", + "val_top5_accuracy", + ]].apply(pd.to_numeric, errors="coerce") + if not np.isfinite(always_finite.to_numpy(dtype=float)).all(): + raise CampaignValidationError( + f"{metrics_path} contains non-finite train/validation metrics" + ) + if metrics[list(HELD_OUT_CURVE_COLUMNS)].notna().any().any(): + raise CampaignValidationError( + f"{metrics_path} leaks held-out test outcomes into training curves" + ) + validation_minimum_index = metrics["val_loss"].idxmin() + validation_minimum = metrics.loc[validation_minimum_index] + if int(validation_minimum["step"]) != int( + completion["best_validation_step"] + ) or not math.isclose( + float(validation_minimum["val_loss"]), + float(completion["best_validation_loss"]), + rel_tol=1e-10, + abs_tol=1e-12, + ): + raise CampaignValidationError( + f"validation-selected checkpoint metadata disagrees with {metrics_path}" + ) + + epoch_path = run_dir / "epoch_metrics.csv" + epoch_metrics = _numeric( + _read_csv(epoch_path), + ( + "step", + "tokens_seen", + "epoch", + "nominal_epoch", + "tokens_per_sec", + "val_generalization_gap", + "test_generalization_gap", + *FINAL_EPOCH_METRICS, + ), + ) + _require_columns( + epoch_metrics, + ( + "step", + "epoch", + "nominal_epoch", + "checkpoint_path", + "val_loss", + "test_monitoring_only", + "test_held_out", + ), + str(epoch_path), + ) + _validate_step_span( + epoch_metrics, total_steps=total_steps, label=str(epoch_path) + ) + if epoch_metrics["step"].duplicated().any(): + raise CampaignValidationError( + f"duplicate permanent-checkpoint steps in {epoch_path}" + ) + observed_permanent_steps = tuple( + sorted(pd.to_numeric(epoch_metrics["step"], errors="raise").astype(int)) + ) + if observed_permanent_steps != EXPECTED_PERMANENT_STEPS: + raise CampaignValidationError( + f"{epoch_path} does not contain the exact frozen 17-state grid" + ) + monitoring_policy = pd.to_numeric( + epoch_metrics["test_monitoring_only"], errors="coerce" + ) + if monitoring_policy.isna().any() or not monitoring_policy.eq(1).all(): + raise CampaignValidationError( + f"{epoch_path} violates the held-out test policy" + ) + held_out_policy = pd.to_numeric( + epoch_metrics["test_held_out"], errors="coerce" + ) + if held_out_policy.isna().any() or not held_out_policy.eq(1).all(): + raise CampaignValidationError( + f"{epoch_path} does not mark every test curve as held out" + ) + epoch_outcomes = epoch_metrics[list(FINITE_EPOCH_METRICS)].apply( + pd.to_numeric, + errors="coerce", + ) + if not np.isfinite(epoch_outcomes.to_numpy(dtype=float)).all(): + raise CampaignValidationError( + f"{epoch_path} contains non-finite permanent-state outcomes" + ) + if epoch_metrics[list(HELD_OUT_CURVE_COLUMNS)].notna().any().any(): + raise CampaignValidationError( + f"{epoch_path} leaks held-out test outcomes into permanent-state curves" + ) + if len(epoch_metrics) < MIN_PERMANENT_CHECKPOINTS: + raise CampaignValidationError( + f"{epoch_path} contains {len(epoch_metrics)} permanent checkpoints; " + f"at least {MIN_PERMANENT_CHECKPOINTS} are required" + ) + + layers_path = run_dir / "spectral" / "layers.csv" + layers = _normalize_spectral_layers( + _read_csv(layers_path), str(layers_path) + ) + spectral_summary_path = run_dir / "spectral" / "summary.csv" + spectral_summary = _numeric( + _read_csv(spectral_summary_path), + ("step", "tokens_seen", "epoch"), + ) + _require_columns(spectral_summary, ("step", "epoch"), str(spectral_summary_path)) + _validate_spectral_inventory( + layers, + epoch_metrics, + spectral_summary, + label=f"optimizer={optimizer} seed={seed}", + ) + for step_value in sorted(epoch_metrics["step"].astype(int)): + status_path = ( + run_dir + / "spectral" + / f"status_step_{step_value:07d}.json" + ) + if not status_path.is_file() or status_path.stat().st_size == 0: + raise CampaignValidationError( + f"missing WeightWatcher completion record: {status_path}" + ) + status = _read_json(status_path) + if status.get("completed") is not True: + raise CampaignValidationError( + f"WeightWatcher completion record is incomplete: {status_path}" + ) + if int(status.get("weightwatcher_analysis_calls", -1)) != 1: + raise CampaignValidationError( + f"WeightWatcher was not called exactly once: {status_path}" + ) + if status.get("finger_policy") != "fix_fingers=clip_xmax": + raise CampaignValidationError( + f"WeightWatcher finger policy mismatch: {status_path}" + ) + for count_key in ( + "alpha_raw_valid_matrices", + "alpha_clip_xmax_valid_matrices", + ): + if int(status.get(count_key, -1)) != len(MATRIX_TYPES): + raise CampaignValidationError( + f"{status_path} does not record six valid " + f"matrices for {count_key}" + ) + + test_path = run_dir / "test_results.json" + test_results = _normalize_test_results( + _read_json(test_path), + optimizer=optimizer, + seed=seed, + source_path=test_path, + ) + final_step = int( + test_results.loc[ + test_results["checkpoint"].eq("final"), "step" + ].iloc[0] + ) + selected_step = int( + test_results.loc[ + test_results["checkpoint"].eq("validation_selected"), "step" + ].iloc[0] + ) + if final_step != total_steps: + raise CampaignValidationError( + f"final test step={final_step}, expected {total_steps}: {test_path}" + ) + if selected_step != int(completion["best_validation_step"]): + raise CampaignValidationError( + f"validation-selected test step mismatch: {test_path}" + ) + final_test_row = test_results[ + test_results["checkpoint"].eq("final") + ].iloc[0] + for metric, completion_key in { + "test_loss": "final_test_loss", + "test_perplexity": "final_test_perplexity", + "test_bits_per_token": "final_test_bits_per_token", + "test_accuracy": "final_test_accuracy", + "test_top5_accuracy": "final_test_top5_accuracy", + "test_bleu": "final_test_bleu", + "test_continuation_token_accuracy": ( + "final_test_continuation_token_accuracy" + ), + "test_continuation_exact_match": ( + "final_test_continuation_exact_match" + ), + }.items(): + observed = float(final_test_row[metric]) + recorded = float(completion[completion_key]) + if not math.isfinite(observed) or not math.isclose( + observed, + recorded, + rel_tol=1e-10, + abs_tol=1e-12, + ): + raise CampaignValidationError( + f"completion/test result mismatch for {metric}: {test_path}" + ) + + qk_path = run_dir / "muonclip_qk.csv" + if optimizer == "muon_clip": + if not qk_path.is_file() or qk_path.stat().st_size == 0: + raise CampaignValidationError( + f"MuonClip run lacks required QK diagnostics: {qk_path}" + ) + qk = _numeric( + _read_csv(qk_path), + ( + "step", + "threshold", + "steps_in_interval", + "head_observations", + "active_heads", + "active_fraction", + "mean_max_logit", + "max_logit", + "mean_gamma", + "min_gamma", + ), + ) + _require_columns( + qk, + ( + "step", + "threshold", + "head_observations", + "active_fraction", + "mean_max_logit", + "max_logit", + "mean_gamma", + "min_gamma", + ), + str(qk_path), + ) + qk_required = qk[[ + "step", + "threshold", + "steps_in_interval", + "head_observations", + "active_heads", + "active_fraction", + "mean_max_logit", + "max_logit", + "mean_gamma", + "min_gamma", + ]].apply(pd.to_numeric, errors="coerce") + if not np.isfinite(qk_required.to_numpy(dtype=float)).all(): + raise CampaignValidationError( + f"{qk_path} contains non-finite QK diagnostics" + ) + observed_steps = qk_required["step"].to_numpy(dtype=float) + if ( + not np.allclose(observed_steps, np.rint(observed_steps)) + or tuple(int(value) for value in observed_steps) + != EXPECTED_QK_STEPS + ): + raise CampaignValidationError( + f"{qk_path} does not cover the exact 500-step QK grid" + ) + expected_intervals = np.diff( + np.asarray([0, *EXPECTED_QK_STEPS], dtype=int) + ) + observed_intervals = qk_required[ + "steps_in_interval" + ].to_numpy(dtype=float) + if ( + not np.array_equal( + observed_intervals, + expected_intervals.astype(float), + ) + or int(observed_intervals.sum()) != EXPECTED_TOTAL_STEPS + ): + raise CampaignValidationError( + f"{qk_path} QK intervals do not cover the training horizon" + ) + observations = qk_required[ + "head_observations" + ].to_numpy(dtype=float) + active_heads = qk_required["active_heads"].to_numpy( + dtype=float + ) + active_fraction = qk_required[ + "active_fraction" + ].to_numpy(dtype=float) + if ( + not np.array_equal(observations, observed_intervals) + or (active_heads < 0.0).any() + or (active_heads > observations).any() + or not np.allclose( + active_fraction, + active_heads / observations, + rtol=1e-12, + atol=1e-12, + ) + ): + raise CampaignValidationError( + f"{qk_path} QK head counts/fractions are inconsistent" + ) + if ( + not qk_required["threshold"].eq(100.0).all() + or not qk_required["active_fraction"].between(0.0, 1.0).all() + or not qk_required["mean_gamma"].between(0.0, 1.0).all() + or not qk_required["min_gamma"].between(0.0, 1.0).all() + or (qk_required["min_gamma"] > qk_required["mean_gamma"]).any() + or ( + qk_required["mean_max_logit"] + > qk_required["max_logit"] + ).any() + ): + raise CampaignValidationError( + f"{qk_path} violates the registered QK diagnostic bounds" + ) + qk_frames.append( + _with_identity( + qk, + optimizer=optimizer, + seed=seed, + source_path=qk_path, + ) + ) + + metrics_frames.append( + _with_identity( + metrics, + optimizer=optimizer, + seed=seed, + source_path=metrics_path, + ) + ) + epoch_frames.append( + _with_identity( + epoch_metrics, + optimizer=optimizer, + seed=seed, + source_path=epoch_path, + ) + ) + layer_frames.append( + _with_identity( + layers, + optimizer=optimizer, + seed=seed, + source_path=layers_path, + ) + ) + spectral_summary_frames.append( + _with_identity( + spectral_summary, + optimizer=optimizer, + seed=seed, + source_path=spectral_summary_path, + ) + ) + test_frames.append(test_results) + manifests.append(manifest) + completions[(optimizer, seed)] = completion + + runtime = manifest.get("runtime_environment", {}) + runtime_mapping = runtime if isinstance(runtime, Mapping) else {} + run_rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "seed": seed, + "run_dir": str(run_dir), + "complete": True, + "validation_error": "", + "optimizer_steps": total_steps, + "train_epochs": completion.get("train_epochs"), + "elapsed_seconds": completion.get("elapsed_seconds"), + "best_validation_step": completion.get("best_validation_step"), + "best_validation_loss": completion.get("best_validation_loss"), + "final_test_loss": completion.get("final_test_loss"), + "final_test_perplexity": completion.get("final_test_perplexity"), + "final_test_accuracy": completion.get("final_test_accuracy"), + "final_test_bleu": completion.get("final_test_bleu"), + "protocol_fingerprint": completion.get( + "fingerprint", manifest.get("protocol_fingerprint") + ), + "accelerator": runtime_mapping.get("accelerator"), + "device": manifest.get("device", runtime_mapping.get("device")), + "torch_version": manifest.get( + "torch_version", runtime_mapping.get("torch_version") + ), + "git_commit": _campaign_invariants(manifest).get("git_commit"), + "git_dirty": _campaign_invariants(manifest).get("git_dirty"), + }) + except (CampaignValidationError, KeyError, TypeError, ValueError) as exc: + message = f"optimizer={optimizer} seed={seed}: {exc}" + errors.append(message) + run_rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "seed": seed, + "run_dir": str(run_dir), + "complete": False, + "validation_error": message, + }) + + if errors and require_complete: + raise CampaignValidationError( + "the exact 2 x 5 campaign is incomplete or invalid:\n- " + + "\n- ".join(errors) + ) + if require_complete and len(manifests) != len(OPTIMIZERS) * len(SEEDS): + raise CampaignValidationError( + f"loaded {len(manifests)} valid runs, expected exactly " + f"{len(OPTIMIZERS) * len(SEEDS)}" + ) + _validate_matched_campaign(manifests, allow_mixed_runtime=allow_mixed_runtime) + + def combine(frames: Sequence[pd.DataFrame]) -> pd.DataFrame: + return ( + pd.concat(frames, ignore_index=True, sort=False) + if frames + else pd.DataFrame() + ) + + return { + "campaign_runs": pd.DataFrame(run_rows), + "metrics": combine(metrics_frames), + "epoch_metrics": combine(epoch_frames), + "spectral_layers": combine(layer_frames), + "spectral_summary": combine(spectral_summary_frames), + "test_results": combine(test_frames), + "qk": combine(qk_frames), + "manifests": manifests, + "completions": completions, + "run_dirs": run_dirs, + "validation_errors": errors, + "unexpected_runs": unexpected, + } + + +def _stats(values: Iterable[float]) -> dict[str, float | int]: + array = np.asarray(list(values), dtype=float) + array = array[np.isfinite(array)] + n = int(array.size) + if n == 0: + return { + "n": 0, + "mean": np.nan, + "sd": np.nan, + "sem": np.nan, + "ci95_half_width": np.nan, + "ci95_low": np.nan, + "ci95_high": np.nan, + } + mean = float(array.mean()) + if n == 1: + return { + "n": 1, + "mean": mean, + "sd": np.nan, + "sem": np.nan, + "ci95_half_width": np.nan, + "ci95_low": np.nan, + "ci95_high": np.nan, + } + sd = float(array.std(ddof=1)) + sem = sd / math.sqrt(n) + critical_values = { + 2: T_975_DF1, + 3: T_975_DF2, + 4: T_975_DF3, + 5: T_975_DF4, + } + if n not in critical_values: + raise CampaignValidationError( + f"Student-t summary expected at most five seeds, observed n={n}" + ) + critical = critical_values[n] + half = critical * sem + return { + "n": n, + "mean": mean, + "sd": sd, + "sem": sem, + "ci95_half_width": half, + "ci95_low": mean - half, + "ci95_high": mean + half, + } + + +def _summarize_performance( + test_results: pd.DataFrame, + epoch_metrics: pd.DataFrame, +) -> pd.DataFrame: + rows: list[dict[str, Any]] = [] + if not test_results.empty: + for (optimizer, checkpoint), group in test_results.groupby( + ["optimizer", "checkpoint"], sort=True + ): + for metric in TEST_RESULT_METRICS: + statistics = _stats(group[metric]) + rows.append({ + "source": "test_results", + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "checkpoint": checkpoint, + "metric": metric, + "valid_exact_n5": bool(statistics["n"] == 5), + **statistics, + }) + + if not epoch_metrics.empty: + final_rows = ( + epoch_metrics.sort_values(["optimizer", "seed", "step"]) + .groupby(["optimizer", "seed"], as_index=False, sort=True) + .tail(1) + ) + for optimizer, group in final_rows.groupby("optimizer", sort=True): + for metric in FINITE_EPOCH_METRICS: + if metric in group.columns: + statistics = _stats(group[metric]) + rows.append({ + "source": "final_permanent_checkpoint", + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "checkpoint": "final", + "metric": metric, + "valid_exact_n5": bool(statistics["n"] == 5), + **statistics, + }) + return pd.DataFrame(rows) + + +def _paired_seed_differences( + test_results: pd.DataFrame, + epoch_metrics: pd.DataFrame, + *, + require_complete: bool, +) -> pd.DataFrame: + sources: list[pd.DataFrame] = [] + if not test_results.empty: + melted = test_results.melt( + id_vars=["optimizer", "seed", "checkpoint"], + value_vars=list(TEST_RESULT_METRICS), + var_name="metric", + value_name="value", + ) + melted["source"] = "test_results" + sources.append(melted) + + if not epoch_metrics.empty: + final_rows = ( + epoch_metrics.sort_values(["optimizer", "seed", "step"]) + .groupby(["optimizer", "seed"], as_index=False, sort=True) + .tail(1) + ) + value_columns = [ + name for name in FINITE_EPOCH_METRICS if name in final_rows.columns + ] + melted = final_rows.melt( + id_vars=["optimizer", "seed"], + value_vars=value_columns, + var_name="metric", + value_name="value", + ) + melted["checkpoint"] = "final" + melted["source"] = "final_permanent_checkpoint" + sources.append(melted) + + if not sources: + return pd.DataFrame() + values = pd.concat(sources, ignore_index=True, sort=False) + contrasts = tuple(combinations(OPTIMIZERS, 2)) + rows: list[dict[str, Any]] = [] + for (source, checkpoint, metric), group in values.groupby( + ["source", "checkpoint", "metric"], sort=True + ): + for optimizer_a, optimizer_b in contrasts: + left = group[group["optimizer"].eq(optimizer_a)][["seed", "value"]] + right = group[group["optimizer"].eq(optimizer_b)][["seed", "value"]] + paired = left.merge( + right, + on="seed", + how="inner", + suffixes=("_a", "_b"), + ).sort_values("seed") + paired["difference_b_minus_a"] = paired["value_b"] - paired["value_a"] + finite = paired[np.isfinite(pd.to_numeric( + paired["difference_b_minus_a"], errors="coerce" + ))] + if require_complete and tuple(finite["seed"].astype(int)) != SEEDS: + raise CampaignValidationError( + f"paired contrast {optimizer_b}-{optimizer_a}, source={source}, " + f"checkpoint={checkpoint}, metric={metric} lacks exact seeds {SEEDS}" + ) + statistics = _stats(finite["difference_b_minus_a"]) + row: dict[str, Any] = { + "source": source, + "checkpoint": checkpoint, + "metric": metric, + "optimizer_a": optimizer_a, + "optimizer_b": optimizer_b, + "contrast": f"{optimizer_b} minus {optimizer_a}", + "difference_definition": "optimizer_b - optimizer_a", + "paired_seeds": ",".join(str(value) for value in finite["seed"]), + "valid_exact_n5": bool(statistics["n"] == 5), + **statistics, + } + for _, paired_row in finite.iterrows(): + seed = int(paired_row["seed"]) + row[f"seed_{seed}_a"] = float(paired_row["value_a"]) + row[f"seed_{seed}_b"] = float(paired_row["value_b"]) + row[f"seed_{seed}_difference"] = float( + paired_row["difference_b_minus_a"] + ) + rows.append(row) + return pd.DataFrame(rows) + + +def _alpha_run_medians(layers: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]: + rows: list[dict[str, Any]] = [] + keys = ["optimizer", "seed", "step", "epoch"] + for values, group in layers.groupby(keys, sort=True): + optimizer, seed, step, epoch = values + observed = set(group["matrix_type"].astype(str)) + if observed != set(MATRIX_TYPES) or len(group) != len(MATRIX_TYPES): + raise CampaignValidationError( + f"cannot form six-matrix run median for {values}: {sorted(observed)}" + ) + raw = pd.to_numeric(group["alpha_raw"], errors="coerce") + clipped = pd.to_numeric(group["alpha_clip_xmax"], errors="coerce") + rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "seed": int(seed), + "step": int(step), + "epoch": float(epoch), + "matrix_count": int(len(group)), + "alpha_raw_finite_count": int(np.isfinite(raw).sum()), + "alpha_clip_xmax_finite_count": int(np.isfinite(clipped).sum()), + "alpha_raw_six_matrix_median": float(raw.median(skipna=True)), + "alpha_clip_xmax_six_matrix_median": float(clipped.median(skipna=True)), + "alpha_clip_minus_raw_six_matrix_median": float( + clipped.median(skipna=True) - raw.median(skipna=True) + ), + }) + run_medians = pd.DataFrame(rows) + + summary_rows: list[dict[str, Any]] = [] + value_columns = ( + "alpha_raw_six_matrix_median", + "alpha_clip_xmax_six_matrix_median", + "alpha_clip_minus_raw_six_matrix_median", + ) + for (optimizer, step, epoch), group in run_medians.groupby( + ["optimizer", "step", "epoch"], sort=True + ): + for metric in value_columns: + statistics = _stats(group[metric]) + summary_rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "step": int(step), + "epoch": float(epoch), + "metric": metric, + "aggregation_order": ( + "median across six matrices within each seeded run, " + "then Student-t summary across seeds" + ), + "valid_exact_n5": bool(statistics["n"] == 5), + **statistics, + }) + return run_medians, pd.DataFrame(summary_rows) + + +def _validation_saturation( + epoch_metrics: pd.DataFrame, +) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: + """Diagnose plateaus from permanent-state validation loss only.""" + + run_rows: list[dict[str, Any]] = [] + integer_rows: list[dict[str, Any]] = [] + for (optimizer, seed), group in epoch_metrics.groupby( + ["optimizer", "seed"], sort=True + ): + observations = group[["step", "epoch", "val_loss"]].copy() + observations = _numeric(observations, ("step", "epoch", "val_loss")) + observations = observations.dropna(subset=["epoch", "val_loss"]) + observations = observations[np.isfinite(observations["val_loss"])] + observations = observations.sort_values(["epoch", "step"]) + if observations.empty: + raise CampaignValidationError( + f"no finite validation loss for optimizer={optimizer} seed={seed}" + ) + maximum_epoch = float(observations["epoch"].max()) + targets = range(0, int(math.floor(maximum_epoch + 1e-9)) + 1) + selected: list[dict[str, Any]] = [] + used_steps: set[int] = set() + for target in targets: + distance = (observations["epoch"] - float(target)).abs() + index = distance.idxmin() + row = observations.loc[index] + if float(distance.loc[index]) > 0.15: + continue + step = int(row["step"]) + if step in used_steps: + continue + used_steps.add(step) + selected.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "seed": int(seed), + "target_epoch": int(target), + "actual_epoch": float(row["epoch"]), + "step": step, + "val_loss": float(row["val_loss"]), + }) + selected_frame = pd.DataFrame(selected).sort_values("target_epoch") + if selected_frame.empty: + raise CampaignValidationError( + f"could not select integer-epoch validation points for {optimizer}/{seed}" + ) + selected_frame["validation_improvement_nats"] = ( + selected_frame["val_loss"].shift(1) - selected_frame["val_loss"] + ) + selected_frame["near_flat_interval"] = ( + selected_frame["validation_improvement_nats"].abs() + <= SATURATION_DELTA_NATS + ) + selected_frame["degrading_interval"] = ( + selected_frame["validation_improvement_nats"] + < -SATURATION_DELTA_NATS + ) + selected_frame["two_consecutive_near_flat"] = ( + selected_frame["near_flat_interval"] + & selected_frame["near_flat_interval"].shift(1, fill_value=False) + ) + selected_frame["two_consecutive_degrading"] = ( + selected_frame["degrading_interval"] + & selected_frame["degrading_interval"].shift(1, fill_value=False) + ) + integer_rows.extend(selected_frame.to_dict("records")) + + final_two_intervals_flat = bool( + len(selected_frame) >= 3 + and selected_frame["near_flat_interval"].iloc[-2:].all() + ) + degradation_rows = selected_frame[selected_frame["two_consecutive_degrading"]] + best_index = observations["val_loss"].idxmin() + best = observations.loc[best_index] + final = observations.iloc[-1] + run_rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "seed": int(seed), + "criterion": ( + "validation-only: abs(one-epoch NLL improvement) <= " + f"{SATURATION_DELTA_NATS:.3f} nat/token for each of the final " + "two complete intervals" + ), + "plateau_detected": final_two_intervals_flat, + "plateau_assessment_end_epoch": ( + float(selected_frame.iloc[-1]["target_epoch"]) + if final_two_intervals_flat + else np.nan + ), + "degradation_detected": bool(not degradation_rows.empty), + "first_degradation_end_epoch": ( + float(degradation_rows.iloc[0]["target_epoch"]) + if not degradation_rows.empty + else np.nan + ), + "best_validation_step": int(best["step"]), + "best_validation_epoch": float(best["epoch"]), + "best_validation_loss": float(best["val_loss"]), + "final_validation_step": int(final["step"]), + "final_validation_epoch": float(final["epoch"]), + "final_validation_loss": float(final["val_loss"]), + "validation_observations": int(len(observations)), + "integer_epoch_observations": int(len(selected_frame)), + "test_metrics_used": False, + }) + + diagnostics = pd.DataFrame(run_rows) + integer_diagnostics = pd.DataFrame(integer_rows) + summary_rows: list[dict[str, Any]] = [] + for optimizer, group in diagnostics.groupby("optimizer", sort=True): + for metric in ( + "plateau_assessment_end_epoch", + "first_degradation_end_epoch", + "best_validation_epoch", + "best_validation_loss", + "final_validation_loss", + ): + summary_rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "metric": metric, + **_stats(group[metric]), + }) + return diagnostics, integer_diagnostics, pd.DataFrame(summary_rows) + + +def _qk_summary(qk: pd.DataFrame) -> pd.DataFrame: + if qk.empty: + return pd.DataFrame(columns=( + "optimizer", + "seed", + "diagnostic_rows", + "last_step", + "active_fraction_weighted", + "max_logit_observed", + "min_gamma_observed", + )) + rows: list[dict[str, Any]] = [] + for (optimizer, seed), group in qk.groupby(["optimizer", "seed"], sort=True): + weights = ( + pd.to_numeric(group["head_observations"], errors="coerce") + if "head_observations" in group.columns + else pd.Series(1.0, index=group.index) + ).fillna(0.0).clip(lower=0.0) + active = pd.to_numeric(group["active_fraction"], errors="coerce") + finite = np.isfinite(active) & np.isfinite(weights) + weighted_active = ( + float(np.average(active[finite], weights=weights[finite])) + if finite.any() and float(weights[finite].sum()) > 0.0 + else float(active[finite].mean()) if finite.any() else np.nan + ) + rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "seed": int(seed), + "diagnostic_rows": int(len(group)), + "last_step": float(pd.to_numeric(group["step"], errors="coerce").max()), + "active_fraction_weighted": weighted_active, + "max_logit_observed": float( + pd.to_numeric(group["max_logit"], errors="coerce").max() + ), + "min_gamma_observed": float( + pd.to_numeric(group["min_gamma"], errors="coerce").min() + ), + }) + return pd.DataFrame(rows) + + +def _resolve_checkpoint(recorded: Any, run_dir: Path) -> Path: + path = Path(str(recorded)) + candidates = ( + path, + run_dir / "epoch_checkpoints" / path.name, + run_dir / path.name, + ) + for candidate in candidates: + if candidate.is_file() and _is_within(candidate, run_dir): + return candidate.resolve() + raise CampaignValidationError( + f"could not resolve recorded checkpoint {recorded!r} below {run_dir}" + ) + + +def _checkpoint_index( + results_root: Path, + epoch_metrics: pd.DataFrame, + run_dirs: Mapping[tuple[str, int], Path], + completions: Mapping[tuple[str, int], Mapping[str, Any]], + *, + require_complete: bool, +) -> pd.DataFrame: + inventory: dict[tuple[str, int, Path], dict[str, Any]] = {} + + def register( + optimizer: str, + seed: int, + path: Path, + role: str, + *, + step: int | None = None, + epoch: float | None = None, + ) -> None: + resolved = path.resolve(strict=False) + key = (optimizer, int(seed), resolved) + record = inventory.setdefault(key, { + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "seed": int(seed), + "checkpoint_path": str(resolved), + "roles": set(), + "steps": set(), + "epochs": set(), + }) + record["roles"].add(role) + if step is not None: + record["steps"].add(int(step)) + if epoch is not None and math.isfinite(float(epoch)): + record["epochs"].add(float(epoch)) + + for (optimizer, seed), run_dir in run_dirs.items(): + if not run_dir.is_dir(): + continue + completion = completions.get((optimizer, seed), {}) + standard = { + "initial": run_dir / "checkpoint_initial.pt", + "latest": run_dir / "checkpoint_latest.pt", + "best": run_dir / "checkpoint_best.pt", + "final": run_dir / "checkpoint_final.pt", + } + for role, path in standard.items(): + if not path.is_file(): + if require_complete: + raise CampaignValidationError( + f"missing required {role} checkpoint: {path}" + ) + continue + step: int | None = None + if role in {"latest", "final"} and "optimizer_steps" in completion: + step = int(completion["optimizer_steps"]) + elif role == "best" and "best_validation_step" in completion: + step = int(completion["best_validation_step"]) + elif role == "initial": + step = 0 + register(optimizer, seed, path, role, step=step) + + for _, row in epoch_metrics.iterrows(): + optimizer = str(row["optimizer"]) + seed = int(row["seed"]) + run_dir = run_dirs[(optimizer, seed)] + try: + path = _resolve_checkpoint(row["checkpoint_path"], run_dir) + except CampaignValidationError: + if require_complete: + raise + continue + register( + optimizer, + seed, + path, + "permanent_epoch", + step=int(row["step"]), + epoch=float(row.get("nominal_epoch", row.get("epoch", np.nan))), + ) + + permanent_counts: defaultdict[tuple[str, int], int] = defaultdict(int) + rows: list[dict[str, Any]] = [] + for record in inventory.values(): + path = Path(record["checkpoint_path"]) + if not path.is_file(): + if require_complete: + raise CampaignValidationError(f"checkpoint disappeared: {path}") + continue + roles = sorted(record.pop("roles")) + steps = sorted(record.pop("steps")) + epochs = sorted(record.pop("epochs")) + if "permanent_epoch" in roles: + permanent_counts[(record["optimizer"], record["seed"])] += 1 + try: + relative = str(path.resolve().relative_to(results_root.resolve())) + except ValueError: + relative = str(path.resolve()) + rows.append({ + **record, + "checkpoint_relative_path": relative, + "roles": ",".join(roles), + "steps": ",".join(str(value) for value in steps), + "epochs": ",".join(f"{value:.12g}" for value in epochs), + "bytes": int(path.stat().st_size), + "sha256": _sha256(path), + }) + if require_complete: + for optimizer in OPTIMIZERS: + for seed in SEEDS: + count = permanent_counts[(optimizer, seed)] + if count < MIN_PERMANENT_CHECKPOINTS: + raise CampaignValidationError( + f"checkpoint index found {count} permanent checkpoints for " + f"{optimizer}/{seed}, expected at least {MIN_PERMANENT_CHECKPOINTS}" + ) + return pd.DataFrame(rows).sort_values( + ["optimizer", "seed", "checkpoint_relative_path"] + ).reset_index(drop=True) + + +def _curve_summary(frame: pd.DataFrame, metric: str) -> pd.DataFrame: + rows: list[dict[str, Any]] = [] + selected = frame[["epoch", "seed", metric]].copy() + selected[metric] = pd.to_numeric(selected[metric], errors="coerce") + selected = selected[np.isfinite(selected[metric])] + for epoch, group in selected.groupby("epoch", sort=True): + rows.append({"epoch": float(epoch), **_stats(group[metric])}) + return pd.DataFrame(rows) + + +def _plot_metric_curve( + axis: plt.Axes, + frame: pd.DataFrame, + metric: str, + *, + label: str, + color: str, + scale: float = 1.0, + linestyle: str = "-", +) -> None: + if metric not in frame.columns: + return + finite = frame[np.isfinite(pd.to_numeric(frame[metric], errors="coerce"))].copy() + if finite.empty: + return + finite[metric] = pd.to_numeric(finite[metric], errors="coerce") * float(scale) + for _, seed_frame in finite.groupby("seed", sort=True): + seed_frame = seed_frame.sort_values("epoch") + axis.plot( + seed_frame["epoch"].to_numpy(dtype=float), + seed_frame[metric].to_numpy(dtype=float), + color=color, + alpha=0.20, + linewidth=0.8, + linestyle=linestyle, + ) + summary = _curve_summary(finite, metric) + if summary.empty: + return + x = summary["epoch"].to_numpy(dtype=float) + mean = summary["mean"].to_numpy(dtype=float) + low = summary["ci95_low"].to_numpy(dtype=float) + high = summary["ci95_high"].to_numpy(dtype=float) + axis.plot( + x, + mean, + color=color, + linewidth=2.0, + linestyle=linestyle, + label=label, + ) + valid = np.isfinite(low) & np.isfinite(high) + if valid.any(): + axis.fill_between(x[valid], low[valid], high[valid], color=color, alpha=0.12) + + +def _plot_posthoc_test_metric( + axis: plt.Axes, + test_results: pd.DataFrame, + metric: str, + *, + title: str, + ylabel: str, +) -> None: + order = ("validation_selected", "final") + labels = ("validation-selected", "final") + for index, checkpoint in enumerate(order): + values = pd.to_numeric( + test_results.loc[ + test_results["checkpoint"].eq(checkpoint), metric + ], + errors="coerce", + ) + values = values[np.isfinite(values)] + if values.empty: + continue + x_values = np.full(len(values), float(index)) + axis.scatter( + x_values, + values.to_numpy(dtype=float), + color="#6B7280", + alpha=0.65, + s=24, + zorder=3, + ) + statistics = _stats(values) + axis.errorbar( + [index], + [statistics["mean"]], + yerr=[[statistics["ci95_half_width"]], [statistics["ci95_half_width"]]], + color="#D55E00", + marker="o", + capsize=4, + linewidth=1.5, + zorder=4, + ) + axis.set_xticks(range(len(order)), labels) + axis.set_title(title) + axis.set_ylabel(ylabel) + axis.grid(axis="y", alpha=0.25) + + +def _optimizer_performance_plot( + metrics: pd.DataFrame, + test_results: pd.DataFrame, + *, + optimizer: str, + path: Path, +) -> None: + frame = metrics[metrics["optimizer"].eq(optimizer)].copy() + if frame.empty: + return + figure, axes = plt.subplots(2, 3, figsize=(16, 9), sharex=False) + flat_axes = axes.ravel() + panels = ( + ( + "Cross-entropy loss", + ( + ("train_loss", "train", SPLIT_COLORS["train"], 1.0, "-"), + ("val_loss", "validation", SPLIT_COLORS["val"], 1.0, "-"), + ), + "NLL (nats/token)", + ), + ( + "Perplexity", + ( + ("train_perplexity", "train", SPLIT_COLORS["train"], 1.0, "-"), + ("val_perplexity", "validation", SPLIT_COLORS["val"], 1.0, "-"), + ), + "perplexity", + ), + ( + "Next-token top-1 accuracy", + ( + ("train_accuracy", "train", SPLIT_COLORS["train"], 100.0, "-"), + ("val_accuracy", "validation", SPLIT_COLORS["val"], 100.0, "-"), + ), + "token accuracy (%)", + ), + ( + "Generalization gaps", + ( + ("val_generalization_gap", "validation - train", SPLIT_COLORS["val"], 1.0, "-"), + ), + "loss gap (nats/token)", + ), + ) + for axis, (title, curves, ylabel) in zip(flat_axes[:4], panels, strict=True): + for metric, label, color, scale, linestyle in curves: + _plot_metric_curve( + axis, + frame, + metric, + label=label, + color=color, + scale=scale, + linestyle=linestyle, + ) + axis.set_title(title) + axis.set_ylabel(ylabel) + axis.set_xlabel("corpus-equivalent epoch") + axis.grid(alpha=0.25) + if axis.lines: + axis.legend(frameon=False, fontsize=8) + optimizer_tests = test_results[test_results["optimizer"].eq(optimizer)] + _plot_posthoc_test_metric( + flat_axes[4], + optimizer_tests, + "test_loss", + title="Held-out post-training test NLL", + ylabel="NLL (nats/token)", + ) + _plot_posthoc_test_metric( + flat_axes[5], + optimizer_tests, + "test_bleu", + title="Held-out greedy-continuation BLEU", + ylabel="corpus BLEU", + ) + figure.suptitle( + f"{OPTIMIZER_LABELS[optimizer]}: seeded runs (mean and 95% Student-t CI)", + fontsize=15, + ) + figure.tight_layout(rect=(0, 0, 1, 0.96)) + _atomic_figure(path, figure) + plt.close(figure) + + +def _alpha_variant_summary( + frame: pd.DataFrame, + metric: str, +) -> pd.DataFrame: + return _curve_summary(frame, metric) + + +def _optimizer_alpha_plot( + layers: pd.DataFrame, + *, + optimizer: str, + path: Path, +) -> None: + selected = layers[layers["optimizer"].eq(optimizer)].copy() + if selected.empty: + return + figure, axes = plt.subplots(2, 3, figsize=(16, 9), sharex=True) + variants = ( + ("alpha_raw", "raw alpha", "#6B7280", "--"), + ("alpha_clip_xmax", "clip_xmax alpha", OPTIMIZER_COLORS[optimizer], "-"), + ) + for axis, matrix in zip(axes.flat, MATRIX_TYPES, strict=True): + matrix_frame = selected[selected["matrix_type"].eq(matrix)].copy() + for metric, label, color, linestyle in variants: + for _, seed_frame in matrix_frame.groupby("seed", sort=True): + seed_frame = seed_frame.sort_values("epoch") + axis.plot( + seed_frame["epoch"].to_numpy(dtype=float), + seed_frame[metric].to_numpy(dtype=float), + color=color, + alpha=0.20, + linewidth=0.8, + linestyle=linestyle, + ) + summary = _alpha_variant_summary(matrix_frame, metric) + if summary.empty: + continue + x = summary["epoch"].to_numpy(dtype=float) + mean = summary["mean"].to_numpy(dtype=float) + low = summary["ci95_low"].to_numpy(dtype=float) + high = summary["ci95_high"].to_numpy(dtype=float) + axis.plot( + x, + mean, + color=color, + linewidth=2.0, + linestyle=linestyle, + label=label, + ) + valid = np.isfinite(low) & np.isfinite(high) + if valid.any(): + axis.fill_between( + x[valid], low[valid], high[valid], color=color, alpha=0.12 + ) + axis.axhline(2.0, color="black", linestyle=":", linewidth=1.0) + axis.set_title(matrix) + axis.set_ylabel("WeightWatcher alpha") + axis.grid(alpha=0.25) + axis.legend(frameon=False, fontsize=8) + for axis in axes[-1, :]: + axis.set_xlabel("corpus-equivalent epoch") + figure.suptitle( + f"{OPTIMIZER_LABELS[optimizer]}: raw versus clip_xmax alpha by matrix", + fontsize=15, + ) + figure.tight_layout(rect=(0, 0, 1, 0.96)) + _atomic_figure(path, figure) + plt.close(figure) + + +def _optimizer_erg_plot( + layers: pd.DataFrame, + *, + optimizer: str, + path: Path, +) -> None: + selected = layers[layers["optimizer"].eq(optimizer)].copy() + if selected.empty: + return + figure, axes = plt.subplots(2, 3, figsize=(16, 9), sharex=True) + diagnostics = ( + ("ERG_gap", "ERG gap", OPTIMIZER_COLORS[optimizer], "-"), + ("num_traps", "correlation traps", "#6B7280", "--"), + ) + for axis, matrix in zip(axes.flat, MATRIX_TYPES, strict=True): + matrix_frame = selected[selected["matrix_type"].eq(matrix)].copy() + for metric, label, color, linestyle in diagnostics: + _plot_metric_curve( + axis, + matrix_frame, + metric, + label=label, + color=color, + linestyle=linestyle, + ) + axis.set_title(matrix) + axis.set_ylabel("WeightWatcher diagnostic") + axis.grid(alpha=0.25) + if axis.lines: + axis.legend(frameon=False, fontsize=8) + for axis in axes[-1, :]: + axis.set_xlabel("corpus-equivalent epoch") + figure.suptitle( + f"{OPTIMIZER_LABELS[optimizer]}: ERG gap and correlation traps by matrix", + fontsize=15, + ) + figure.tight_layout(rect=(0, 0, 1, 0.96)) + _atomic_figure(path, figure) + plt.close(figure) + + +def _format_float(value: Any) -> str: + try: + numeric = float(value) + except (TypeError, ValueError): + return str(value) + if not math.isfinite(numeric): + return "" + if abs(numeric) >= 1_000 or (0 < abs(numeric) < 1e-3): + return f"{numeric:.4g}" + return f"{numeric:.5f}" + + +def _markdown_table(frame: pd.DataFrame, columns: Sequence[str]) -> str: + if frame.empty: + return "_No rows available._" + selected = frame.loc[:, [column for column in columns if column in frame.columns]].copy() + headers = list(selected.columns) + lines = [ + "| " + " | ".join(headers) + " |", + "| " + " | ".join("---" for _ in headers) + " |", + ] + for _, row in selected.iterrows(): + values = [ + str(_format_float(row[column])).replace("|", "\\|") + for column in headers + ] + lines.append("| " + " | ".join(values) + " |") + return "\n".join(lines) + + +def _table_html(frame: pd.DataFrame, columns: Sequence[str] | None = None) -> str: + if frame.empty: + return "

No rows available.

" + selected = frame if columns is None else frame[ + [column for column in columns if column in frame.columns] + ] + return selected.to_html( + index=False, + border=0, + classes="dataframe", + na_rep="", + float_format=lambda value: _format_float(value), + escape=True, + ) + + +def _latest_alpha_summary(alpha_summary: pd.DataFrame) -> pd.DataFrame: + if alpha_summary.empty: + return alpha_summary + latest_rows: list[pd.DataFrame] = [] + for optimizer, group in alpha_summary.groupby("optimizer", sort=True): + latest_epoch = group["epoch"].max() + latest_rows.append(group[group["epoch"].eq(latest_epoch)]) + return pd.concat(latest_rows, ignore_index=True, sort=False) + + +def _write_summary( + output_root: Path, + *, + campaign_runs: pd.DataFrame, + performance: pd.DataFrame, + paired: pd.DataFrame, + alpha_summary: pd.DataFrame, + saturation: pd.DataFrame, + qk_summary: pd.DataFrame, + warnings: Sequence[str], +) -> Path: + primary_performance = performance[ + performance["source"].eq("test_results") + & performance["checkpoint"].eq("validation_selected") + ].copy() + primary_performance = primary_performance[ + primary_performance["metric"].isin(TEST_RESULT_METRICS) + ] + primary_pairs = paired[ + paired["source"].eq("test_results") + & paired["checkpoint"].eq("validation_selected") + & paired["metric"].isin(("test_loss", "test_accuracy")) + ].copy() + latest_alpha = _latest_alpha_summary(alpha_summary) + lines = [ + "# One-head nanoGPT baseline — 2026-08-21", + "", + "This report validates the exact AdamW/MuonClip × five-seed " + "campaign. The complete seeded run is the unit of replication. Test " + "results are reported only as post-run outcomes; the saturation " + "diagnostic uses validation loss exclusively.", + "", + "Metric definitions: accuracy is fixed-probe next-token top-1 " + "accuracy (not classification or whole-sequence accuracy); top-5 " + "accuracy is reported separately. Perplexity is exp(mean token NLL) " + "and bits/token is NLL/log(2). BLEU is fixed greedy held-out " + "continuation BLEU, a secondary language-model diagnostic—not " + "translation BLEU. Continuation token accuracy and exact match are " + "reported alongside it.", + "", + "## Campaign status", + "", + _markdown_table( + campaign_runs, + ( + "optimizer", + "seed", + "complete", + "optimizer_steps", + "train_epochs", + "best_validation_loss", + "final_test_loss", + "accelerator", + "git_commit", + ), + ), + "", + "## Validation-selected test performance", + "", + _markdown_table( + primary_performance, + ( + "optimizer", + "metric", + "n", + "mean", + "sd", + "ci95_low", + "ci95_high", + ), + ), + "", + "## Paired seed differences", + "", + "Differences are always `optimizer_b - optimizer_a`; every valid row " + "uses the same five seeds and a df=4 Student-t interval.", + "", + _markdown_table( + primary_pairs, + ( + "metric", + "contrast", + "n", + "mean", + "sd", + "ci95_low", + "ci95_high", + ), + ), + "", + "## Validation-only saturation", + "", + _markdown_table( + saturation, + ( + "optimizer", + "seed", + "plateau_detected", + "plateau_assessment_end_epoch", + "degradation_detected", + "first_degradation_end_epoch", + "best_validation_epoch", + "best_validation_loss", + "final_validation_loss", + "test_metrics_used", + ), + ), + "", + "## Final six-matrix alpha summaries", + "", + "Each seeded value is first the median over exactly six transformer " + "matrices. The displayed interval is then computed across the five " + "seeded run medians.", + "", + _markdown_table( + latest_alpha, + ( + "optimizer", + "epoch", + "metric", + "n", + "mean", + "sd", + "ci95_low", + "ci95_high", + ), + ), + "", + "## MuonClip QK diagnostics", + "", + _markdown_table( + qk_summary, + ( + "optimizer", + "seed", + "active_fraction_weighted", + "max_logit_observed", + "min_gamma_observed", + ), + ), + "", + "## Artifacts", + "", + ] + for filename in CSV_ARTIFACTS: + lines.append(f"- `{filename}`") + lines.extend([ + "- `plots/`", + "- `report.html`", + "- `results_manifest.json`", + ]) + if warnings: + lines.extend(["", "## Warnings", ""]) + lines.extend(f"- {warning}" for warning in warnings) + destination = output_root / "SUMMARY.md" + _atomic_write_text(destination, "\n".join(lines) + "\n") + return destination + + +def _write_html_report( + output_root: Path, + *, + campaign_runs: pd.DataFrame, + performance: pd.DataFrame, + paired: pd.DataFrame, + alpha_summary: pd.DataFrame, + saturation: pd.DataFrame, + qk_summary: pd.DataFrame, + checkpoint_index: pd.DataFrame, + warnings: Sequence[str], +) -> Path: + selected_performance = performance[ + performance["source"].eq("test_results") + ] + selected_pairs = paired[ + paired["source"].eq("test_results") + & paired["metric"].isin(("test_loss", "test_accuracy")) + ] + latest_alpha = _latest_alpha_summary(alpha_summary) + artifact_links = "\n".join( + f'
  • {html.escape(filename)}
    • ' + for filename in CSV_ARTIFACTS + ) + image_sections: list[str] = [] + for optimizer in OPTIMIZERS: + label = html.escape(OPTIMIZER_LABELS[optimizer]) + performance_path = f"plots/{optimizer}_performance.png" + alpha_path = f"plots/{optimizer}_alpha_raw_vs_clip_xmax.png" + erg_path = f"plots/{optimizer}_erg_gap_num_traps.png" + image_sections.append( + f"

    {label}

    " + f'{label} performance trajectories' + f'{label} raw versus clipped alpha' + f'{label} ERG gap and correlation traps' + ) + warning_html = ( + "

    Warnings

      " + + "".join(f"
    • {html.escape(value)}
      • " for value in warnings) + + "
    " + if warnings + else "" + ) + document = f""" + + + + +One-head nanoGPT baseline — 2026-08-21 + + + +

    One-head nanoGPT baseline — 2026-08-21

    +

    Exact campaign: AdamW and MuonClip; seeds 1337, 2027, 4099, 31415, and 271828. The run is the replicate. Test quantities never enter the saturation calculation. A 95% interval with n=5 uses Student-t critical value 2.77645.

    +

    Campaign validation

    +{_table_html(campaign_runs, ("optimizer", "seed", "complete", "optimizer_steps", "train_epochs", "best_validation_loss", "final_test_loss", "accelerator", "torch_version", "git_commit"))} +

    Test performance

    +

    Accuracy is next-token top-1 accuracy on a fixed held-out probe, not a classification or sequence-level accuracy. Top-5 accuracy is separate. Perplexity is exp(mean token NLL), and bits/token is NLL/log(2). BLEU is a secondary fixed greedy held-out continuation diagnostic—not translation BLEU—and is accompanied by continuation token accuracy and exact match.

    +{_table_html(selected_performance, ("optimizer", "checkpoint", "metric", "n", "mean", "sd", "ci95_low", "ci95_high"))} +

    Paired seeded differences

    +

    Every difference is optimizer_b - optimizer_a. Positive and negative values must be interpreted according to whether higher or lower is preferable for the metric.

    +{_table_html(selected_pairs, ("checkpoint", "metric", "contrast", "n", "mean", "sd", "ci95_low", "ci95_high", "valid_exact_n5"))} +

    Validation-only saturation diagnostic

    +

    Plateau means |one-epoch validation-NLL improvement| ≤ {SATURATION_DELTA_NATS:.3f} nat/token for two consecutive intervals. Runs are not stopped by this diagnostic.

    +{_table_html(saturation, ("optimizer", "seed", "plateau_detected", "plateau_assessment_end_epoch", "degradation_detected", "first_degradation_end_epoch", "best_validation_epoch", "best_validation_loss", "final_validation_loss", "test_metrics_used"))} +

    Final WeightWatcher alpha

    +

    The six layer values are reduced to one median inside each seeded run before any across-seed mean or confidence interval is computed.

    +{_table_html(latest_alpha, ("optimizer", "epoch", "metric", "n", "mean", "sd", "ci95_low", "ci95_high", "valid_exact_n5"))} +

    MuonClip QK diagnostics

    +{_table_html(qk_summary)} +

    Plots

    +{''.join(image_sections)} +

    Checkpoint integrity

    +

    {len(checkpoint_index):,} checkpoint files were indexed by byte size and SHA-256. See checkpoint_sha256.csv.

    +

    Machine-readable artifacts

    +
      {artifact_links}
    +{warning_html} + + +""" + destination = output_root / "report.html" + _atomic_write_text(destination, document) + return destination + + +def _input_artifact_manifest( + results_root: Path, + *, + run_dirs: Mapping[tuple[str, int], Path], + checkpoint_index: pd.DataFrame, +) -> list[dict[str, Any]]: + candidates: set[Path] = set() + for _identity, run_dir in run_dirs.items(): + if not run_dir.is_dir(): + continue + for relative in REQUIRED_RUN_FILES: + path = run_dir / relative + if path.is_file(): + candidates.add(path.resolve()) + qk_path = run_dir / "muonclip_qk.csv" + if qk_path.is_file(): + candidates.add(qk_path.resolve()) + candidates.update(path.resolve() for path in (run_dir / "spectral").glob( + "status_step_*.json" + ) if path.is_file()) + if "checkpoint_path" in checkpoint_index.columns: + candidates.update( + Path(str(value)).resolve() + for value in checkpoint_index["checkpoint_path"] + if Path(str(value)).is_file() + ) + + rows: list[dict[str, Any]] = [] + resolved_root = results_root.resolve() + for path in sorted(candidates): + if not _is_within(path, resolved_root): + raise CampaignValidationError( + f"report input artifact escapes results root: {path}" + ) + rows.append({ + "path": str(path.relative_to(resolved_root)), + "bytes": int(path.stat().st_size), + "sha256": _sha256(path), + }) + return rows + + +def _artifact_manifest( + output_root: Path, + *, + results_root: Path, + args: argparse.Namespace, + campaign_runs: pd.DataFrame, + checkpoint_index: pd.DataFrame, + input_artifacts: Sequence[Mapping[str, Any]], + source_git_commit: str, + warnings: Sequence[str], +) -> dict[str, Any]: + manifest_path = output_root / "results_manifest.json" + artifacts: list[dict[str, Any]] = [] + for path in sorted(output_root.rglob("*")): + relative = path.relative_to(output_root) + if ( + not path.is_file() + or path == manifest_path + or path.name.endswith(".tmp") + or (relative.parts and relative.parts[0] == "notebooks") + ): + continue + artifacts.append({ + "path": str(relative), + "bytes": int(path.stat().st_size), + "sha256": _sha256(path), + }) + run_records = campaign_runs.to_dict("records") + return { + "schema_version": 2, + "campaign": "nanogpt_one_head_2026_08_21_baseline", + "generated_at_utc": _utc_now(), + "report_builder": { + "path": str(Path(__file__).resolve()), + "sha256": _sha256(Path(__file__).resolve()), + }, + "source_git_commit": source_git_commit, + "frozen_config": { + "path": str(FROZEN_CONFIG.relative_to(REPOSITORY_ROOT)), + "canonical_sha256": FROZEN_CONFIG_SHA256, + }, + "results_root": str(results_root.resolve()), + "output_root": str(output_root.resolve()), + "exact_campaign": { + "optimizers": list(OPTIMIZERS), + "seeds": list(SEEDS), + "expected_run_count": len(OPTIMIZERS) * len(SEEDS), + "require_complete": _requires_complete(args), + "allow_extra_runs": bool(args.allow_extra_runs), + "allow_mixed_runtime": False, + }, + "statistical_contract": { + "replicate": "complete seeded run", + "paired_seed_differences": True, + "paired_n": 5, + "student_t_critical_95_df4": T_975_DF4, + "alpha_aggregation": ( + "six-matrix median within run, then across-seed Student-t summary" + ), + "saturation_uses_test_metrics": False, + "saturation_delta_nats": SATURATION_DELTA_NATS, + }, + "run_records": run_records, + "valid_run_count": int( + campaign_runs.get("complete", pd.Series(dtype=bool)) + .fillna(False) + .astype(bool) + .sum() + ), + "checkpoint_file_count": int(len(checkpoint_index)), + "checkpoint_total_bytes": int( + pd.to_numeric(checkpoint_index.get("bytes", pd.Series(dtype=float)), errors="coerce") + .fillna(0) + .sum() + ), + "input_artifacts": [dict(record) for record in input_artifacts], + "artifacts": artifacts, + "warnings": list(warnings), + "command": list(sys.argv), + } + + +def build_report(args: argparse.Namespace) -> Path: + results_root = Path(args.results_root).resolve(strict=False) + output_root = Path(args.output_root).resolve(strict=False) + _validate_paths(results_root, output_root) + if bool(getattr(args, "allow_mixed_runtime", False)): + raise CampaignValidationError( + "mixed-runtime reporting is disabled because the registered " + "paired statistics require one homogeneous hardware/runtime block" + ) + require_complete = _requires_complete(args) + output_root.mkdir(parents=True, exist_ok=True) + (output_root / "plots").mkdir(parents=True, exist_ok=True) + _initialize_matplotlib(output_root) + + campaign = _load_campaign( + results_root, + require_complete=require_complete, + allow_extra_runs=args.allow_extra_runs, + allow_mixed_runtime=False, + ) + campaign_runs = campaign["campaign_runs"] + metrics = campaign["metrics"] + epoch_metrics = campaign["epoch_metrics"] + layers = campaign["spectral_layers"] + spectral_summary = campaign["spectral_summary"] + test_results = campaign["test_results"] + qk = campaign["qk"] + + performance = _summarize_performance(test_results, epoch_metrics) + paired = _paired_seed_differences( + test_results, + epoch_metrics, + require_complete=require_complete, + ) + alpha_run_medians, alpha_summary = _alpha_run_medians(layers) + saturation, saturation_integer, saturation_summary = _validation_saturation( + epoch_metrics + ) + qk_summary = _qk_summary(qk) + + print("[report] hashing checkpoints; this can take several minutes", flush=True) + checkpoint_index = _checkpoint_index( + results_root, + epoch_metrics, + campaign["run_dirs"], + campaign["completions"], + require_complete=require_complete, + ) + + table_map = { + "campaign_runs.csv": campaign_runs, + "metrics_all.csv": metrics, + "epoch_metrics_all.csv": epoch_metrics, + "spectral_layers_all.csv": layers, + "spectral_summary_all.csv": spectral_summary, + "test_results_all.csv": test_results, + "qk_diagnostics_all.csv": qk, + "qk_summary.csv": qk_summary, + "performance_summary.csv": performance, + "paired_seed_differences.csv": paired, + "alpha_run_medians.csv": alpha_run_medians, + "alpha_across_seed_summary.csv": alpha_summary, + "saturation_diagnostics.csv": saturation, + "saturation_integer_epoch_validation.csv": saturation_integer, + "saturation_across_seed_summary.csv": saturation_summary, + "checkpoint_sha256.csv": checkpoint_index, + } + for filename, frame in table_map.items(): + _atomic_csv(output_root / filename, frame) + + for optimizer in OPTIMIZERS: + _optimizer_performance_plot( + metrics, + test_results, + optimizer=optimizer, + path=output_root / "plots" / f"{optimizer}_performance.png", + ) + _optimizer_alpha_plot( + layers, + optimizer=optimizer, + path=( + output_root + / "plots" + / f"{optimizer}_alpha_raw_vs_clip_xmax.png" + ), + ) + _optimizer_erg_plot( + layers, + optimizer=optimizer, + path=output_root / "plots" / f"{optimizer}_erg_gap_num_traps.png", + ) + + warnings = [*campaign["validation_errors"]] + if campaign["unexpected_runs"]: + warnings.append( + "Unexpected seed directories were ignored: " + + ", ".join(campaign["unexpected_runs"]) + ) + if not require_complete: + warnings.append( + "Incomplete runs were permitted; rows with n<5 are not valid exact " + "five-seed comparisons." + ) + + _write_summary( + output_root, + campaign_runs=campaign_runs, + performance=performance, + paired=paired, + alpha_summary=alpha_summary, + saturation=saturation, + qk_summary=qk_summary, + warnings=warnings, + ) + report_path = _write_html_report( + output_root, + campaign_runs=campaign_runs, + performance=performance, + paired=paired, + alpha_summary=alpha_summary, + saturation=saturation, + qk_summary=qk_summary, + checkpoint_index=checkpoint_index, + warnings=warnings, + ) + input_artifacts = _input_artifact_manifest( + results_root, + run_dirs=campaign["run_dirs"], + checkpoint_index=checkpoint_index, + ) + source_git_commit = str( + _campaign_invariants(campaign["manifests"][0]).get("git_commit", "") + ) + manifest = _artifact_manifest( + output_root, + results_root=results_root, + args=args, + campaign_runs=campaign_runs, + checkpoint_index=checkpoint_index, + input_artifacts=input_artifacts, + source_git_commit=source_git_commit, + warnings=warnings, + ) + _atomic_json(output_root / "results_manifest.json", manifest) + print(f"[report] complete: {report_path}", flush=True) + return report_path + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Validate and report the exact 2026-08-21 one-head nanoGPT " + "AdamW/MuonClip five-seed campaign" + ) + ) + parser.add_argument( + "--results-root", + required=True, + help=( + "input results directory containing optimizer/seed_ runs; " + "must be strictly below resolved /tmp or /private/tmp" + ), + ) + parser.add_argument( + "--output-root", + required=True, + help="report destination strictly below resolved /tmp or /private/tmp", + ) + completion_group = parser.add_mutually_exclusive_group() + completion_group.add_argument( + "--require-complete", + dest="require_complete", + action="store_true", + help=( + "require all exact twenty completed runs (the default; this explicit " + "form is used by the executed notebook contract)" + ), + ) + completion_group.add_argument( + "--allow-incomplete", + dest="require_complete", + action="store_false", + help=( + "build a diagnostic report from available runs; default behavior " + "requires all exact twenty completed runs" + ), + ) + parser.set_defaults(require_complete=True) + parser.add_argument( + "--allow-extra-runs", + action="store_true", + help="ignore additional seed_* directories under the two optimizer roots", + ) + return parser + + +def main() -> None: + args = _parser().parse_args() + try: + build_report(args) + except CampaignValidationError as exc: + print(f"[report] ERROR: {exc}", file=sys.stderr, flush=True) + raise SystemExit(2) from exc + + +if __name__ == "__main__": + main() diff --git a/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/run_experiment.py b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/run_experiment.py new file mode 100644 index 00000000..34d0ddf8 --- /dev/null +++ b/baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/run_experiment.py @@ -0,0 +1,3269 @@ +#!/usr/bin/env python3 +"""Safe, reproducible launcher for the 2026-08-21 one-head nanoGPT campaign. + +The launcher intentionally imports only the Python standard library at module +load time. This keeps path-policy tests and ``--help`` usable before the +scientific environment has been installed. Scientific dependencies are +loaded only by the command that needs them. +""" + +from __future__ import annotations + +import argparse +import csv +from datetime import datetime, timezone +import hashlib +import importlib.metadata +import json +import math +import os +from pathlib import Path +import platform +import re +import shlex +import shutil +import subprocess +import sys +import tempfile +import time +from typing import Any, Mapping, Sequence + + +EXPERIMENT_ROOT_ENV = "RG_NANOGPT_EXPERIMENT_ROOT" +CANONICAL_OPTIMIZERS = ("adamw", "muon_clip") +CANONICAL_SEEDS = (1337, 2027, 4099, 31415, 271828) +EXPECTED_REPLICATES = 10 +EXPECTED_MATRICES = 6 +MINIMUM_PERMANENT_STATES = 10 +PINNED_WEIGHTWATCHER = "0.7.7" +PINNED_PACKAGE_VERSION = "0.5.1" +FROZEN_CONFIG_SHA256 = ( + "ebbbdfa30efe96b0b0c1c68ae4fc81909361502d89ad336d1181d00fcb85876a" +) +MANIFEST_PACKAGES = ( + "python", + "rg-nanogpt-one-head", + "torch", + "torch-xla", + "numpy", + "pandas", + "scipy", + "PyYAML", + "datasets", + "tiktoken", + "sacrebleu", + "weightwatcher", + "powerlaw", + "papermill", + "packaging", +) + +SCRIPT_PATH = Path(__file__).resolve() +EXPERIMENT_DIR = SCRIPT_PATH.parents[1] +REPOSITORY_ROOT = SCRIPT_PATH.parents[4] +NANOGPT_ROOT = REPOSITORY_ROOT / "baseline" / "nanogpt_one_head" +DEFAULT_CONFIG = EXPERIMENT_DIR / "configs" / "baseline.yaml" + +REQUIRED_RUN_FILES = ( + "run_complete.json", + "manifest.json", + "metrics.csv", + "epoch_metrics.csv", + "checkpoint_initial.pt", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "test_results.json", + "spectral/layers.csv", + "spectral/summary.csv", +) + +REQUIRED_ANALYSIS_FILES = ( + "SUMMARY.md", + "report.html", + "results_manifest.json", + "campaign_runs.csv", + "metrics_all.csv", + "epoch_metrics_all.csv", + "spectral_layers_all.csv", + "spectral_summary_all.csv", + "test_results_all.csv", + "qk_diagnostics_all.csv", + "qk_summary.csv", + "performance_summary.csv", + "paired_seed_differences.csv", + "alpha_run_medians.csv", + "alpha_across_seed_summary.csv", + "saturation_diagnostics.csv", + "saturation_integer_epoch_validation.csv", + "saturation_across_seed_summary.csv", + "checkpoint_sha256.csv", + "plots/adamw_performance.png", + "plots/adamw_alpha_raw_vs_clip_xmax.png", + "plots/adamw_erg_gap_num_traps.png", + "plots/muon_clip_performance.png", + "plots/muon_clip_alpha_raw_vs_clip_xmax.png", + "plots/muon_clip_erg_gap_num_traps.png", + "notebooks/01_Performance_and_Spectra.executed.ipynb", +) + +DEPENDENCIES = { + "rg-nanogpt-one-head": "rg_nanogpt_one_head", + "torch": "torch", + "numpy": "numpy", + "pandas": "pandas", + "scipy": "scipy", + "PyYAML": "yaml", + "weightwatcher": "weightwatcher", + "powerlaw": "powerlaw", + "datasets": "datasets", + "tiktoken": "tiktoken", + "sacrebleu": "sacrebleu", + "matplotlib": "matplotlib", + "jupyter": "jupyter", + "ipykernel": "ipykernel", + "nbformat": "nbformat", + "papermill": "papermill", + "packaging": "packaging", +} + + +class CampaignError(ValueError): + """An actionable campaign-policy or artifact-validation failure.""" + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _strict_descendant(path: Path, parent: Path) -> bool: + try: + relative = path.relative_to(parent) + except ValueError: + return False + return relative != Path(".") + + +def resolve_experiment_root( + env: Mapping[str, str] | None = None, +) -> Path: + """Resolve and validate the required campaign output root. + + The value must be an absolute *strict descendant* of the resolved ``/tmp`` + or ``/private/tmp`` root. Resolving before validation prevents an existing + symlink or ``..`` component from escaping. A root that is HOME (or below + HOME) is rejected independently, even if a caller has placed HOME in tmp. + No directory is created by this function. + """ + + source = os.environ if env is None else env + raw = source.get(EXPERIMENT_ROOT_ENV) + if raw is None or not str(raw).strip(): + raise CampaignError( + f"{EXPERIMENT_ROOT_ENV} is required and must name an absolute " + "directory strictly below /tmp or /private/tmp" + ) + value = str(raw).strip() + if "~" in value: + raise CampaignError( + f"{EXPERIMENT_ROOT_ENV} must not contain '~': {value!r}" + ) + candidate = Path(value) + if not candidate.is_absolute(): + raise CampaignError( + f"{EXPERIMENT_ROOT_ENV} must be absolute, observed {value!r}" + ) + resolved = candidate.resolve(strict=False) + allowed_roots = { + Path("/tmp").resolve(strict=False), + Path("/private/tmp").resolve(strict=False), + } + if not any(_strict_descendant(resolved, root) for root in allowed_roots): + raise CampaignError( + f"{EXPERIMENT_ROOT_ENV} resolved to {resolved}; it must be " + "strictly below /tmp or /private/tmp" + ) + + home_raw = source.get("HOME") + if home_raw: + home = Path(str(home_raw)).expanduser().resolve(strict=False) + if resolved == home or _strict_descendant(resolved, home): + raise CampaignError( + f"{EXPERIMENT_ROOT_ENV} must never be HOME or below HOME" + ) + return resolved + + +def _within(path: Path, root: Path) -> bool: + resolved = path.resolve(strict=False) + return resolved == root or _strict_descendant(resolved, root) + + +def _require_within(path: Path, root: Path, label: str) -> Path: + resolved = path.resolve(strict=False) + if not _within(resolved, root): + raise CampaignError(f"{label} escapes experiment root: {resolved}") + return resolved + + +def _paths(root: Path) -> dict[str, Path]: + paths = { + "root": root, + "data": root / "data", + "results": root / "results", + "logs": root / "logs", + "analysis": root / "analysis", + "tables": root / "analysis" / "tables", + "plots": root / "analysis" / "plots", + "provenance": root / "provenance", + "cache": root / "cache", + "tmp": root / "tmp", + } + for label, path in paths.items(): + _require_within(path, root, label) + return paths + + +def _create_runtime_directories(paths: Mapping[str, Path]) -> None: + for path in paths.values(): + path.mkdir(parents=True, exist_ok=True) + for relative in ( + "huggingface/datasets", + "huggingface/hub", + "huggingface/assets", + "huggingface/modules", + "huggingface/transformers", + "tiktoken", + "matplotlib", + "xdg/cache", + "xdg/config", + "xdg/data", + "xdg/state", + "torch", + "torch_extensions", + "torchinductor", + "cuda", + "triton", + "cupy", + "xla", + "pip", + "uv", + "numba", + "joblib", + "keras", + "sacrebleu", + "wandb/cache", + "wandb/config", + "wandb/data", + "pycache", + "jupyter/config", + "jupyter/data", + "jupyter/runtime", + "ipython", + "home", + ): + (paths["cache"] / relative).mkdir(parents=True, exist_ok=True) + + +def _child_environment(root: Path, paths: Mapping[str, Path]) -> dict[str, str]: + cache = paths["cache"] + additions = { + EXPERIMENT_ROOT_ENV: str(root), + "RG_NANOGPT_ONE_HEAD_ROOT": str(root), + "RG_NANOGPT_ONE_HEAD_DATA_ROOT": str(paths["data"]), + "RG_NANOGPT_ONE_HEAD_RESULTS_ROOT": str(paths["results"]), + "RG_NANOGPT_ONE_HEAD_PLOTS_ROOT": str(paths["plots"]), + "HF_HOME": str(cache / "huggingface"), + "HF_DATASETS_CACHE": str(cache / "huggingface" / "datasets"), + "HUGGINGFACE_HUB_CACHE": str(cache / "huggingface" / "hub"), + "HF_HUB_CACHE": str(cache / "huggingface" / "hub"), + "HF_ASSETS_CACHE": str(cache / "huggingface" / "assets"), + "HF_MODULES_CACHE": str(cache / "huggingface" / "modules"), + "TRANSFORMERS_CACHE": str(cache / "huggingface" / "transformers"), + "TIKTOKEN_CACHE_DIR": str(cache / "tiktoken"), + "MPLCONFIGDIR": str(cache / "matplotlib"), + "XDG_CACHE_HOME": str(cache / "xdg" / "cache"), + "XDG_CONFIG_HOME": str(cache / "xdg" / "config"), + "XDG_DATA_HOME": str(cache / "xdg" / "data"), + "XDG_STATE_HOME": str(cache / "xdg" / "state"), + "TORCH_HOME": str(cache / "torch"), + "TORCH_EXTENSIONS_DIR": str(cache / "torch_extensions"), + "TORCHINDUCTOR_CACHE_DIR": str(cache / "torchinductor"), + "CUDA_CACHE_PATH": str(cache / "cuda"), + "TRITON_CACHE_DIR": str(cache / "triton"), + "CUPY_CACHE_DIR": str(cache / "cupy"), + "XLA_PERSISTENT_CACHE_PATH": str(cache / "xla"), + "PIP_CACHE_DIR": str(cache / "pip"), + "UV_CACHE_DIR": str(cache / "uv"), + "NUMBA_CACHE_DIR": str(cache / "numba"), + "JOBLIB_TEMP_FOLDER": str(cache / "joblib"), + "KERAS_HOME": str(cache / "keras"), + "SACREBLEU": str(cache / "sacrebleu"), + "WANDB_CACHE_DIR": str(cache / "wandb" / "cache"), + "WANDB_CONFIG_DIR": str(cache / "wandb" / "config"), + "WANDB_DATA_DIR": str(cache / "wandb" / "data"), + "PYTHONPYCACHEPREFIX": str(cache / "pycache"), + "JUPYTER_CONFIG_DIR": str(cache / "jupyter" / "config"), + "JUPYTER_DATA_DIR": str(cache / "jupyter" / "data"), + "JUPYTER_RUNTIME_DIR": str(cache / "jupyter" / "runtime"), + "IPYTHONDIR": str(cache / "ipython"), + # Some third-party libraries ignore XDG/cache-specific variables. Give + # child processes an ephemeral HOME inside the required /tmp campaign + # root so even those fallbacks can never touch the user's real home. + "HOME": str(cache / "home"), + "TMPDIR": str(paths["tmp"]), + "TMP": str(paths["tmp"]), + "TEMP": str(paths["tmp"]), + "PYTHONUNBUFFERED": "1", + "PYTORCH_ENABLE_MPS_FALLBACK": "1", + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + "MPLBACKEND": "Agg", + "TOKENIZERS_PARALLELISM": "false", + } + for name, value in additions.items(): + if name in { + "PYTHONUNBUFFERED", + "PYTORCH_ENABLE_MPS_FALLBACK", + "CUBLAS_WORKSPACE_CONFIG", + "MPLBACKEND", + "TOKENIZERS_PARALLELISM", + }: + continue + _require_within(Path(value), root, f"environment variable {name}") + + child = os.environ.copy() + child.update(additions) + source_path = str(NANOGPT_ROOT / "src") + # Do not inherit arbitrary module-shadowing paths whose distribution + # metadata can look pinned while different source code is imported. + child["PYTHONPATH"] = source_path + return child + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(1024 * 1024): + digest.update(block) + return digest.hexdigest() + + +def _canonical_json(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), default=str) + + +def _canonical_sha256(value: Any) -> str: + return hashlib.sha256(_canonical_json(value).encode("utf-8")).hexdigest() + + +def _runtime_block_identity(runtime: Mapping[str, Any]) -> dict[str, Any]: + """Return the comparison identity for one declared hardware block. + + Auto identities remain exact, including a CUDA UUID. A collaborator may + deliberately assign the same user block ID to homogeneous machines; in + that case host/install-path and accelerator-instance fields are recorded + but do not prevent pooling distinct complete seeds. + """ + + identity = dict(runtime) + if str(identity.get("hardware_block_id_source", "")) == "user": + for field in ( + "python_executable", + "processor", + "cuda_device_uuid", + "cuda_device_count", + "xla_process_index", + ): + identity.pop(field, None) + return identity + + +def _atomic_json(path: Path, payload: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + prefix=path.name + ".", suffix=".tmp", dir=path.parent + ) + temporary = Path(temporary_name) + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, sort_keys=True, default=str) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def _atomic_text(path: Path, value: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + prefix=path.name + ".", suffix=".tmp", dir=path.parent + ) + temporary = Path(temporary_name) + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as handle: + handle.write(value) + if value and not value.endswith("\n"): + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def _append_jsonl(path: Path, payload: Mapping[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + record = json.dumps(dict(payload), sort_keys=True, default=str) + "\n" + descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600) + try: + os.write(descriptor, record.encode("utf-8")) + finally: + os.close(descriptor) + + +def _acquire_exclusive_lock(path: Path): + try: + import fcntl + except ImportError as exc: # pragma: no cover - campaign targets Unix hosts + raise CampaignError( + "campaign subprocess locks require a Unix fcntl implementation" + ) from exc + path.parent.mkdir(parents=True, exist_ok=True) + handle = path.open("a+", encoding="utf-8") + try: + fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as exc: + handle.seek(0) + owner = handle.read().strip() or "unknown owner" + handle.close() + raise CampaignError( + f"another campaign process holds {path}: {owner}" + ) from exc + handle.seek(0) + handle.truncate() + handle.write( + json.dumps( + {"pid": os.getpid(), "started_at_utc": _utc_now()}, + sort_keys=True, + ) + + "\n" + ) + handle.flush() + os.fsync(handle.fileno()) + return handle + + +def _release_exclusive_lock(handle) -> None: + import fcntl + + try: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + finally: + handle.close() + + +def _git(arguments: Sequence[str], *, check: bool = True) -> subprocess.CompletedProcess[str]: + try: + return subprocess.run( + ["git", "-C", str(REPOSITORY_ROOT), *arguments], + check=check, + capture_output=True, + text=True, + timeout=30, + ) + except (OSError, subprocess.SubprocessError) as exc: + raise CampaignError(f"git command failed: {exc}") from exc + + +def _git_provenance() -> dict[str, Any]: + try: + commit = _git(("rev-parse", "HEAD")).stdout.strip() + branch = _git(("rev-parse", "--abbrev-ref", "HEAD")).stdout.strip() + describe = _git(("describe", "--tags", "--always", "--dirty")).stdout.strip() + tags = [ + line + for line in _git(("tag", "--points-at", "HEAD")).stdout.splitlines() + if line.strip() + ] + status = _git(("status", "--porcelain=v1", "--untracked-files=all")).stdout + remote_result = _git(("remote", "get-url", "origin"), check=False) + return { + "available": True, + "commit": commit, + "branch": branch, + "describe": describe, + "tags_at_commit": tags, + "tag_status": ",".join(tags) if tags else "untagged", + "origin_url": remote_result.stdout.strip() if remote_result.returncode == 0 else None, + "clean": not bool(status.strip()), + "dirty": bool(status.strip()), + "status_sha256": hashlib.sha256(status.encode("utf-8")).hexdigest(), + "status_lines": len(status.splitlines()), + } + except CampaignError: + return {"available": False, "clean": False, "dirty": None} + + +def _require_clean_git() -> dict[str, Any]: + provenance = _git_provenance() + if not provenance.get("available"): + raise CampaignError("a readable Git checkout is required") + if not provenance.get("clean"): + status = _git(("status", "--short", "--untracked-files=all")).stdout.strip() + preview = "\n".join(status.splitlines()[:20]) + raise CampaignError( + "production commands require a clean Git worktree. Commit or " + f"remove all changes first. Current status:\n{preview}" + ) + return provenance + + +def _dependency_versions() -> dict[str, str | None]: + versions: dict[str, str | None] = {} + for distribution in DEPENDENCIES: + try: + versions[distribution] = importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + versions[distribution] = None + return versions + + +def _normalized_distribution_name(value: str) -> str: + return re.sub(r"[-_.]+", "-", str(value).strip()).lower() + + +def _installed_distribution_lock( + scientific_packages: Mapping[str, str], +) -> tuple[dict[str, Any], str]: + """Lock the installed dependency closure and reject opaque origins. + + Unrelated packages in a long-lived conda environment are retained by raw + ``pip freeze`` but do not belong in the replay contract. Conversely, every + installed transitive dependency reachable from a campaign requirement is + locked. PEP 610 direct/VCS/file installs are rejected because ``Name==ver`` + would silently discard the only information capable of reproducing them. + """ + + try: + from packaging.requirements import InvalidRequirement, Requirement + except ImportError as exc: + raise CampaignError( + "packaging is required to construct the dependency closure lock" + ) from exc + + project_name = _normalized_distribution_name("rg-nanogpt-one-head") + packages: dict[str, dict[str, str]] = {} + pending = [ + name + for name, version in scientific_packages.items() + if name != "python" + and str(version) != "not-installed" + and _normalized_distribution_name(name) != project_name + ] + visited: set[str] = set() + while pending: + requested_name = pending.pop() + requested_normalized = _normalized_distribution_name(requested_name) + if requested_normalized in visited: + continue + visited.add(requested_normalized) + try: + distribution = importlib.metadata.distribution(requested_name) + except importlib.metadata.PackageNotFoundError: + # Optional requirements and platform markers may name packages not + # installed on this hardware block. + continue + name = str(distribution.metadata.get("Name", "")).strip() + version = str(distribution.version).strip() + if not name or not version or "\n" in name or "\n" in version: + raise CampaignError( + "installed distribution has unsafe or missing name/version metadata" + ) + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*", name): + raise CampaignError( + f"installed distribution name is not replay-safe: {name!r}" + ) + normalized = _normalized_distribution_name(name) + if normalized == project_name: + continue + raw_direct_url = distribution.read_text("direct_url.json") + if raw_direct_url: + try: + direct_url = json.loads(raw_direct_url) + except json.JSONDecodeError as exc: + raise CampaignError( + f"installed dependency {name} has invalid direct_url.json" + ) from exc + origin = str(direct_url.get("url", "unknown")) + raise CampaignError( + f"installed dependency {name}=={version} uses a direct origin " + f"({origin}); a portable exact replay cannot replace it with " + "a name/version guess. Reinstall that dependency from a " + "replayable package index or conda channel before production." + ) + record = {"name": name, "version": version} + previous = packages.get(normalized) + if previous is not None and previous != record: + raise CampaignError( + "conflicting installed distribution metadata for " + f"{normalized}: {previous!r} versus {record!r}" + ) + packages[normalized] = record + for requirement_text in distribution.requires or (): + try: + requirement = Requirement(requirement_text) + except InvalidRequirement as exc: + raise CampaignError( + f"installed dependency {name} has an invalid requirement: " + f"{requirement_text!r}" + ) from exc + dependency_normalized = _normalized_distribution_name( + requirement.name + ) + if dependency_normalized not in visited: + pending.append(requirement.name) + if not packages: + raise CampaignError("could not enumerate the installed dependency closure") + requirements = "\n".join( + f"{record['name']}=={record['version']}" + for _, record in sorted(packages.items()) + ) + "\n" + payload = { + "schema_version": 2, + "python_version": platform.python_version(), + "packages": packages, + "scientific_packages": dict(scientific_packages), + "excluded_project": "rg-nanogpt-one-head", + "replay_policy": ( + "installed campaign dependency closure pinned by metadata " + "name/version; opaque direct origins rejected; checked-out " + "project installed separately with --no-deps" + ), + } + return payload, requirements + + +def _load_yaml(path: Path) -> dict[str, Any]: + try: + import yaml + except ImportError as exc: + raise CampaignError("PyYAML is required to validate the campaign config") from exc + try: + payload = yaml.safe_load(path.read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError) as exc: + raise CampaignError(f"could not read valid YAML from {path}: {exc}") from exc + if not isinstance(payload, dict): + raise CampaignError(f"config is not a YAML mapping: {path}") + return payload + + +def _resolve_config(value: str | Path) -> Path: + path = Path(value) + if not path.is_absolute(): + path = (Path.cwd() / path).resolve(strict=False) + else: + path = path.resolve(strict=False) + if not path.is_file(): + raise CampaignError(f"config does not exist: {path}") + if not _within(path, REPOSITORY_ROOT): + raise CampaignError(f"config must be tracked inside the repository: {path}") + try: + relative = path.relative_to(REPOSITORY_ROOT) + except ValueError as exc: # pragma: no cover - guarded by _within above + raise CampaignError(f"config escapes the repository: {path}") from exc + tracked = _git(("ls-files", "--error-unmatch", "--", str(relative)), check=False) + if tracked.returncode != 0: + raise CampaignError( + "config must be a Git-tracked file (ignored/untracked configs are " + f"not reproducible): {path}" + ) + return path + + +def _validate_protocol_config(config: Path) -> dict[str, Any]: + cfg = _load_yaml(config) + observed_config_sha = _canonical_sha256(cfg) + if observed_config_sha != FROZEN_CONFIG_SHA256: + raise CampaignError( + "config does not exactly match the frozen dated campaign; " + f"canonical_sha256={observed_config_sha}, " + f"expected={FROZEN_CONFIG_SHA256}" + ) + dataset = cfg.get("dataset", {}) + model = cfg.get("model", {}) + training = cfg.get("training", {}) + profiles = cfg.get("optimizer_profiles", {}) + ww = cfg.get("weightwatcher", {}) + runtime = cfg.get("runtime", {}) + expected_dataset = { + "name": "HuggingFaceFW/fineweb-edu", + "config": "sample-10BT", + "split": "train", + "revision": "593b3a867298afb8ce42625a270ef20ddcad28f9", + "tokenizer": "gpt2", + "train_tokens": 80_000_000, + "val_tokens": 1_000_000, + "test_tokens": 1_000_000, + } + if any(dataset.get(key) != value for key, value in expected_dataset.items()): + raise CampaignError("config does not match the frozen FineWeb-Edu corpus protocol") + expected_model = { + "vocab_size": 50_257, + "block_size": 256, + "n_layer": 1, + "n_head": 1, + "n_embd": 128, + "dropout": 0.0, + "bias": False, + "tie_weights": True, + } + if any(model.get(key) != value for key, value in expected_model.items()): + raise CampaignError("config does not match the frozen one-head model") + if tuple(training.get("seeds", ())) != CANONICAL_SEEDS: + raise CampaignError( + f"config training.seeds must be exactly {list(CANONICAL_SEEDS)}" + ) + missing_profiles = set(CANONICAL_OPTIMIZERS).difference(profiles) + if missing_profiles: + raise CampaignError( + "config lacks campaign optimizer profiles: " + + ", ".join(sorted(missing_profiles)) + ) + if ww.get("fix_fingers") != "clip_xmax": + raise CampaignError("weightwatcher.fix_fingers must equal clip_xmax") + if int(ww.get("max_fingers", 0)) < 1: + raise CampaignError("weightwatcher.max_fingers must be positive") + if ww.get("require_raw_alpha") is not True: + raise CampaignError("weightwatcher.require_raw_alpha must be true") + if ww.get("enabled") is not True or ww.get("strict") is not True: + raise CampaignError("WeightWatcher must be enabled and strict") + target_epochs = float(training.get("target_epochs", 0.0)) + epoch_interval = float(training.get("epoch_interval", 0.0)) + if target_epochs != 4.0 or epoch_interval != 0.25: + raise CampaignError("campaign horizon must be 4.0 epochs at 0.25-epoch states") + expected_training = { + "batch_size": 4, + "grad_accum_steps": 8, + "eval_interval_steps": 500, + "eval_batches": 64, + "checkpoint_interval_steps": 500, + "grad_clip": 1.0, + } + if any(training.get(key) != value for key, value in expected_training.items()): + raise CampaignError("config does not match the frozen training/evaluation cadence") + permanent_states = int(round(target_epochs / epoch_interval)) + 1 + if permanent_states < MINIMUM_PERMANENT_STATES: + raise CampaignError( + f"config yields only {permanent_states} permanent states; at least " + f"{MINIMUM_PERMANENT_STATES} are required" + ) + if ( + runtime.get("matmul_precision") != "highest" + or runtime.get("allow_tf32") is not False + or runtime.get("cudnn_benchmark") is not False + or runtime.get("deterministic_algorithms") is not True + or runtime.get("deterministic_warn_only") is not False + ): + raise CampaignError("runtime must use the frozen strict deterministic settings") + return cfg + + +def _expected_total_steps(cfg: Mapping[str, Any]) -> int: + training = cfg["training"] + model = cfg["model"] + dataset = cfg["dataset"] + tokens_per_update = ( + int(training["batch_size"]) + * int(training["grad_accum_steps"]) + * int(model["block_size"]) + ) + return max( + 1, + int( + math.ceil( + float(training["target_epochs"]) + * int(dataset["train_tokens"]) + / tokens_per_update + ) + ), + ) + + +def _expected_permanent_steps(cfg: Mapping[str, Any]) -> tuple[int, ...]: + training = cfg["training"] + train_tokens = int(cfg["dataset"]["train_tokens"]) + tokens_per_update = ( + int(training["batch_size"]) + * int(training["grad_accum_steps"]) + * int(cfg["model"]["block_size"]) + ) + target_epochs = float(training["target_epochs"]) + interval = float(training["epoch_interval"]) + total_steps = _expected_total_steps(cfg) + epochs = [0.0] + current = interval + while current < target_epochs - 1e-12: + epochs.append(round(current, 12)) + current += interval + epochs.append(target_epochs) + steps: dict[int, float] = {} + for epoch in epochs: + if epoch == 0.0: + step = 0 + elif math.isclose(epoch, target_epochs, rel_tol=0.0, abs_tol=1e-12): + step = total_steps + else: + step = int(round(epoch * train_tokens / tokens_per_update)) + steps[min(total_steps, max(0, step))] = epoch + steps[total_steps] = target_epochs + return tuple(sorted(steps)) + + +def _require_dependency_contract() -> dict[str, str]: + versions = _dependency_versions() + missing = [name for name, version in versions.items() if version is None] + if missing: + raise CampaignError( + "missing experiment dependencies: " + ", ".join(missing) + + f". Install with: {sys.executable} -m pip install -e {NANOGPT_ROOT}" + ) + if versions.get("weightwatcher") != PINNED_WEIGHTWATCHER: + raise CampaignError( + "WeightWatcher must be pinned exactly to " + f"{PINNED_WEIGHTWATCHER}; observed {versions.get('weightwatcher')}" + ) + if versions.get("rg-nanogpt-one-head") != PINNED_PACKAGE_VERSION: + raise CampaignError( + "the editable campaign package must be reinstalled at version " + f"{PINNED_PACKAGE_VERSION}; observed " + f"{versions.get('rg-nanogpt-one-head')}. Run: " + f"{sys.executable} -m pip install -e {NANOGPT_ROOT}" + ) + resolved = { + name: str(version) for name, version in versions.items() if version is not None + } + resolved["python"] = platform.python_version() + try: + resolved["torch-xla"] = importlib.metadata.version("torch-xla") + except importlib.metadata.PackageNotFoundError: + resolved["torch-xla"] = "not-installed" + # Enforce the same replayability policy before a multi-day run, rather + # than discovering an opaque direct/VCS/file dependency only at archive. + # Return the full closure in the spelling used by the run manifest so a + # second host cannot begin a disjoint replicate with a silent transitive + # dependency difference. + dependency_lock, _ = _installed_distribution_lock(resolved) + for record in dependency_lock["packages"].values(): + resolved.setdefault(str(record["name"]), str(record["version"])) + return resolved + + +def _verified_data_metadata( + paths: Mapping[str, Path], cfg: Mapping[str, Any] +) -> dict[str, Any]: + data_root = paths["data"] + metadata_path = data_root / "meta.json" + metadata = _read_json(metadata_path) + expected_splits = { + "train": int(cfg["dataset"]["train_tokens"]), + "val": int(cfg["dataset"]["val_tokens"]), + "test": int(cfg["dataset"]["test_tokens"]), + } + expected_metadata = { + "schema_version": 2, + "dataset_name": cfg["dataset"]["name"], + "dataset_config": cfg["dataset"]["config"], + "dataset_split": cfg["dataset"]["split"], + "dataset_revision": cfg["dataset"]["revision"], + "tokenizer": cfg["dataset"]["tokenizer"], + "vocab_size": int(cfg["model"]["vocab_size"]), + "eot_token": 50_256, + "dtype": "uint16", + "document_disjoint_splits": True, + "splits": expected_splits, + } + mismatches = { + key: (metadata.get(key), expected) + for key, expected in expected_metadata.items() + if metadata.get(key) != expected + } + if mismatches: + raise CampaignError( + "prepared corpus metadata does not match the frozen campaign: " + + _canonical_json(mismatches) + ) + file_metadata = metadata.get("files") + if not isinstance(file_metadata, dict): + raise CampaignError("prepared corpus metadata has no file-hash inventory") + for split, token_count in expected_splits.items(): + record = file_metadata.get(split) + if not isinstance(record, dict): + raise CampaignError(f"prepared corpus metadata lacks files.{split}") + if record.get("path") != f"{split}.bin": + raise CampaignError(f"prepared corpus files.{split}.path is not canonical") + path = data_root / f"{split}.bin" + expected_bytes = token_count * 2 + if ( + not path.is_file() + or path.stat().st_size != expected_bytes + or int(record.get("bytes", -1)) != expected_bytes + ): + raise CampaignError(f"prepared {split} token file has an invalid byte count") + observed_hash = _sha256(path) + if not record.get("sha256") or observed_hash != str(record["sha256"]): + raise CampaignError(f"prepared {split} token file SHA-256 mismatch") + return metadata + + +def _inspect_runtime( + config: Path, + device: str, + child_env: Mapping[str, str], +) -> dict[str, Any]: + probe = """ +import json +import sys +from rg_nanogpt_one_head.config import load_config +from rg_nanogpt_one_head.runtime import choose_device, configure_runtime, runtime_metadata +cfg = load_config(sys.argv[1]) +resolved = choose_device(sys.argv[2]) +configure_runtime(resolved, cfg) +print('RG_RUNTIME_JSON=' + json.dumps(runtime_metadata(resolved), sort_keys=True)) +""" + try: + completed = subprocess.run( + [sys.executable, "-c", probe, str(config), str(device)], + cwd=REPOSITORY_ROOT, + env=dict(child_env), + check=True, + capture_output=True, + text=True, + timeout=180, + ) + except (OSError, subprocess.SubprocessError) as exc: + raise CampaignError(f"could not inspect requested runtime: {exc}") from exc + prefix = "RG_RUNTIME_JSON=" + lines = [line for line in completed.stdout.splitlines() if line.startswith(prefix)] + if len(lines) != 1: + raise CampaignError( + "runtime inspection did not emit exactly one metadata record; " + f"stderr={completed.stderr[-2000:]}" + ) + try: + payload = json.loads(lines[0][len(prefix) :]) + except json.JSONDecodeError as exc: + raise CampaignError("runtime inspection emitted invalid JSON") from exc + if not isinstance(payload, dict): + raise CampaignError("runtime inspection metadata is not an object") + return payload + + +def _run_context( + *, + cfg: Mapping[str, Any], + git: Mapping[str, Any], + data_metadata: Mapping[str, Any] | None = None, + dependencies: Mapping[str, str] | None = None, + runtime: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + return { + "cfg": dict(cfg), + "config_sha256": _canonical_sha256(cfg), + "git_commit": str(git.get("commit", "")), + "data_metadata": dict(data_metadata) if data_metadata is not None else None, + "dependencies": dict(dependencies) if dependencies is not None else None, + "runtime": dict(runtime) if runtime is not None else None, + "total_steps": _expected_total_steps(cfg), + "permanent_steps": _expected_permanent_steps(cfg), + } + + +def _parse_optimizers(value: str) -> tuple[str, ...]: + requested = tuple(part.strip() for part in value.split(",") if part.strip()) + if not requested: + raise CampaignError("at least one optimizer is required") + if len(requested) != len(set(requested)): + raise CampaignError("optimizer list contains duplicates") + invalid = set(requested).difference(CANONICAL_OPTIMIZERS) + if invalid: + raise CampaignError("unsupported campaign optimizers: " + ", ".join(sorted(invalid))) + return tuple(name for name in CANONICAL_OPTIMIZERS if name in requested) + + +def _parse_seeds(value: str) -> tuple[int, ...]: + try: + requested = tuple(int(part.strip()) for part in value.split(",") if part.strip()) + except ValueError as exc: + raise CampaignError("seeds must be comma-separated integers") from exc + if not requested: + raise CampaignError("at least one seed is required") + if len(requested) != len(set(requested)): + raise CampaignError("seed list contains duplicates") + invalid = set(requested).difference(CANONICAL_SEEDS) + if invalid: + raise CampaignError("unsupported campaign seeds: " + ", ".join(map(str, sorted(invalid)))) + return tuple(seed for seed in CANONICAL_SEEDS if seed in requested) + + +def _stream_command( + command: Sequence[str], + *, + log_path: Path, + environment: Mapping[str, str], +) -> int: + log_path.parent.mkdir(parents=True, exist_ok=True) + # A stable log path identifies one artifact-producing subprocess. Adjacent + # locks allow different optimizer/seed jobs to run in parallel, while two + # launchers cannot write the same replicate, corpus, report, or preflight. + lock = _acquire_exclusive_lock( + log_path.with_name(log_path.name + ".lock") + ) + try: + started = _utc_now() + header = f"\n[{started}] $ {' '.join(command)}\n" + print(header.rstrip(), flush=True) + with log_path.open("a", encoding="utf-8", buffering=1) as log: + log.write(header) + try: + process = subprocess.Popen( + list(command), + cwd=REPOSITORY_ROOT, + env=dict(environment), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + except OSError as exc: + log.write(f"launcher error: {exc}\n") + raise CampaignError(f"could not start command: {exc}") from exc + assert process.stdout is not None + try: + for line in process.stdout: + print(line, end="", flush=True) + log.write(line) + except KeyboardInterrupt: + process.terminate() + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + raise + return_code = process.wait() + footer = f"[{_utc_now()}] exit_code={return_code}\n" + print(footer.rstrip(), flush=True) + log.write(footer) + return return_code + finally: + _release_exclusive_lock(lock) + + +def _read_json(path: Path) -> dict[str, Any]: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise CampaignError(f"invalid JSON artifact {path}: {exc}") from exc + if not isinstance(payload, dict): + raise CampaignError(f"JSON artifact is not an object: {path}") + return payload + + +def _read_csv_header(path: Path) -> set[str]: + try: + with path.open("r", encoding="utf-8", newline="") as handle: + reader = csv.reader(handle) + return set(next(reader)) + except (OSError, StopIteration, csv.Error) as exc: + raise CampaignError(f"invalid or empty CSV artifact {path}: {exc}") from exc + + +def _run_dir(paths: Mapping[str, Path], optimizer: str, seed: int) -> Path: + return paths["results"] / optimizer / f"seed_{int(seed)}" + + +def _validate_manifest_binding( + manifest: Mapping[str, Any], + completion: Mapping[str, Any], + *, + optimizer: str, + seed: int, + expected_context: Mapping[str, Any] | None, + run_dir: Path, +) -> None: + try: + manifest_seed = int(manifest.get("seed", -1)) + except (TypeError, ValueError) as exc: + raise CampaignError(f"manifest seed is invalid in {run_dir}") from exc + if manifest.get("optimizer") != optimizer or manifest_seed != seed: + raise CampaignError(f"manifest identity mismatch in {run_dir}") + initial_model_hash = str(manifest.get("initial_model_sha256", "")) + if ( + len(initial_model_hash) != 64 + or any( + character not in "0123456789abcdef" + for character in initial_model_hash.lower() + ) + ): + raise CampaignError(f"manifest has no initial-model tensor hash in {run_dir}") + fingerprint = str(manifest.get("protocol_fingerprint", "")) + if not fingerprint or str(completion.get("fingerprint", "")) != fingerprint: + raise CampaignError(f"manifest/completion fingerprint mismatch in {run_dir}") + + source = manifest.get("source_repository") + if not isinstance(source, dict): + raise CampaignError(f"manifest has no source_repository object in {run_dir}") + if source.get("available") is not True or source.get("dirty") is not False: + raise CampaignError(f"run was not produced from a clean readable Git source: {run_dir}") + if not str(source.get("commit", "")) or source.get("commit") == "unknown": + raise CampaignError(f"manifest has no exact source commit in {run_dir}") + + packages = manifest.get("package_versions") + if not isinstance(packages, dict): + raise CampaignError(f"manifest has no package_versions object in {run_dir}") + if str(packages.get("weightwatcher", "")) != PINNED_WEIGHTWATCHER: + raise CampaignError( + f"manifest did not use weightwatcher=={PINNED_WEIGHTWATCHER}: {run_dir}" + ) + if str(packages.get("rg-nanogpt-one-head", "")) != PINNED_PACKAGE_VERSION: + raise CampaignError( + f"manifest did not use rg-nanogpt-one-head=={PINNED_PACKAGE_VERSION}: " + f"{run_dir}" + ) + + runtime = manifest.get("runtime_environment") + if not isinstance(runtime, dict): + raise CampaignError(f"manifest has no runtime_environment object in {run_dir}") + if ( + runtime.get("float32_matmul_precision") != "highest" + or runtime.get("deterministic_algorithms") is not True + or runtime.get("deterministic_warn_only") is not False + ): + raise CampaignError(f"manifest violates deterministic runtime policy in {run_dir}") + if not str(runtime.get("hardware_block_id", "")).strip() or not str( + runtime.get("hardware_block_id_source", "") + ).strip(): + raise CampaignError(f"manifest has no hardware-block identity in {run_dir}") + if runtime.get("accelerator") == "cuda" and ( + runtime.get("cuda_matmul_allow_tf32") is not False + or runtime.get("cudnn_allow_tf32") is not False + ): + raise CampaignError(f"manifest enabled CUDA TF32 in {run_dir}") + + if expected_context is None: + return + cfg = expected_context["cfg"] + expected_mappings = { + "protocol": cfg["protocol"], + "model": cfg["model"], + "training": cfg["training"], + "evaluation": cfg["evaluation"], + "weightwatcher": cfg["weightwatcher"], + } + expected_profile = dict(cfg["optimizer_profiles"][optimizer]) + expected_profile["name"] = optimizer + expected_mappings["optimizer_profile"] = expected_profile + for key, expected in expected_mappings.items(): + if _canonical_json(manifest.get(key)) != _canonical_json(expected): + raise CampaignError(f"manifest {key} differs from frozen config in {run_dir}") + if str(manifest.get("config_sha256", "")) != str( + expected_context["config_sha256"] + ): + raise CampaignError(f"manifest config SHA-256 mismatch in {run_dir}") + if int(manifest.get("max_steps", -1)) != int(expected_context["total_steps"]): + raise CampaignError(f"manifest optimizer-step horizon mismatch in {run_dir}") + tokens_per_step = ( + int(cfg["training"]["batch_size"]) + * int(cfg["training"]["grad_accum_steps"]) + * int(cfg["model"]["block_size"]) + ) + if int(manifest.get("tokens_per_step", -1)) != tokens_per_step: + raise CampaignError(f"manifest tokens_per_step mismatch in {run_dir}") + if str(source.get("commit", "")) != str(expected_context["git_commit"]): + raise CampaignError( + f"run source commit differs from the checked-out commit in {run_dir}" + ) + expected_data = expected_context.get("data_metadata") + if expected_data is not None and _canonical_json( + manifest.get("data_metadata") + ) != _canonical_json(expected_data): + raise CampaignError(f"manifest data metadata/hash inventory mismatch in {run_dir}") + expected_dependencies = expected_context.get("dependencies") + if isinstance(expected_dependencies, Mapping) and _canonical_json( + packages + ) != _canonical_json(expected_dependencies): + expected_names = set(expected_dependencies) + observed_names = set(packages) + mismatched = sorted( + name + for name in expected_names & observed_names + if str(expected_dependencies[name]) != str(packages[name]) + ) + raise CampaignError( + "manifest dependency closure differs from the current production " + f"environment in {run_dir}; missing={sorted(expected_names - observed_names)[:10]}, " + f"extra={sorted(observed_names - expected_names)[:10]}, " + f"version_mismatches={mismatched[:10]}" + ) + expected_runtime = expected_context.get("runtime") + if expected_runtime is not None and _canonical_json( + _runtime_block_identity(runtime) + ) != _canonical_json(_runtime_block_identity(expected_runtime)): + raise CampaignError( + f"manifest runtime/hardware differs from the requested runtime in {run_dir}" + ) + + +def _validate_completed_run( + run_dir: Path, + optimizer: str, + seed: int, + *, + expected_context: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + failure_marker = run_dir / "run_failed.json" + if failure_marker.exists(): + raise CampaignError(f"failure marker exists: {failure_marker}") + missing = [ + relative + for relative in REQUIRED_RUN_FILES + if not (run_dir / relative).is_file() or (run_dir / relative).stat().st_size == 0 + ] + if missing: + raise CampaignError(f"{run_dir} lacks required artifacts: {', '.join(missing)}") + completion = _read_json(run_dir / "run_complete.json") + manifest = _read_json(run_dir / "manifest.json") + if completion.get("completed") is not True: + raise CampaignError(f"{run_dir}/run_complete.json does not declare completed=true") + try: + recorded_seed = int(completion.get("seed", -1)) + except (TypeError, ValueError) as exc: + raise CampaignError(f"completion seed is invalid in {run_dir}") from exc + if completion.get("optimizer") != optimizer or recorded_seed != seed: + raise CampaignError(f"completion identity mismatch for {optimizer}/seed_{seed}") + _validate_manifest_binding( + manifest, + completion, + optimizer=optimizer, + seed=seed, + expected_context=expected_context, + run_dir=run_dir, + ) + + epoch_path = run_dir / "epoch_metrics.csv" + metrics_path = run_dir / "metrics.csv" + spectral_path = run_dir / "spectral" / "layers.csv" + epoch_columns = _read_csv_header(epoch_path) + metric_columns = _read_csv_header(metrics_path) + spectral_columns = _read_csv_header(spectral_path) + required_metrics = { + f"{split}_{metric}" + for split in ("train", "val", "test") + for metric in ( + "loss", + "perplexity", + "bits_per_token", + "accuracy", + "top5_accuracy", + ) + } + required_metrics.update({ + "test_bleu", + "test_continuation_token_accuracy", + "test_continuation_exact_match", + }) + for label, columns in ( + (str(metrics_path), metric_columns), + (str(epoch_path), epoch_columns), + ): + missing_metrics = required_metrics.difference(columns) + if missing_metrics: + raise CampaignError( + f"{label} lacks required campaign metrics: " + + ", ".join(sorted(missing_metrics)) + ) + required_spectral = { + "step", "matrix_name", "alpha", "raw_alpha", "alpha_raw", + "alpha_clip_xmax", "alpha_delta", "num_fingers", "finger_policy", + "primary_alpha_variant", "weightwatcher_analysis_calls", + } + missing_spectral = required_spectral.difference(spectral_columns) + if missing_spectral: + raise CampaignError( + f"{spectral_path} lacks one-pass WeightWatcher columns: " + + ", ".join(sorted(missing_spectral)) + ) + if "checkpoint_path" not in epoch_columns: + raise CampaignError(f"{epoch_path} lacks checkpoint_path") + if "test_held_out" not in epoch_columns: + raise CampaignError(f"{epoch_path} lacks test_held_out policy markers") + + with metrics_path.open("r", encoding="utf-8", newline="") as handle: + metric_rows = list(csv.DictReader(handle)) + if not metric_rows: + raise CampaignError(f"{metrics_path} has no metric rows") + with epoch_path.open("r", encoding="utf-8", newline="") as handle: + epoch_rows = list(csv.DictReader(handle)) + held_out_curve_columns = ( + "test_loss", + "test_perplexity", + "test_bits_per_token", + "test_accuracy", + "test_top5_accuracy", + "test_bleu", + "test_continuation_token_accuracy", + "test_continuation_exact_match", + "test_generalization_gap", + ) + for label, rows in ((metrics_path, metric_rows), (epoch_path, epoch_rows)): + for row in rows: + for column in held_out_curve_columns: + raw = str(row.get(column, "")).strip() + if not raw: + continue + try: + value = float(raw) + except ValueError as exc: + raise CampaignError( + f"invalid held-out placeholder {column} in {label}" + ) from exc + if not math.isnan(value): + raise CampaignError( + f"{label} leaks held-out {column} into a training curve" + ) + if any(str(row.get("test_held_out", "")).strip() not in {"1", "1.0"} for row in epoch_rows): + raise CampaignError(f"{epoch_path} does not mark every test curve held out") + if len(epoch_rows) < MINIMUM_PERMANENT_STATES: + raise CampaignError( + f"{run_dir} has {len(epoch_rows)} permanent states; " + f"at least {MINIMUM_PERMANENT_STATES} are required" + ) + checkpoint_paths: set[Path] = set() + for row in epoch_rows: + recorded = Path(str(row.get("checkpoint_path", ""))) + candidate = recorded if recorded.is_file() else run_dir / "epoch_checkpoints" / recorded.name + if not candidate.is_file() or candidate.stat().st_size == 0: + raise CampaignError(f"missing permanent checkpoint: {recorded}") + checkpoint_paths.add(candidate.resolve()) + if len(checkpoint_paths) < MINIMUM_PERMANENT_STATES: + raise CampaignError(f"{run_dir} has fewer than ten distinct permanent checkpoints") + + with spectral_path.open("r", encoding="utf-8", newline="") as handle: + layer_rows = list(csv.DictReader(handle)) + by_step: dict[int, set[str]] = {} + rows_by_step: dict[int, int] = {} + for row in layer_rows: + try: + step_value = float(row.get("step", "")) + step = int(step_value) + except (TypeError, ValueError) as exc: + raise CampaignError(f"invalid WeightWatcher step in {spectral_path}") from exc + if not math.isfinite(step_value) or step_value != step: + raise CampaignError(f"non-integer WeightWatcher step in {spectral_path}") + by_step.setdefault(step, set()).add(str(row.get("matrix_name", ""))) + rows_by_step[step] = rows_by_step.get(step, 0) + 1 + if row.get("finger_policy") != "fix_fingers=clip_xmax": + raise CampaignError(f"incorrect WeightWatcher finger policy in {spectral_path}") + if row.get("primary_alpha_variant") != "clip_xmax": + raise CampaignError(f"incorrect primary alpha variant in {spectral_path}") + try: + alpha_alias = float(row["alpha"]) + raw_alias = float(row["raw_alpha"]) + alpha = float(row["alpha_clip_xmax"]) + raw_alpha = float(row["alpha_raw"]) + alpha_delta = float(row["alpha_delta"]) + num_fingers = float(row["num_fingers"]) + calls = int(float(row["weightwatcher_analysis_calls"])) + except (KeyError, TypeError, ValueError) as exc: + raise CampaignError(f"invalid WeightWatcher row in {spectral_path}") from exc + values = (alpha_alias, raw_alias, alpha, raw_alpha, alpha_delta, num_fingers) + if not all(math.isfinite(value) for value in values) or calls != 1: + raise CampaignError(f"non-finite alpha or repeated WeightWatcher call in {spectral_path}") + if alpha_alias != alpha or raw_alias != raw_alpha: + raise CampaignError(f"WeightWatcher alpha aliases disagree in {spectral_path}") + if not math.isclose( + alpha_delta, raw_alpha - alpha, rel_tol=1e-12, abs_tol=1e-12 + ) or num_fingers < 0: + raise CampaignError(f"invalid WeightWatcher finger correction in {spectral_path}") + try: + epoch_step_values = [float(row.get("step", "")) for row in epoch_rows] + epoch_steps = {int(value) for value in epoch_step_values} + except (TypeError, ValueError) as exc: + raise CampaignError(f"invalid permanent checkpoint step in {epoch_path}") from exc + if any( + not math.isfinite(value) or value != int(value) + for value in epoch_step_values + ) or len(epoch_steps) != len(epoch_rows): + raise CampaignError(f"duplicate or non-integer permanent steps in {epoch_path}") + if expected_context is not None: + expected_steps = set(int(value) for value in expected_context["permanent_steps"]) + if epoch_steps != expected_steps: + raise CampaignError( + f"permanent checkpoint grid differs from the frozen 17-state grid in {run_dir}" + ) + if set(by_step) != epoch_steps: + raise CampaignError(f"WeightWatcher steps do not match permanent states in {run_dir}") + if any( + len(matrices) != EXPECTED_MATRICES or rows_by_step[step] != EXPECTED_MATRICES + for step, matrices in by_step.items() + ): + raise CampaignError(f"not every WeightWatcher state has six matrices in {run_dir}") + + status_files = list((run_dir / "spectral").glob("status_step_*.json")) + expected_status = { + run_dir / "spectral" / f"status_step_{step:07d}.json" + for step in epoch_steps + } + if set(status_files) != expected_status: + raise CampaignError(f"WeightWatcher completion-record count mismatch in {run_dir}") + for path in status_files: + status = _read_json(path) + try: + calls = int(status.get("weightwatcher_analysis_calls", -1)) + except (TypeError, ValueError) as exc: + raise CampaignError( + f"invalid WeightWatcher call count in status: {path}" + ) from exc + if status.get("completed") is not True or calls != 1: + raise CampaignError(f"incomplete WeightWatcher status: {path}") + + test_results = _read_json(run_dir / "test_results.json") + policy = str(test_results.get("policy", "")).lower() + if "held out" not in policy or "validation" not in policy or "never" not in policy: + raise CampaignError( + f"test_results.json does not declare the protected test policy in {run_dir}" + ) + required_test_fields = { + "step", + "loss", + "perplexity", + "bits_per_token", + "accuracy", + "top5_accuracy", + "bleu", + "continuation_token_accuracy", + "continuation_exact_match", + } + parsed_test_results: dict[str, dict[str, float]] = {} + for checkpoint in ("final", "validation_selected"): + values = test_results.get(checkpoint) + if not isinstance(values, dict): + raise CampaignError( + f"test_results.json lacks object {checkpoint!r} in {run_dir}" + ) + missing_test = required_test_fields.difference(values) + if missing_test: + raise CampaignError( + f"test_results.json {checkpoint} lacks fields: " + + ", ".join(sorted(missing_test)) + ) + try: + numeric_test = [float(values[field]) for field in required_test_fields] + except (TypeError, ValueError) as exc: + raise CampaignError( + f"test_results.json {checkpoint} contains nonnumeric metrics" + ) from exc + if not all(math.isfinite(value) for value in numeric_test): + raise CampaignError( + f"test_results.json {checkpoint} contains non-finite metrics" + ) + parsed = {field: float(values[field]) for field in required_test_fields} + if parsed["loss"] < 0.0 or parsed["bits_per_token"] < 0.0: + raise CampaignError(f"test_results.json {checkpoint} has negative NLL/bits") + if parsed["perplexity"] <= 0.0 or not math.isclose( + math.log(parsed["perplexity"]), + parsed["loss"], + rel_tol=1e-10, + abs_tol=1e-10, + ): + raise CampaignError( + f"test_results.json {checkpoint} perplexity is inconsistent with loss" + ) + if not math.isclose( + parsed["bits_per_token"] * math.log(2.0), + parsed["loss"], + rel_tol=1e-10, + abs_tol=1e-10, + ): + raise CampaignError( + f"test_results.json {checkpoint} bits/token is inconsistent with loss" + ) + bounded = ( + "accuracy", + "top5_accuracy", + "continuation_token_accuracy", + "continuation_exact_match", + ) + if any(not 0.0 <= parsed[field] <= 1.0 for field in bounded): + raise CampaignError( + f"test_results.json {checkpoint} has an accuracy outside [0, 1]" + ) + if ( + parsed["top5_accuracy"] < parsed["accuracy"] + or parsed["continuation_exact_match"] + > parsed["continuation_token_accuracy"] + 1e-12 + or not 0.0 <= parsed["bleu"] <= 100.0 + ): + raise CampaignError( + f"test_results.json {checkpoint} violates metric bounds/order" + ) + parsed_test_results[checkpoint] = parsed + try: + total_steps = int(completion["optimizer_steps"]) + selected_step = int(completion["best_validation_step"]) + final_test_step = int(test_results["final"]["step"]) + selected_test_step = int(test_results["validation_selected"]["step"]) + except (KeyError, TypeError, ValueError) as exc: + raise CampaignError(f"completion/test checkpoint steps are invalid in {run_dir}") from exc + if expected_context is not None and total_steps != int( + expected_context["total_steps"] + ): + raise CampaignError(f"optimizer-step horizon differs from frozen config in {run_dir}") + try: + metric_step_values = [float(row.get("step", "")) for row in metric_rows] + except (TypeError, ValueError) as exc: + raise CampaignError(f"invalid metric step in {metrics_path}") from exc + if ( + any(not math.isfinite(value) or value != int(value) for value in metric_step_values) + or 0 not in metric_step_values + or float(total_steps) not in metric_step_values + or max(metric_step_values) != float(total_steps) + ): + raise CampaignError(f"metrics do not span step zero through {total_steps} in {run_dir}") + if final_test_step != total_steps or selected_test_step != selected_step: + raise CampaignError( + f"final or validation-selected test checkpoint is inconsistent in {run_dir}" + ) + final_completion_fields = { + "loss": "final_test_loss", + "perplexity": "final_test_perplexity", + "bits_per_token": "final_test_bits_per_token", + "accuracy": "final_test_accuracy", + "top5_accuracy": "final_test_top5_accuracy", + "bleu": "final_test_bleu", + "continuation_token_accuracy": "final_test_continuation_token_accuracy", + "continuation_exact_match": "final_test_continuation_exact_match", + } + for metric, completion_key in final_completion_fields.items(): + try: + recorded = float(completion[completion_key]) + except (KeyError, TypeError, ValueError) as exc: + raise CampaignError(f"completion lacks numeric {completion_key} in {run_dir}") from exc + if not math.isfinite(recorded) or not math.isclose( + recorded, + parsed_test_results["final"][metric], + rel_tol=1e-12, + abs_tol=1e-12, + ): + raise CampaignError( + f"completion {completion_key} disagrees with final test results in {run_dir}" + ) + try: + validation_rows = [ + (int(float(row["step"])), float(row["val_loss"])) for row in metric_rows + ] + recorded_best_loss = float(completion["best_validation_loss"]) + except (KeyError, TypeError, ValueError) as exc: + raise CampaignError(f"validation-selection metadata is invalid in {run_dir}") from exc + if not validation_rows or not all(math.isfinite(loss) for _, loss in validation_rows): + raise CampaignError(f"metrics.csv has non-finite validation loss in {run_dir}") + observed_best_step, observed_best_loss = min(validation_rows, key=lambda item: item[1]) + if observed_best_step != selected_step or not math.isclose( + observed_best_loss, + recorded_best_loss, + rel_tol=1e-10, + abs_tol=1e-12, + ): + raise CampaignError( + f"validation-selected checkpoint does not match metrics.csv in {run_dir}" + ) + + # The package validator torch-loads every standard and permanent + # checkpoint, checks its embedded run identity and step/epoch metadata, + # and verifies the exact MuonClip QK interval grid. Keep the launcher + # import-free at startup, but never declare a production run complete from + # filenames and CSVs alone. + try: + from rg_nanogpt_one_head.completion import ( + CompletedRunValidationError, + validate_completed_run, + ) + except ImportError as exc: + raise CampaignError( + "cannot validate checkpoint payloads; install the dated nanoGPT " + "package and its PyTorch dependencies" + ) from exc + try: + validate_completed_run( + run_dir, + expected_fingerprint=str(completion["fingerprint"]), + expected_optimizer=optimizer, + expected_seed=seed, + expected_total_steps=total_steps, + verify_checkpoints=True, + ) + except CompletedRunValidationError as exc: + raise CampaignError(str(exc)) from exc + return completion + + +def _status_rows( + paths: Mapping[str, Path], + optimizers: Sequence[str], + seeds: Sequence[int], + *, + expected_context: Mapping[str, Any] | None = None, +) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for optimizer in optimizers: + for seed in seeds: + run_dir = _run_dir(paths, optimizer, seed) + status = "missing" + detail = "run directory does not exist" + if (run_dir / "run_failed.json").is_file(): + status = "failed" + detail = "run_failed.json exists" + elif run_dir.exists(): + try: + completion = _validate_completed_run( + run_dir, + optimizer, + seed, + expected_context=expected_context, + ) + except CampaignError as exc: + status = "incomplete" + detail = str(exc) + else: + status = "complete" + detail = f"steps={completion.get('optimizer_steps', 'unknown')}" + rows.append({ + "optimizer": optimizer, + "seed": seed, + "status": status, + "detail": detail, + "run_dir": str(run_dir), + }) + return rows + + +def _preflight_existing_campaign( + paths: Mapping[str, Path], + *, + expected_context: Mapping[str, Any], +) -> None: + """Reject a split-provenance campaign before launching another replicate.""" + + ignored_pre_manifest_names = {"muonclip_walk_location.json"} + for optimizer in CANONICAL_OPTIMIZERS: + for seed in CANONICAL_SEEDS: + run_dir = _run_dir(paths, optimizer, seed) + if not run_dir.is_dir(): + continue + substantive = [ + path for path in run_dir.iterdir() + if path.name not in ignored_pre_manifest_names + ] + if not substantive: + continue + manifest_path = run_dir / "manifest.json" + if not manifest_path.is_file(): + raise CampaignError( + "existing campaign artifacts have no manifest and cannot be " + f"safely resumed: {run_dir}" + ) + manifest = _read_json(manifest_path) + completion_path = run_dir / "run_complete.json" + if completion_path.is_file(): + _validate_completed_run( + run_dir, + optimizer, + seed, + expected_context=expected_context, + ) + else: + completion = {"fingerprint": manifest.get("protocol_fingerprint")} + _validate_manifest_binding( + manifest, + completion, + optimizer=optimizer, + seed=seed, + expected_context=expected_context, + run_dir=run_dir, + ) + + +def _write_provenance( + paths: Mapping[str, Path], config: Path, child_env: Mapping[str, str] +) -> dict[str, Any]: + path_names = ( + EXPERIMENT_ROOT_ENV, "RG_NANOGPT_ONE_HEAD_DATA_ROOT", + "RG_NANOGPT_ONE_HEAD_RESULTS_ROOT", "RG_NANOGPT_ONE_HEAD_PLOTS_ROOT", + "HF_HOME", "HF_DATASETS_CACHE", "HUGGINGFACE_HUB_CACHE", + "HF_ASSETS_CACHE", "HF_MODULES_CACHE", "TRANSFORMERS_CACHE", + "TIKTOKEN_CACHE_DIR", "MPLCONFIGDIR", "XDG_CACHE_HOME", + "XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_STATE_HOME", "TORCH_HOME", + "TORCH_EXTENSIONS_DIR", "TORCHINDUCTOR_CACHE_DIR", "CUDA_CACHE_PATH", + "TRITON_CACHE_DIR", "CUPY_CACHE_DIR", "XLA_PERSISTENT_CACHE_PATH", + "PIP_CACHE_DIR", "UV_CACHE_DIR", "NUMBA_CACHE_DIR", + "JOBLIB_TEMP_FOLDER", "KERAS_HOME", "SACREBLEU", + "WANDB_CACHE_DIR", "WANDB_CONFIG_DIR", "WANDB_DATA_DIR", + "PYTHONPYCACHEPREFIX", + "JUPYTER_CONFIG_DIR", "JUPYTER_DATA_DIR", "JUPYTER_RUNTIME_DIR", + "IPYTHONDIR", "TMPDIR", "TMP", "TEMP", + "CONDA_PREFIX", "PYTHONPATH", "HOME", + ) + try: + freeze = subprocess.run( + [sys.executable, "-m", "pip", "freeze", "--all"], + cwd=REPOSITORY_ROOT, + env=dict(child_env), + check=True, + capture_output=True, + text=True, + timeout=120, + ).stdout + except (OSError, subprocess.SubprocessError) as exc: + raise CampaignError(f"could not capture pip freeze: {exc}") from exc + freeze_path = paths["provenance"] / "pip_freeze.txt" + _atomic_text(freeze_path, freeze) + conda_record: dict[str, Any] | None = None + conda_prefix = str(child_env.get("CONDA_PREFIX", "")).strip() + if conda_prefix: + conda_executable = shutil.which("conda", path=child_env.get("PATH")) + if conda_executable is None: + raise CampaignError( + "CONDA_PREFIX is set but the conda executable is unavailable; " + "cannot capture an exact conda replay lock" + ) + try: + explicit = subprocess.run( + [conda_executable, "list", "--explicit", "--prefix", conda_prefix], + cwd=REPOSITORY_ROOT, + env=dict(child_env), + check=True, + capture_output=True, + text=True, + timeout=120, + ).stdout + conda_json_text = subprocess.run( + [conda_executable, "list", "--json", "--prefix", conda_prefix], + cwd=REPOSITORY_ROOT, + env=dict(child_env), + check=True, + capture_output=True, + text=True, + timeout=120, + ).stdout + conda_packages = json.loads(conda_json_text) + except (OSError, subprocess.SubprocessError, json.JSONDecodeError) as exc: + raise CampaignError(f"could not capture exact conda locks: {exc}") from exc + if not explicit.lstrip().startswith("# This file may be used to create"): + raise CampaignError("conda explicit lock has an unexpected format") + if not isinstance(conda_packages, list): + raise CampaignError("conda package inventory is not a JSON list") + explicit_path = paths["provenance"] / "conda_explicit.txt" + packages_path = paths["provenance"] / "conda_packages.json" + _atomic_text(explicit_path, explicit) + _atomic_text( + packages_path, + json.dumps(conda_packages, indent=2, sort_keys=True) + "\n", + ) + conda_record = { + "prefix": conda_prefix, + "explicit_path": str(explicit_path), + "explicit_sha256": _sha256(explicit_path), + "packages_path": str(packages_path), + "packages_sha256": _sha256(packages_path), + } + payload = { + "schema_version": 1, + "captured_at_utc": _utc_now(), + "campaign_id": "nanogpt_one_head_2026_08_21_baseline_v3", + "python": { + "version": sys.version, + "executable": sys.executable, + "implementation": platform.python_implementation(), + }, + "platform": { + "system": platform.system(), + "release": platform.release(), + "machine": platform.machine(), + "platform": platform.platform(), + }, + "git": _git_provenance(), + "config": {"path": str(config), "sha256": _sha256(config)}, + "launcher": {"path": str(SCRIPT_PATH), "sha256": _sha256(SCRIPT_PATH)}, + "dependencies": _dependency_versions(), + "pip_freeze": { + "path": str(freeze_path), + "sha256": _sha256(freeze_path), + "line_count": len(freeze.splitlines()), + }, + "conda": conda_record, + "paths": {name: child_env.get(name) for name in path_names}, + "runtime_environment": { + name: child_env.get(name) + for name in ( + "PYTORCH_ENABLE_MPS_FALLBACK", + "CUBLAS_WORKSPACE_CONFIG", + "MPLBACKEND", + "TOKENIZERS_PARALLELISM", + ) + }, + } + _atomic_json(paths["provenance"] / "provenance.json", payload) + return payload + + +def _command_doctor( + args: argparse.Namespace, + root: Path, + paths: Mapping[str, Path], + child_env: Mapping[str, str], +) -> int: + config = _resolve_config(args.config) + _require_clean_git() + cfg = _validate_protocol_config(config) + versions = _require_dependency_contract() + + smoke_imports = ", ".join(DEPENDENCIES.values()) + import_command = [ + sys.executable, + "-u", + "-c", + f"import {smoke_imports}; import rg_nanogpt_one_head; " + "print('scientific imports: OK')", + ] + log_path = paths["logs"] / "doctor.log" + if _stream_command(import_command, log_path=log_path, environment=child_env) != 0: + raise CampaignError(f"dependency import smoke test failed; inspect {log_path}") + backend_summary_path = paths["provenance"] / "doctor_backend_smoke.json" + backend_summary_path.unlink(missing_ok=True) + backend_command = [ + sys.executable, + "-u", + str( + NANOGPT_ROOT + / "src" + / "rg_nanogpt_one_head" + / "doctor_smoke.py" + ), + "--config", + str(config), + "--work-dir", + str(paths["tmp"] / "doctor-smoke"), + "--summary", + str(backend_summary_path), + "--device", + args.device, + ] + if _stream_command(backend_command, log_path=log_path, environment=child_env) != 0: + raise CampaignError(f"backend smoke test failed; inspect {log_path}") + backend_summary = _read_json(backend_summary_path) + if backend_summary.get("completed") is not True: + raise CampaignError("backend smoke test did not declare completion") + optimizer_smokes = backend_summary.get("optimizers") + if ( + not isinstance(optimizer_smokes, list) + or not all(isinstance(item, dict) for item in optimizer_smokes) + or tuple(item.get("optimizer") for item in optimizer_smokes) + != CANONICAL_OPTIMIZERS + or not all( + item.get("checkpoint_roundtrip") is True + and item.get("resumed_optimizer_step") is True + for item in optimizer_smokes + ) + ): + raise CampaignError("backend smoke optimizer/checkpoint inventory is incomplete") + ww_smoke = backend_summary.get("weightwatcher") + if ( + not isinstance(ww_smoke, dict) + or int(ww_smoke.get("analysis_calls", -1)) != 1 + or int(ww_smoke.get("matrix_count", -1)) != EXPECTED_MATRICES + or int(ww_smoke.get("raw_alpha_count", -1)) != EXPECTED_MATRICES + or int(ww_smoke.get("clipped_alpha_count", -1)) != EXPECTED_MATRICES + ): + raise CampaignError("backend smoke WeightWatcher inventory is incomplete") + + dataset = cfg.get("dataset", {}) + training = cfg.get("training", {}) + effective_tokens = ( + int(training.get("batch_size", 0)) + * int(training.get("grad_accum_steps", 0)) + * int(cfg.get("model", {}).get("block_size", 0)) + ) + doctor = { + "schema_version": 1, + "completed": True, + "checked_at_utc": _utc_now(), + "experiment_root": str(root), + "config": str(config), + "config_sha256": _sha256(config), + "device_request": args.device, + "resolved_device": backend_summary.get("resolved_device"), + "accelerator": backend_summary.get("accelerator"), + "runtime": backend_summary.get("runtime"), + "backend_smoke": { + "summary_path": str(backend_summary_path), + "summary_sha256": _sha256(backend_summary_path), + "optimizer_count": len(optimizer_smokes), + "weightwatcher_analysis_calls": int(ww_smoke["analysis_calls"]), + "weightwatcher_matrix_count": int(ww_smoke["matrix_count"]), + }, + "git": _git_provenance(), + "dependencies": versions, + "campaign": { + "optimizers": list(CANONICAL_OPTIMIZERS), + "seeds": list(CANONICAL_SEEDS), + "replicates": EXPECTED_REPLICATES, + "train_tokens": int(dataset.get("train_tokens", 0)), + "target_epochs": float(training.get("target_epochs", 0.0)), + "effective_tokens_per_update": effective_tokens, + }, + } + _atomic_json(paths["provenance"] / "doctor.json", doctor) + _write_provenance(paths, config, child_env) + print(json.dumps(doctor, indent=2, sort_keys=True)) + return 0 + + +def _require_doctor_gate( + paths: Mapping[str, Path], + *, + config: Path, + git: Mapping[str, Any], + dependencies: Mapping[str, str], + runtime: Mapping[str, Any], +) -> dict[str, Any]: + """Require a successful smoke test bound to this production runtime.""" + + doctor_path = paths["provenance"] / "doctor.json" + backend_path = paths["provenance"] / "doctor_backend_smoke.json" + if not doctor_path.is_file() or not backend_path.is_file(): + raise CampaignError( + "the backend doctor gate has not completed for this campaign root; " + "run `doctor --device ` before `run`" + ) + doctor = _read_json(doctor_path) + if int(doctor.get("schema_version", -1)) != 1 or doctor.get("completed") is not True: + raise CampaignError(f"invalid or incomplete doctor gate: {doctor_path}") + if Path(str(doctor.get("experiment_root", ""))).resolve(strict=False) != paths[ + "root" + ].resolve(strict=False): + raise CampaignError("doctor gate belongs to a different campaign root") + if str(doctor.get("config_sha256", "")) != _sha256(config): + raise CampaignError( + "doctor gate used a different frozen config; rerun `doctor`" + ) + + doctor_git = doctor.get("git") + if not isinstance(doctor_git, Mapping) or ( + doctor_git.get("available") is not True + or doctor_git.get("clean") is not True + or doctor_git.get("dirty") is not False + or str(doctor_git.get("commit", "")) != str(git.get("commit", "")) + or str(doctor_git.get("origin_url", "")) != str(git.get("origin_url", "")) + ): + raise CampaignError( + "doctor gate is not bound to this clean source commit/origin; " + "rerun `doctor`" + ) + if _canonical_json(doctor.get("dependencies")) != _canonical_json(dependencies): + raise CampaignError( + "scientific dependencies changed after the doctor gate; rerun `doctor`" + ) + + doctor_runtime = doctor.get("runtime") + if not isinstance(doctor_runtime, Mapping) or _canonical_json( + _runtime_block_identity(doctor_runtime) + ) != _canonical_json(_runtime_block_identity(runtime)): + raise CampaignError( + "runtime/hardware block differs from the successful doctor gate; " + "rerun `doctor --device `" + ) + if ( + str(doctor.get("resolved_device", "")) != str(doctor_runtime.get("device", "")) + or str(doctor.get("accelerator", "")) + != str(doctor_runtime.get("accelerator", "")) + ): + raise CampaignError("doctor gate has inconsistent resolved-runtime metadata") + + backend_record = doctor.get("backend_smoke") + if not isinstance(backend_record, Mapping): + raise CampaignError("doctor gate has no backend-smoke binding") + recorded_backend_path = Path( + str(backend_record.get("summary_path", "")) + ).resolve(strict=False) + if recorded_backend_path != backend_path.resolve(strict=False): + raise CampaignError("doctor gate points at an unexpected backend summary") + if str(backend_record.get("summary_sha256", "")) != _sha256(backend_path): + raise CampaignError("doctor backend summary changed after the gate completed") + + backend = _read_json(backend_path) + optimizer_smokes = backend.get("optimizers") + weightwatcher = backend.get("weightwatcher") + if ( + int(backend.get("schema_version", -1)) != 1 + or backend.get("completed") is not True + or _canonical_json(backend.get("runtime")) + != _canonical_json(doctor_runtime) + or not isinstance(optimizer_smokes, list) + or tuple( + item.get("optimizer") for item in optimizer_smokes if isinstance(item, dict) + ) + != CANONICAL_OPTIMIZERS + or not all( + isinstance(item, dict) + and item.get("checkpoint_roundtrip") is True + and item.get("resumed_optimizer_step") is True + for item in optimizer_smokes + ) + or not isinstance(weightwatcher, Mapping) + or int(weightwatcher.get("analysis_calls", -1)) != 1 + or int(weightwatcher.get("matrix_count", -1)) != EXPECTED_MATRICES + or int(weightwatcher.get("raw_alpha_count", -1)) != EXPECTED_MATRICES + or int(weightwatcher.get("clipped_alpha_count", -1)) != EXPECTED_MATRICES + ): + raise CampaignError("doctor backend summary no longer satisfies the smoke gate") + return doctor + + +def _command_verify_lock( + args: argparse.Namespace, + root: Path, + paths: Mapping[str, Path], + child_env: Mapping[str, str], +) -> int: + del root, paths, child_env + lock_path = Path(args.lock).resolve(strict=False) + if not lock_path.is_file(): + raise CampaignError(f"dependency lock does not exist: {lock_path}") + lock = _read_json(lock_path) + if int(lock.get("schema_version", -1)) != 2: + raise CampaignError(f"unsupported dependency-lock schema: {lock_path}") + expected_python = str(lock.get("python_version", "")) + if expected_python != platform.python_version(): + raise CampaignError( + "Python version differs from the archived lock: " + f"expected={expected_python!r}, observed={platform.python_version()!r}" + ) + expected_scientific = lock.get("scientific_packages") + expected_packages = lock.get("packages") + if not isinstance(expected_scientific, Mapping) or not isinstance( + expected_packages, Mapping + ): + raise CampaignError(f"dependency lock is incomplete: {lock_path}") + observed_scientific = _require_dependency_contract() + if _canonical_json(observed_scientific) != _canonical_json( + expected_scientific + ): + raise CampaignError( + "scientific dependency map differs from the archived lock" + ) + observed_lock, _ = _installed_distribution_lock(observed_scientific) + observed_packages = observed_lock["packages"] + if _canonical_json(observed_packages) != _canonical_json(expected_packages): + expected_names = set(expected_packages) + observed_names = set(observed_packages) + mismatched = sorted( + name + for name in expected_names & observed_names + if expected_packages[name] != observed_packages[name] + ) + raise CampaignError( + "installed distributions differ from the archived lock; " + f"missing={sorted(expected_names - observed_names)[:10]}, " + f"extra={sorted(observed_names - expected_names)[:10]}, " + f"version_mismatches={mismatched[:10]}" + ) + print(f"Dependency lock verified exactly: {lock_path}") + return 0 + + +def _command_prepare( + args: argparse.Namespace, + root: Path, + paths: Mapping[str, Path], + child_env: Mapping[str, str], +) -> int: + del root + config = _resolve_config(args.config) + _require_clean_git() + cfg = _validate_protocol_config(config) + command = [ + sys.executable, + "-u", + "-m", + "rg_nanogpt_one_head.data", + "--config", + str(config), + "--output-dir", + str(paths["data"]), + ] + if args.force: + command.append("--force") + log_path = paths["logs"] / "prepare.log" + return_code = _stream_command(command, log_path=log_path, environment=child_env) + if return_code != 0: + raise CampaignError( + f"corpus preparation exited with {return_code}; inspect {log_path}" + ) + _verified_data_metadata(paths, cfg) + _write_provenance(paths, config, child_env) + print(f"Prepared and verified corpus: {paths['data']}") + return 0 + + +def _require_prepared_data( + paths: Mapping[str, Path], cfg: Mapping[str, Any] +) -> dict[str, Any]: + required = [paths["data"] / "meta.json"] + [ + paths["data"] / f"{split}.bin" for split in ("train", "val", "test") + ] + missing = [str(path) for path in required if not path.is_file() or path.stat().st_size == 0] + if missing: + raise CampaignError( + "prepared corpus is missing or empty; run `prepare` first: " + + ", ".join(missing) + ) + return _verified_data_metadata(paths, cfg) + + +def _command_run( + args: argparse.Namespace, + root: Path, + paths: Mapping[str, Path], + child_env: Mapping[str, str], +) -> int: + del root + config = _resolve_config(args.config) + git = _require_clean_git() + cfg = _validate_protocol_config(config) + dependencies = _require_dependency_contract() + runtime = _inspect_runtime(config, args.device, child_env) + _require_doctor_gate( + paths, + config=config, + git=git, + dependencies=dependencies, + runtime=runtime, + ) + data_metadata = _require_prepared_data(paths, cfg) + expected_context = _run_context( + cfg=cfg, + git=git, + data_metadata=data_metadata, + dependencies=dependencies, + runtime=runtime, + ) + _preflight_existing_campaign(paths, expected_context=expected_context) + optimizers = _parse_optimizers(args.optimizers) + seeds = _parse_seeds(args.seeds) + failures: list[dict[str, Any]] = [] + history = paths["provenance"] / "command_history.jsonl" + + for optimizer in optimizers: + for seed in seeds: + run_dir = _run_dir(paths, optimizer, seed) + try: + _validate_completed_run( + run_dir, + optimizer, + seed, + expected_context=expected_context, + ) + except CampaignError: + pass + else: + print(f"Already complete and verified: {optimizer}/seed_{seed}") + continue + + module = ( + "rg_nanogpt_one_head.muonclip" + if optimizer == "muon_clip" + else "rg_nanogpt_one_head.training" + ) + command = [ + sys.executable, + "-u", + "-m", + module, + "--config", + str(config), + "--optimizer", + optimizer, + "--seeds", + str(seed), + "--data-root", + str(paths["data"]), + "--results-root", + str(paths["results"]), + "--device", + args.device, + "--mps-retries", + str(args.mps_retries), + "--fail-fast", + ] + log_path = paths["logs"] / "runs" / optimizer / f"seed_{seed}.log" + started = time.monotonic() + _append_jsonl(history, { + "event": "replicate_start", + "at_utc": _utc_now(), + "optimizer": optimizer, + "seed": seed, + "device_request": args.device, + "command": command, + "log": str(log_path), + }) + replicate_env = dict(child_env) + replicate_env["RG_NANOGPT_CAMPAIGN_COMMAND"] = shlex.join(command) + try: + return_code = _stream_command( + command, log_path=log_path, environment=replicate_env + ) + if return_code != 0: + raise CampaignError(f"training process exited with {return_code}") + _validate_completed_run( + run_dir, + optimizer, + seed, + expected_context=expected_context, + ) + except (CampaignError, KeyboardInterrupt) as exc: + failure = { + "optimizer": optimizer, + "seed": seed, + "error": str(exc), + "log": str(log_path), + } + failures.append(failure) + _append_jsonl(history, { + "event": "replicate_end", + "at_utc": _utc_now(), + "optimizer": optimizer, + "seed": seed, + "completed": False, + "elapsed_seconds": time.monotonic() - started, + "error": str(exc), + }) + print( + f"FAILED {optimizer}/seed_{seed}: {exc}; log={log_path}", + file=sys.stderr, + flush=True, + ) + if isinstance(exc, KeyboardInterrupt): + raise + if args.stop_on_error: + break + else: + _append_jsonl(history, { + "event": "replicate_end", + "at_utc": _utc_now(), + "optimizer": optimizer, + "seed": seed, + "completed": True, + "elapsed_seconds": time.monotonic() - started, + }) + print(f"COMPLETE {optimizer}/seed_{seed}: {run_dir}") + if failures and args.stop_on_error: + break + + _write_provenance(paths, config, child_env) + if failures: + _atomic_json(paths["provenance"] / "last_run_failures.json", { + "recorded_at_utc": _utc_now(), "failures": failures + }) + raise CampaignError( + f"{len(failures)} requested replicate(s) failed or remained incomplete" + ) + return 0 + + +def _command_monitor( + args: argparse.Namespace, + root: Path, + paths: Mapping[str, Path], + child_env: Mapping[str, str], +) -> int: + del root + command = [ + sys.executable, + "-u", + "-m", + "rg_nanogpt_one_head.monitor", + "--results-root", + str(paths["results"]), + "--optimizer", + str(args.optimizer), + "--seed", + str(args.seed), + "--interval", + str(args.interval), + "--recent", + str(args.recent), + ] + if args.once: + command.append("--once") + if args.no_clear: + command.append("--no-clear") + try: + completed = subprocess.run( + command, + cwd=REPOSITORY_ROOT, + env=dict(child_env), + check=False, + ) + except OSError as exc: + raise CampaignError(f"could not start live monitor: {exc}") from exc + if completed.returncode != 0: + raise CampaignError(f"live monitor exited with {completed.returncode}") + return 0 + + +def _print_status_table(rows: Sequence[Mapping[str, Any]]) -> None: + widths = { + "optimizer": max(len("optimizer"), *(len(str(row["optimizer"])) for row in rows)), + "seed": max(len("seed"), *(len(str(row["seed"])) for row in rows)), + "status": max(len("status"), *(len(str(row["status"])) for row in rows)), + } + print( + f"{'optimizer':<{widths['optimizer']}} " + f"{'seed':>{widths['seed']}} {'status':<{widths['status']}} detail" + ) + for row in rows: + print( + f"{str(row['optimizer']):<{widths['optimizer']}} " + f"{str(row['seed']):>{widths['seed']}} " + f"{str(row['status']):<{widths['status']}} {row['detail']}" + ) + + +def _command_status( + args: argparse.Namespace, + root: Path, + paths: Mapping[str, Path], + child_env: Mapping[str, str], +) -> int: + del root, child_env + config = _resolve_config(args.config) + cfg = _validate_protocol_config(config) + expected_context = _run_context(cfg=cfg, git=_git_provenance()) + optimizers = _parse_optimizers(args.optimizers) + seeds = _parse_seeds(args.seeds) + rows = _status_rows( + paths, + optimizers, + seeds, + expected_context=expected_context, + ) + if args.json: + print(json.dumps(rows, indent=2, sort_keys=True)) + else: + _print_status_table(rows) + incomplete = [row for row in rows if row["status"] != "complete"] + if incomplete: + print( + f"{len(incomplete)}/{len(rows)} requested replicates are not complete", + file=sys.stderr, + ) + return 1 + return 0 + + +def _require_exact_campaign( + paths: Mapping[str, Path], + *, + expected_context: Mapping[str, Any], +) -> list[dict[str, Any]]: + rows = _status_rows( + paths, + CANONICAL_OPTIMIZERS, + CANONICAL_SEEDS, + expected_context=expected_context, + ) + incomplete = [row for row in rows if row["status"] != "complete"] + if incomplete: + details = "; ".join( + f"{row['optimizer']}/seed_{row['seed']}={row['status']}" + for row in incomplete + ) + raise CampaignError( + f"the exact 2 x 5 campaign is not complete ({details})" + ) + if len(rows) != EXPECTED_REPLICATES: + raise CampaignError( + f"campaign inventory contains {len(rows)} runs, expected {EXPECTED_REPLICATES}" + ) + runtime_reference: str | None = None + initial_hashes: dict[int, set[str]] = { + seed: set() for seed in CANONICAL_SEEDS + } + for optimizer in CANONICAL_OPTIMIZERS: + for seed in CANONICAL_SEEDS: + manifest = _read_json(_run_dir(paths, optimizer, seed) / "manifest.json") + runtime = manifest.get("runtime_environment") + if not isinstance(runtime, Mapping): + raise CampaignError( + f"run has no runtime identity: {optimizer}/seed_{seed}" + ) + serialized = _canonical_json(_runtime_block_identity(runtime)) + initial_hashes[seed].add(str(manifest["initial_model_sha256"])) + if runtime_reference is None: + runtime_reference = serialized + elif serialized != runtime_reference: + raise CampaignError( + "the exact 2 x 5 campaign mixes hardware-block identities; " + "run each accelerator as a separate campaign root" + ) + mismatched_initializations = { + seed: sorted(values) + for seed, values in initial_hashes.items() + if len(values) != 1 + } + if mismatched_initializations: + raise CampaignError( + "optimizer arms do not share identical step-zero tensors by seed: " + + _canonical_json(mismatched_initializations) + ) + return rows + + +def _command_analyze( + args: argparse.Namespace, + root: Path, + paths: Mapping[str, Path], + child_env: Mapping[str, str], +) -> int: + del root + config = _resolve_config(args.config) + git = _require_clean_git() + cfg = _validate_protocol_config(config) + dependencies = _require_dependency_contract() + data_metadata = _require_prepared_data(paths, cfg) + expected_context = _run_context( + cfg=cfg, + git=git, + data_metadata=data_metadata, + dependencies=dependencies, + ) + require_complete = not bool(args.allow_incomplete) + if require_complete: + _require_exact_campaign(paths, expected_context=expected_context) + else: + _preflight_existing_campaign(paths, expected_context=expected_context) + rows = _status_rows( + paths, + CANONICAL_OPTIMIZERS, + CANONICAL_SEEDS, + expected_context=expected_context, + ) + complete = [row for row in rows if row["status"] == "complete"] + completed_optimizers = {str(row["optimizer"]) for row in complete} + missing_optimizers = set(CANONICAL_OPTIMIZERS).difference( + completed_optimizers + ) + if missing_optimizers: + raise CampaignError( + "provisional analysis requires at least one complete replicate " + "from every optimizer; missing: " + + ", ".join(sorted(missing_optimizers)) + ) + complete_seed_sets = { + optimizer: { + int(row["seed"]) + for row in complete + if row["optimizer"] == optimizer + } + for optimizer in CANONICAL_OPTIMIZERS + } + common_seeds = set.intersection(*complete_seed_sets.values()) + if not common_seeds: + raise CampaignError( + "provisional analysis requires at least one paired seed " + "completed across both optimizers" + ) + + report_builder = EXPERIMENT_DIR / "scripts" / "build_report.py" + source_notebook = ( + EXPERIMENT_DIR / "notebooks" / "01_Performance_and_Spectra.ipynb" + ) + if not report_builder.is_file(): + raise CampaignError(f"report builder is missing: {report_builder}") + if not source_notebook.is_file(): + raise CampaignError(f"source analysis notebook is missing: {source_notebook}") + + executed_dir = paths["analysis"] / "notebooks" + executed_dir.mkdir(parents=True, exist_ok=True) + executed_notebook = executed_dir / "01_Performance_and_Spectra.executed.ipynb" + command = [ + sys.executable, + "-u", + "-m", + "papermill", + str(source_notebook), + str(executed_notebook), + "--cwd", + str(REPOSITORY_ROOT), + "-p", + "RESULTS_ROOT", + str(paths["results"]), + "-p", + "OUTPUT_ROOT", + str(paths["analysis"]), + "-p", + "REQUIRE_COMPLETE", + str(require_complete), + ] + log_path = paths["logs"] / "analyze.log" + return_code = _stream_command(command, log_path=log_path, environment=child_env) + if return_code != 0: + raise CampaignError( + f"analysis notebook exited with {return_code}; inspect {log_path}" + ) + required_outputs = ( + paths["analysis"] / "SUMMARY.md", + paths["analysis"] / "report.html", + paths["analysis"] / "results_manifest.json", + paths["analysis"] / "campaign_runs.csv", + paths["analysis"] / "performance_summary.csv", + paths["analysis"] / "paired_seed_differences.csv", + paths["analysis"] / "spectral_layers_all.csv", + paths["analysis"] / "alpha_across_seed_summary.csv", + paths["analysis"] / "saturation_diagnostics.csv", + executed_notebook, + ) + missing = [ + str(path) + for path in required_outputs + if not path.is_file() or path.stat().st_size == 0 + ] + if missing: + raise CampaignError( + "analysis returned zero but required outputs are missing: " + + ", ".join(missing) + ) + for optimizer in CANONICAL_OPTIMIZERS: + for suffix in ( + "performance", + "alpha_raw_vs_clip_xmax", + "erg_gap_num_traps", + ): + figure = paths["plots"] / f"{optimizer}_{suffix}.png" + if not figure.is_file() or figure.stat().st_size == 0: + raise CampaignError(f"analysis did not produce required plot: {figure}") + + results_manifest_path = paths["analysis"] / "results_manifest.json" + results_manifest = _read_json(results_manifest_path) + results_manifest["executed_notebook"] = { + "path": str(executed_notebook.relative_to(paths["analysis"])), + "bytes": executed_notebook.stat().st_size, + "sha256": _sha256(executed_notebook), + } + results_manifest["analysis_finalized_at_utc"] = _utc_now() + _atomic_json(results_manifest_path, results_manifest) + _validate_analysis_bundle( + paths, + expected_context=expected_context, + require_complete=require_complete, + ) + + _write_provenance(paths, config, child_env) + label = ( + "Final analysis report" + if require_complete + else "Provisional analysis report" + ) + print(f"{label}: {paths['analysis'] / 'report.html'}") + print(f"Executed notebook: {executed_notebook}") + return 0 + + +def _validated_artifact_inventory( + records: Any, + *, + root: Path, + label: str, +) -> dict[str, Path]: + if not isinstance(records, list) or not records: + raise CampaignError(f"analysis manifest has no {label} inventory") + inventory: dict[str, Path] = {} + for index, record in enumerate(records): + if not isinstance(record, Mapping): + raise CampaignError(f"analysis manifest {label}[{index}] is not an object") + relative = Path(str(record.get("path", ""))) + if relative.is_absolute() or relative == Path(".") or ".." in relative.parts: + raise CampaignError(f"analysis manifest {label} has unsafe path: {relative}") + unresolved = root / relative + path = unresolved.resolve(strict=False) + if not _within(path, root.resolve(strict=False)): + raise CampaignError(f"analysis manifest {label} escapes its root: {relative}") + key = relative.as_posix() + if key in inventory: + raise CampaignError(f"analysis manifest {label} duplicates {key}") + if not path.is_file() or unresolved.is_symlink(): + raise CampaignError(f"analysis manifest {label} artifact is missing: {path}") + try: + expected_bytes = int(record.get("bytes", -1)) + except (TypeError, ValueError) as exc: + raise CampaignError(f"analysis manifest {label} has invalid size: {key}") from exc + if expected_bytes != path.stat().st_size or str(record.get("sha256", "")) != _sha256(path): + raise CampaignError(f"analysis manifest {label} hash/size is stale: {key}") + inventory[key] = path + return inventory + + +def _expected_analysis_input_paths(paths: Mapping[str, Path]) -> set[str]: + expected: set[str] = set() + results_root = paths["results"].resolve() + for optimizer in CANONICAL_OPTIMIZERS: + for seed in CANONICAL_SEEDS: + run_dir = _run_dir(paths, optimizer, seed) + for relative in ( + "manifest.json", + "run_complete.json", + "metrics.csv", + "epoch_metrics.csv", + "spectral/layers.csv", + "spectral/summary.csv", + "test_results.json", + ): + expected.add(str((run_dir / relative).resolve().relative_to(results_root))) + for step in _expected_permanent_steps(_load_yaml(DEFAULT_CONFIG)): + status = run_dir / "spectral" / f"status_step_{step:07d}.json" + expected.add(str(status.resolve().relative_to(results_root))) + if optimizer == "muon_clip": + expected.add( + str((run_dir / "muonclip_qk.csv").resolve().relative_to(results_root)) + ) + + checkpoint_index = paths["analysis"] / "checkpoint_sha256.csv" + try: + with checkpoint_index.open("r", encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + except OSError as exc: + raise CampaignError(f"could not read checkpoint integrity index: {exc}") from exc + if not rows: + raise CampaignError("checkpoint integrity index is empty") + for row in rows: + relative = Path(str(row.get("checkpoint_relative_path", ""))) + if relative.is_absolute() or relative == Path(".") or ".." in relative.parts: + raise CampaignError(f"checkpoint index has unsafe path: {relative}") + checkpoint = (results_root / relative).resolve(strict=False) + if not _within(checkpoint, results_root): + raise CampaignError(f"checkpoint index escapes results root: {relative}") + if ( + not checkpoint.is_file() + or int(row.get("bytes", -1)) != checkpoint.stat().st_size + or str(row.get("sha256", "")) != _sha256(checkpoint) + ): + raise CampaignError(f"checkpoint index is stale for {relative}") + expected.add(relative.as_posix()) + return expected + + +def _validate_analysis_bundle( + paths: Mapping[str, Path], + *, + expected_context: Mapping[str, Any], + require_complete: bool = True, +) -> dict[str, Any]: + analysis_root = paths["analysis"].resolve() + results_root = paths["results"].resolve() + for relative in REQUIRED_ANALYSIS_FILES: + path = analysis_root / relative + if not path.is_file() or path.stat().st_size == 0: + raise CampaignError(f"analysis bundle lacks required artifact: {path}") + manifest = _read_json(analysis_root / "results_manifest.json") + if int(manifest.get("schema_version", -1)) != 2 or manifest.get("campaign") != ( + "nanogpt_one_head_2026_08_21_baseline" + ): + raise CampaignError("analysis manifest schema/campaign identity is invalid") + exact = manifest.get("exact_campaign") + valid_run_count = int(manifest.get("valid_run_count", -1)) + if not isinstance(exact, Mapping) or ( + exact.get("optimizers") != list(CANONICAL_OPTIMIZERS) + or exact.get("seeds") != list(CANONICAL_SEEDS) + or int(exact.get("expected_run_count", -1)) != EXPECTED_REPLICATES + or bool(exact.get("require_complete")) is not bool(require_complete) + or exact.get("allow_extra_runs") is not False + or exact.get("allow_mixed_runtime") is not False + or valid_run_count < 1 + or valid_run_count > EXPECTED_REPLICATES + or (require_complete and valid_run_count != EXPECTED_REPLICATES) + ): + raise CampaignError( + "analysis manifest does not match the requested 2 x 5 " + "completion policy" + ) + if Path(str(manifest.get("results_root", ""))).resolve(strict=False) != results_root: + raise CampaignError("analysis manifest points at a different results root") + if Path(str(manifest.get("output_root", ""))).resolve(strict=False) != analysis_root: + raise CampaignError("analysis manifest points at a different output root") + if str(manifest.get("source_git_commit", "")) != str(expected_context["git_commit"]): + raise CampaignError("analysis manifest source commit is stale") + frozen = manifest.get("frozen_config") + if not isinstance(frozen, Mapping) or str( + frozen.get("canonical_sha256", "") + ) != str(expected_context["config_sha256"]): + raise CampaignError("analysis manifest frozen config is stale") + builder = manifest.get("report_builder") + report_builder = EXPERIMENT_DIR / "scripts" / "build_report.py" + if not isinstance(builder, Mapping) or str(builder.get("sha256", "")) != _sha256( + report_builder + ): + raise CampaignError("analysis was generated by a different report builder") + + inputs = _validated_artifact_inventory( + manifest.get("input_artifacts"), root=results_root, label="input_artifacts" + ) + if require_complete: + expected_inputs = _expected_analysis_input_paths(paths) + if set(inputs) != expected_inputs: + missing = sorted(expected_inputs.difference(inputs))[:10] + extra = sorted(set(inputs).difference(expected_inputs))[:10] + raise CampaignError( + "analysis input inventory differs from the current campaign; " + f"missing={missing}, extra={extra}" + ) + outputs = _validated_artifact_inventory( + manifest.get("artifacts"), root=analysis_root, label="artifacts" + ) + required_hashed_outputs = set(REQUIRED_ANALYSIS_FILES).difference( + {"results_manifest.json", "notebooks/01_Performance_and_Spectra.executed.ipynb"} + ) + if not required_hashed_outputs.issubset(outputs): + raise CampaignError("analysis output hash inventory is incomplete") + notebook = manifest.get("executed_notebook") + _validated_artifact_inventory( + [notebook] if isinstance(notebook, Mapping) else None, + root=analysis_root, + label="executed_notebook", + ) + return manifest + + +def _copy_small_artifact( + source: Path, + destination: Path, + *, + allowed_source_root: Path, + maximum_bytes: int = 64 * 1024 * 1024, +) -> None: + if source.is_symlink(): + raise CampaignError(f"refusing to copy symlink into run record: {source}") + resolved = source.resolve(strict=True) + if not _within(resolved, allowed_source_root): + raise CampaignError(f"run-record source escapes allowed root: {source}") + if not resolved.is_file(): + raise CampaignError(f"run-record source is not a file: {source}") + if resolved.suffix in {".pt", ".bin"} or ".partial" in resolved.name: + raise CampaignError(f"large/raw payload is prohibited in run records: {source}") + size = resolved.stat().st_size + if size > maximum_bytes: + raise CampaignError( + f"run-record artifact is unexpectedly large ({size} bytes): {source}" + ) + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(resolved, destination) + + +def _command_archive( + args: argparse.Namespace, + root: Path, + paths: Mapping[str, Path], + child_env: Mapping[str, str], +) -> int: + del root + config = _resolve_config(args.config) + git = _require_clean_git() + cfg = _validate_protocol_config(config) + dependencies = _require_dependency_contract() + data_metadata = _require_prepared_data(paths, cfg) + expected_context = _run_context( + cfg=cfg, + git=git, + data_metadata=data_metadata, + dependencies=dependencies, + ) + _require_exact_campaign(paths, expected_context=expected_context) + _validate_analysis_bundle(paths, expected_context=expected_context) + _write_provenance(paths, config, child_env) + _append_jsonl(paths["provenance"] / "command_history.jsonl", { + "event": "run_record_snapshot", + "at_utc": _utc_now(), + "git_commit": git.get("commit"), + }) + + stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + short_commit = str(git.get("commit", "unknown"))[:12] + record_name = f"{stamp}_{short_commit}" + runs_root = EXPERIMENT_DIR / "runs" + if not _within(runs_root, EXPERIMENT_DIR): + raise CampaignError(f"dated runs directory escapes experiment: {runs_root}") + runs_root.mkdir(parents=True, exist_ok=True) + target = runs_root / record_name + if target.exists(): + raise CampaignError(f"refusing to overwrite existing run record: {target}") + stage = Path(tempfile.mkdtemp(prefix=f".{record_name}.partial-", dir=runs_root)) + + try: + protocol_sources = { + EXPERIMENT_DIR / "campaign.yaml": Path("protocol/campaign.yaml"), + config: Path("protocol/baseline.yaml"), + } + for source, relative in protocol_sources.items(): + _copy_small_artifact( + source, + stage / relative, + allowed_source_root=REPOSITORY_ROOT, + ) + + for source in sorted(paths["provenance"].rglob("*")): + if not source.is_file(): + continue + relative = source.resolve().relative_to(paths["provenance"].resolve()) + _copy_small_artifact( + source, + stage / "provenance" / relative, + allowed_source_root=paths["provenance"], + ) + + dependency_lock, replay_requirements = _installed_distribution_lock( + dependencies + ) + _atomic_text( + stage / "provenance" / "requirements_replay.txt", + replay_requirements, + ) + _atomic_json( + stage / "provenance" / "dependency_lock.json", + dependency_lock, + ) + conda_packages_path = stage / "provenance" / "conda_packages.json" + if conda_packages_path.is_file(): + try: + conda_packages = json.loads( + conda_packages_path.read_text(encoding="utf-8") + ) + except (OSError, json.JSONDecodeError) as exc: + raise CampaignError( + f"could not read captured conda package inventory: {exc}" + ) from exc + if not isinstance(conda_packages, list): + raise CampaignError("captured conda package inventory is not a list") + pip_overlay_names = { + _normalized_distribution_name(str(record.get("name", ""))) + for record in conda_packages + if isinstance(record, Mapping) + and ( + str(record.get("channel", "")).lower() == "pypi" + or str(record.get("build", "")).lower() == "pypi_0" + or str(record.get("build_string", "")).lower() == "pypi_0" + ) + } + pip_overlay_names.discard( + _normalized_distribution_name("rg-nanogpt-one-head") + ) + locked_packages = dependency_lock["packages"] + missing_overlay = sorted( + name for name in pip_overlay_names if name not in locked_packages + ) + if missing_overlay: + raise CampaignError( + "conda lists pip overlays absent from Python metadata: " + + ", ".join(missing_overlay[:20]) + ) + overlay_requirements = "\n".join( + f"{locked_packages[name]['name']}=={locked_packages[name]['version']}" + for name in sorted(pip_overlay_names) + ) + if overlay_requirements: + overlay_requirements += "\n" + _atomic_text( + stage / "provenance" / "pip_overlay_replay.txt", + overlay_requirements, + ) + environment_replay_lines = [ + "export REPLAY_CONDA_PREFIX=\"$RG_NANOGPT_EXPERIMENT_ROOT/replay-conda\"", + "conda create --yes --prefix \"$REPLAY_CONDA_PREFIX\" --file \"$RUN_RECORD/provenance/conda_explicit.txt\"", + "export REPLAY_PYTHON=\"$REPLAY_CONDA_PREFIX/bin/python\"", + "\"$REPLAY_PYTHON\" -m pip install -r \"$RUN_RECORD/provenance/pip_overlay_replay.txt\"", + ] + else: + environment_replay_lines = [ + "python -m venv \"$RG_NANOGPT_EXPERIMENT_ROOT/replay-venv\"", + "export REPLAY_PYTHON=\"$RG_NANOGPT_EXPERIMENT_ROOT/replay-venv/bin/python\"", + "\"$REPLAY_PYTHON\" -m pip install -r \"$RUN_RECORD/provenance/requirements_replay.txt\"", + ] + + for source in sorted(paths["analysis"].rglob("*")): + if not source.is_file(): + continue + relative = source.resolve().relative_to(paths["analysis"].resolve()) + _copy_small_artifact( + source, + stage / "analysis" / relative, + allowed_source_root=paths["analysis"], + ) + + run_manifest_rows: list[dict[str, Any]] = [] + for optimizer in CANONICAL_OPTIMIZERS: + for seed in CANONICAL_SEEDS: + run_dir = _run_dir(paths, optimizer, seed) + manifest = _read_json(run_dir / "manifest.json") + completion = _read_json(run_dir / "run_complete.json") + run_manifest_rows.append({ + "optimizer": optimizer, + "seed": seed, + "accelerator": manifest.get("runtime_environment", {}).get( + "accelerator" + ) if isinstance(manifest.get("runtime_environment"), dict) else None, + "hardware_block_id": manifest.get( + "runtime_environment", {} + ).get("hardware_block_id") + if isinstance(manifest.get("runtime_environment"), dict) + else None, + "hardware_block_id_source": manifest.get( + "runtime_environment", {} + ).get("hardware_block_id_source") + if isinstance(manifest.get("runtime_environment"), dict) + else None, + "torch_version": manifest.get( + "torch_version", + manifest.get("runtime_environment", {}).get("torch_version") + if isinstance(manifest.get("runtime_environment"), dict) + else None, + ), + "optimizer_steps": completion.get("optimizer_steps"), + "best_validation_step": completion.get("best_validation_step"), + }) + for filename in ( + "manifest.json", + "run_complete.json", + "test_results.json", + ): + _copy_small_artifact( + run_dir / filename, + stage / "runs" / optimizer / f"seed_{seed}" / filename, + allowed_source_root=run_dir, + ) + + accelerators = {str(row["accelerator"]) for row in run_manifest_rows} + if len(accelerators) != 1: + raise CampaignError( + "archived campaign does not have one accelerator identity" + ) + replay_device = next(iter(accelerators)) + if replay_device not in {"cpu", "mps", "cuda", "tpu"}: + raise CampaignError( + f"unsupported archived accelerator identity: {replay_device!r}" + ) + hardware_blocks = { + ( + str(row["hardware_block_id"]), + str(row["hardware_block_id_source"]), + ) + for row in run_manifest_rows + } + if len(hardware_blocks) != 1: + raise CampaignError( + "archived campaign does not have one hardware-block identity" + ) + replay_hardware_block, replay_hardware_source = next( + iter(hardware_blocks) + ) + hardware_replay_lines = ( + [ + "export RG_NANOGPT_HARDWARE_BLOCK_ID=" + + shlex.quote(replay_hardware_block) + ] + if replay_hardware_source == "user" + else [] + ) + tpu_replay_lines = ( + ["export RG_NANOGPT_ALLOW_EPHEMERAL_TPU_STORAGE=1"] + if replay_device == "tpu" + else [] + ) + tag = str(git.get("tag_status", "untagged")) + record_lines = [ + "# One-head nanoGPT campaign run record", + "", + f"- Archived at (UTC): `{_utc_now()}`", + f"- Source commit: `{git.get('commit', 'unknown')}`", + f"- Git describe: `{git.get('describe', 'unknown')}`", + f"- Tag status: `{tag}`", + f"- Frozen config canonical SHA-256: `{FROZEN_CONFIG_SHA256}`", + f"- Frozen config file SHA-256: `{_sha256(config)}`", + "- Design: AdamW / MuonClip × seeds " + "1337 / 2027 / 4099 / 31415 / 271828", + "- Test policy: held out until post-training; validation NLL selects checkpoints", + "- WeightWatcher: one clip_xmax call per state; clipped and raw alpha retained", + "", + "## Reproduce", + "", + "Start from the repository commit/tag that contains this archive. Copy the", + "record outside Git before checking out the exact source commit used to run it.", + "", + "```bash", + "export RG_NANOGPT_EXPERIMENT_ROOT=/tmp/rg-nanogpt-one-head-20260821", + "mkdir -p \"$RG_NANOGPT_EXPERIMENT_ROOT/cache/home\"", + "export HOME=\"$RG_NANOGPT_EXPERIMENT_ROOT/cache/home\"", + f"export ARCHIVED_RUN_RECORD='baseline/experiments/nanogpt_one_head_2026_08_21_baseline/runs/{record_name}'", + "export RUN_RECORD=\"$RG_NANOGPT_EXPERIMENT_ROOT/archived-run-record\"", + "mkdir -p \"$RUN_RECORD\"", + "cp -R \"$ARCHIVED_RUN_RECORD\"/. \"$RUN_RECORD\"/", + f"git checkout {git.get('commit', 'COMMIT')}", + "mkdir -p \"$RG_NANOGPT_EXPERIMENT_ROOT\"/{cache/{home,pip,xdg/{cache,config,data,state},matplotlib},tmp}", + "export PIP_CACHE_DIR=\"$RG_NANOGPT_EXPERIMENT_ROOT/cache/pip\"", + "export XDG_CACHE_HOME=\"$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/cache\"", + "export XDG_CONFIG_HOME=\"$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/config\"", + "export XDG_DATA_HOME=\"$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/data\"", + "export XDG_STATE_HOME=\"$RG_NANOGPT_EXPERIMENT_ROOT/cache/xdg/state\"", + "export MPLCONFIGDIR=\"$RG_NANOGPT_EXPERIMENT_ROOT/cache/matplotlib\"", + "export TMPDIR=\"$RG_NANOGPT_EXPERIMENT_ROOT/tmp\"", + "export PYTORCH_ENABLE_MPS_FALLBACK=1", + *tpu_replay_lines, + *hardware_replay_lines, + *environment_replay_lines, + "\"$REPLAY_PYTHON\" -m pip install --no-deps -e './baseline/nanogpt_one_head'", + "\"$REPLAY_PYTHON\" baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/run_experiment.py verify-lock --lock \"$RUN_RECORD/provenance/dependency_lock.json\"", + f"\"$REPLAY_PYTHON\" baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/run_experiment.py doctor --device {replay_device}", + "\"$REPLAY_PYTHON\" baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/run_experiment.py prepare", + f"\"$REPLAY_PYTHON\" baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/run_experiment.py run --device {replay_device}", + "\"$REPLAY_PYTHON\" baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/run_experiment.py analyze", + "\"$REPLAY_PYTHON\" baseline/experiments/nanogpt_one_head_2026_08_21_baseline/scripts/run_experiment.py archive", + "```", + "", + "The `analysis/` directory contains aggregate tables, separate optimizer plots,", + "the HTML/Markdown report, and the executed notebook. `runs/` contains only", + "lightweight per-replicate manifests and protected test outcomes. Corpus files,", + "model checkpoints, caches, and raw large data are intentionally excluded.", + "The raw `pip_freeze.txt` is retained for audit. `requirements_replay.txt`", + "pins the complete installed campaign dependency closure by metadata version.", + "Opaque direct/VCS/file origins are rejected instead of being rewritten as a", + "misleading name/version pin; the project is installed separately with", + "`--no-deps -e`. A conda run also includes an explicit platform lock and a", + "pip-overlay lock. Large binary packages are not vendored into Git; preserve", + "the public package-channel configuration or an external wheelhouse needed to", + "resolve them. Replay fails closed unless `verify-lock` passes before training.", + "", + "## Replicate inventory", + "", + "| optimizer | seed | accelerator | torch | steps | best validation step |", + "|---|---:|---|---|---:|---:|", + ] + for row in run_manifest_rows: + record_lines.append( + "| {optimizer} | {seed} | {accelerator} | {torch_version} | " + "{optimizer_steps} | {best_validation_step} |".format(**row) + ) + _atomic_text(stage / "RUN_RECORD.md", "\n".join(record_lines) + "\n") + + artifact_rows = [] + total_bytes = 0 + for artifact in sorted(stage.rglob("*")): + if not artifact.is_file(): + continue + size = artifact.stat().st_size + total_bytes += size + artifact_rows.append({ + "path": str(artifact.relative_to(stage)), + "bytes": size, + "sha256": _sha256(artifact), + }) + maximum_total = 256 * 1024 * 1024 + if total_bytes > maximum_total: + raise CampaignError( + f"tracked run record would be {total_bytes} bytes; limit is {maximum_total}" + ) + _atomic_json(stage / "archive_manifest.json", { + "schema_version": 1, + "created_at_utc": _utc_now(), + "source_git": git, + "config_sha256": _sha256(config), + "config_canonical_sha256": FROZEN_CONFIG_SHA256, + "file_count_excluding_manifest": len(artifact_rows), + "total_bytes_excluding_manifest": total_bytes, + "files": artifact_rows, + "excluded": [ + "tokenized corpus", "model checkpoints", "cache directories", "training logs" + ], + }) + os.replace(stage, target) + except BaseException: + if stage.exists(): + shutil.rmtree(stage) + raise + + print(f"Tracked-ready run record: {target}") + print("This command intentionally made the Git worktree dirty; review and commit the run record.") + return 0 + + +def _add_config_argument(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--config", + default=str(DEFAULT_CONFIG), + help=f"frozen protocol config (default: {DEFAULT_CONFIG})", + ) + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Run, audit, analyze, and archive the exact dated one-head " + "nanoGPT AdamW/MuonClip campaign" + ) + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + doctor = subparsers.add_parser("doctor", help="validate source, dependencies, device, and paths") + _add_config_argument(doctor) + doctor.add_argument( + "--device", choices=("auto", "cpu", "cuda", "mps", "tpu", "xla"), default="auto" + ) + doctor.set_defaults(handler=_command_doctor) + + verify_lock = subparsers.add_parser( + "verify-lock", + help="compare the active environment with an archived exact dependency lock", + ) + verify_lock.add_argument( + "--lock", + required=True, + help="path to an archived provenance/dependency_lock.json", + ) + verify_lock.set_defaults(handler=_command_verify_lock) + + prepare = subparsers.add_parser("prepare", help="download and verify the pinned token corpus") + _add_config_argument(prepare) + prepare.add_argument("--force", action="store_true", help="rebuild even if verified data exist") + prepare.set_defaults(handler=_command_prepare) + + run = subparsers.add_parser("run", help="run or resume requested campaign replicates") + _add_config_argument(run) + run.add_argument( + "--optimizers", + default=",".join(CANONICAL_OPTIMIZERS), + help="canonical comma-separated subset; default is both arms", + ) + run.add_argument( + "--seeds", + default=",".join(map(str, CANONICAL_SEEDS)), + help="canonical comma-separated subset; default is all five", + ) + run.add_argument( + "--device", choices=("auto", "cpu", "cuda", "mps", "tpu", "xla"), default="auto" + ) + run.add_argument( + "--mps-retries", + type=int, + default=2, + help="fresh-process MPS resume attempts after failure (default: 2)", + ) + run.add_argument("--stop-on-error", action="store_true") + run.set_defaults(handler=_command_run) + + monitor = subparsers.add_parser( + "monitor", + help="live-display training plus raw and clip_xmax WeightWatcher alphas", + ) + monitor.add_argument("--optimizer", choices=CANONICAL_OPTIMIZERS, default="muon_clip") + monitor.add_argument("--seed", type=int, choices=CANONICAL_SEEDS, default=1337) + monitor.add_argument("--interval", type=float, default=30.0) + monitor.add_argument("--recent", type=int, default=8) + monitor.add_argument("--once", action="store_true") + monitor.add_argument("--no-clear", action="store_true") + monitor.set_defaults(handler=_command_monitor) + + status = subparsers.add_parser("status", help="strictly inspect requested replicate artifacts") + _add_config_argument(status) + status.add_argument("--optimizers", default=",".join(CANONICAL_OPTIMIZERS)) + status.add_argument("--seeds", default=",".join(map(str, CANONICAL_SEEDS))) + status.add_argument("--json", action="store_true", help="emit machine-readable status") + status.set_defaults(handler=_command_status) + + analyze = subparsers.add_parser( + "analyze", + help="execute the report notebook (strict 2 x 5 by default)", + ) + _add_config_argument(analyze) + analyze.add_argument( + "--allow-incomplete", + action="store_true", + help=( + "build a clearly marked provisional report from completed runs; " + "archive still requires the exact 2 x 5 campaign" + ), + ) + analyze.set_defaults(handler=_command_analyze) + + archive = subparsers.add_parser( + "archive", + help="copy a small, check-in-ready run record into the dated experiment", + ) + _add_config_argument(archive) + archive.set_defaults(handler=_command_archive) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + if getattr(args, "mps_retries", 0) < 0: + parser.error("--mps-retries must be nonnegative") + if getattr(args, "interval", 1.0) <= 0: + parser.error("--interval must be positive") + if getattr(args, "recent", 1) < 1: + parser.error("--recent must be positive") + + try: + root = resolve_experiment_root() + except CampaignError as exc: + print(f"[campaign] ERROR: {exc}", file=sys.stderr, flush=True) + return 2 + paths = _paths(root) + _create_runtime_directories(paths) + child_env = _child_environment(root, paths) + history = paths["provenance"] / "command_history.jsonl" + invocation_id = hashlib.sha256( + f"{os.getpid()}:{time.time_ns()}:{' '.join(sys.argv)}".encode("utf-8") + ).hexdigest()[:16] + _append_jsonl(history, { + "event": "command_start", + "invocation_id": invocation_id, + "at_utc": _utc_now(), + "argv": list(sys.argv if argv is None else [SCRIPT_PATH.name, *argv]), + "cwd": str(Path.cwd()), + "pid": os.getpid(), + "experiment_root": str(root), + "git": _git_provenance(), + }) + exit_code = 1 + error: str | None = None + try: + exit_code = int(args.handler(args, root, paths, child_env)) + except CampaignError as exc: + exit_code = 2 + error = str(exc) + print(f"[campaign] ERROR: {exc}", file=sys.stderr, flush=True) + except KeyboardInterrupt: + exit_code = 130 + error = "interrupted" + print("[campaign] interrupted", file=sys.stderr, flush=True) + finally: + _append_jsonl(history, { + "event": "command_end", + "invocation_id": invocation_id, + "at_utc": _utc_now(), + "exit_code": exit_code, + "error": error, + }) + return exit_code + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/baseline/nanogpt_one_head/README.md b/baseline/nanogpt_one_head/README.md index c6e31d2f..eb760141 100644 --- a/baseline/nanogpt_one_head/README.md +++ b/baseline/nanogpt_one_head/README.md @@ -6,6 +6,13 @@ measurement, and multi-seed conventions in `CalculatedContent/nanogpt-experiments`, but is isolated here so RG optimizer variants can use it as a clean control. +The current AdamW/MuonClip, five-seed, clip-Xmax/raw-alpha campaign is +defined in the +[2026-08-21 dated experiment](../experiments/nanogpt_one_head_2026_08_21_baseline/README.md). +Use that folder's checked runner for new overnight Mac, H100, or single-device +TPU executions. The protocol below remains the historical one-epoch +SGD/AdamW/Muon reference. + It trains the same **one-block, one-attention-head nanoGPT** with: 1. **SGD + Nesterov momentum**; @@ -166,8 +173,11 @@ Each nominal reporting checkpoint records: ```text train / validation / test cross-entropy train / validation / test perplexity +train / validation / test bits per token train / validation / test next-token top-1 accuracy +train / validation / test next-token top-5 accuracy fixed-continuation test BLEU +fixed-continuation token accuracy and exact match validation and test generalization gaps primary and auxiliary learning rates gradient norms @@ -194,6 +204,7 @@ results//seed_/ manifest.json metrics.csv epoch_metrics.csv + checkpoint_initial.pt checkpoint_latest.pt checkpoint_best.pt checkpoint_final.pt @@ -211,11 +222,15 @@ generator, Python/NumPy/Torch RNG state, CUDA RNG state where applicable, MPS RNG state where the installed PyTorch exposes it, elapsed time, validation-best state, and a protocol fingerprint. A mismatched config, verified data identity, optimizer, or seed is rejected rather than silently resumed. Completed runs are -skipped. - -Test measurements are monitoring-only. Validation loss selects -`checkpoint_best.pt`; test loss, test accuracy, test perplexity, and BLEU never -change optimizer updates, schedules, stopping, or checkpoint selection. +skipped. Full-state checkpoints carry exact model and optimizer-state SHA-256 +digests; permanent model-only checkpoints carry model-state digests. Completion +validation recomputes them before reuse or analysis. + +The test split stays held out throughout training. Validation loss selects +`checkpoint_best.pt`; only after optimization ends are the final and +validation-selected checkpoints evaluated for test loss, accuracy, perplexity, +and BLEU. Test outcomes never change updates, schedules, stopping, or +checkpoint selection. ## Conda / local workflow diff --git a/baseline/nanogpt_one_head/TPU.md b/baseline/nanogpt_one_head/TPU.md index 405e06f5..8e86a0e3 100644 --- a/baseline/nanogpt_one_head/TPU.md +++ b/baseline/nanogpt_one_head/TPU.md @@ -162,7 +162,8 @@ checkpoint selection policy, or WeightWatcher definition. on a different accelerator. - WeightWatcher always receives CPU copies of only the six hidden transformer matrices. -- Greedy BLEU is monitoring-only. On TPU it runs on a CPU snapshot to avoid - compiling a different XLA graph for every generated sequence length. +- Greedy BLEU is a post-training secondary audit. On TPU it runs on a CPU + snapshot to avoid compiling a different XLA graph for every generated + sequence length. - Randomized WeightWatcher measurements restore Python, NumPy, CPU Torch, and accelerator RNG state before training resumes. diff --git a/baseline/nanogpt_one_head/notebooks/01_sgd_momentum_baseline.ipynb b/baseline/nanogpt_one_head/notebooks/01_sgd_momentum_baseline.ipynb index 5f528b25..4e81bd30 100644 --- a/baseline/nanogpt_one_head/notebooks/01_sgd_momentum_baseline.ipynb +++ b/baseline/nanogpt_one_head/notebooks/01_sgd_momentum_baseline.ipynb @@ -117,7 +117,7 @@ "source": [ "## Per-epoch task metrics\n", "\n", - "Test measurements are monitoring-only and never select checkpoints, change learning rates, or tune hyperparameters. `test_bleu` is deterministic greedy continuation overlap against fixed held-out continuations; it is a secondary language-model diagnostic, not a translation benchmark." + "The test split stays held out during optimization. After training, the final and validation-selected checkpoints receive the test and deterministic greedy-continuation BLEU audits; these never select checkpoints, change learning rates, or tune hyperparameters. BLEU is a secondary language-model diagnostic, not a translation benchmark." ] }, { diff --git a/baseline/nanogpt_one_head/notebooks/02_adamw_baseline.ipynb b/baseline/nanogpt_one_head/notebooks/02_adamw_baseline.ipynb index b33084cc..da69146a 100644 --- a/baseline/nanogpt_one_head/notebooks/02_adamw_baseline.ipynb +++ b/baseline/nanogpt_one_head/notebooks/02_adamw_baseline.ipynb @@ -117,7 +117,7 @@ "source": [ "## Per-epoch task metrics\n", "\n", - "Test measurements are monitoring-only and never select checkpoints, change learning rates, or tune hyperparameters. `test_bleu` is deterministic greedy continuation overlap against fixed held-out continuations; it is a secondary language-model diagnostic, not a translation benchmark." + "The test split stays held out during optimization. After training, the final and validation-selected checkpoints receive the test and deterministic greedy-continuation BLEU audits; these never select checkpoints, change learning rates, or tune hyperparameters. BLEU is a secondary language-model diagnostic, not a translation benchmark." ] }, { diff --git a/baseline/nanogpt_one_head/notebooks/03_muon_baseline.ipynb b/baseline/nanogpt_one_head/notebooks/03_muon_baseline.ipynb index a1f3a29d..469c8c33 100644 --- a/baseline/nanogpt_one_head/notebooks/03_muon_baseline.ipynb +++ b/baseline/nanogpt_one_head/notebooks/03_muon_baseline.ipynb @@ -117,7 +117,7 @@ "source": [ "## Per-epoch task metrics\n", "\n", - "Test measurements are monitoring-only and never select checkpoints, change learning rates, or tune hyperparameters. `test_bleu` is deterministic greedy continuation overlap against fixed held-out continuations; it is a secondary language-model diagnostic, not a translation benchmark." + "The test split stays held out during optimization. After training, the final and validation-selected checkpoints receive the test and deterministic greedy-continuation BLEU audits; these never select checkpoints, change learning rates, or tune hyperparameters. BLEU is a secondary language-model diagnostic, not a translation benchmark." ] }, { diff --git a/baseline/nanogpt_one_head/notebooks/04_compare_baselines.ipynb b/baseline/nanogpt_one_head/notebooks/04_compare_baselines.ipynb index df1ef905..5ab3d544 100644 --- a/baseline/nanogpt_one_head/notebooks/04_compare_baselines.ipynb +++ b/baseline/nanogpt_one_head/notebooks/04_compare_baselines.ipynb @@ -70,7 +70,7 @@ "source": [ "## Load the complete nine-run suite\n", "\n", - "This fails if any optimizer/seed run is incomplete. The test set is monitoring-only and was never used for optimizer updates, schedule selection, early stopping, or hyperparameter tuning." + "This fails if any optimizer/seed run is incomplete. The test set stayed held out until post-training audits of the final and validation-selected checkpoints; it was never used for updates, schedule selection, early stopping, or hyperparameter tuning." ] }, { diff --git a/baseline/nanogpt_one_head/notebooks/05_muonclip_esd_clip_xmax.ipynb b/baseline/nanogpt_one_head/notebooks/05_muonclip_esd_clip_xmax.ipynb index 2331b1ac..c4e17178 100644 --- a/baseline/nanogpt_one_head/notebooks/05_muonclip_esd_clip_xmax.ipynb +++ b/baseline/nanogpt_one_head/notebooks/05_muonclip_esd_clip_xmax.ipynb @@ -18,7 +18,7 @@ "jupyter lab notebooks/05_muonclip_esd_clip_xmax.ipynb\n", "```\n", "\n", - "The notebook never hard-codes a run path. It loads the latest completed epoch/quarter-epoch checkpoint by default, reruns standard WeightWatcher, reruns `fix_fingers='clip_xmax'`, compares the fits, and plots a 2x3 ESD grid.\n" + "The notebook never hard-codes a run path. It loads the latest completed epoch/quarter-epoch checkpoint by default, runs WeightWatcher exactly once with `fix_fingers='clip_xmax'`, compares `raw_alpha` with the corrected `alpha`, and plots a 2x3 ESD grid.\n" ] }, { @@ -148,9 +148,9 @@ "id": "standard-note", "metadata": {}, "source": [ - "## Standard WeightWatcher analysis\n", + "## One-pass WeightWatcher analysis\n", "\n", - "This reproduces the training-time contract and uses the same six matrix holder.\n" + "This reproduces the training-time contract and uses the same six-matrix holder. WeightWatcher 0.7.7 returns both the corrected `alpha` and pre-finger-removal `raw_alpha` from this single call.\n" ] }, { @@ -179,16 +179,22 @@ " display(frame[cols].sort_values(\"matrix_name\"))\n", "\n", "reset_diagnostic_seed()\n", - "watcher_standard = ww.WeightWatcher(model=holder)\n", - "details_standard_raw = watcher_standard.analyze(\n", + "watcher = ww.WeightWatcher(model=holder)\n", + "details_raw = watcher.analyze(\n", " ERG=ERG,\n", " randomize=RANDOMIZE,\n", " plot=True,\n", " min_evals=MIN_EVALS,\n", + " fix_fingers='clip_xmax',\n", + " max_fingers=MAX_FINGERS,\n", ")\n", - "details_standard = attach(details_standard_raw)\n", - "details_standard.to_csv(OUTPUT_DIR / \"weightwatcher_standard.csv\", index=False)\n", - "show_details(details_standard)" + "details = attach(details_raw)\n", + "details[\"alpha_raw\"] = pd.to_numeric(details[\"raw_alpha\"], errors=\"coerce\")\n", + "details[\"alpha_clip_xmax\"] = pd.to_numeric(details[\"alpha\"], errors=\"coerce\")\n", + "details[\"alpha_reduction\"] = details[\"alpha_raw\"] - details[\"alpha_clip_xmax\"]\n", + "details[\"weightwatcher_analysis_calls\"] = 1\n", + "details.to_csv(OUTPUT_DIR / \"weightwatcher_one_pass_clip_xmax.csv\", index=False)\n", + "show_details(details)" ] }, { @@ -198,7 +204,7 @@ "source": [ "## `fix_fingers='clip_xmax'`\n", "\n", - "Same checkpoint, same diagnostic seed, same ERG/randomization settings; only the finite-size finger correction is changed.\n" + "No second analysis is needed: `raw_alpha` and corrected `alpha` come from the one call above. The separate raw fit boundaries are not claimed because this one-pass contract does not return them.\n" ] }, { @@ -208,19 +214,13 @@ "metadata": {}, "outputs": [], "source": [ - "reset_diagnostic_seed()\n", - "watcher_clip = ww.WeightWatcher(model=holder)\n", - "details_clip_raw = watcher_clip.analyze(\n", - " ERG=ERG,\n", - " randomize=RANDOMIZE,\n", - " plot=True,\n", - " min_evals=MIN_EVALS,\n", - " fix_fingers='clip_xmax',\n", - " max_fingers=MAX_FINGERS,\n", - ")\n", - "details_clip = attach(details_clip_raw)\n", - "details_clip.to_csv(OUTPUT_DIR / \"weightwatcher_clip_xmax.csv\", index=False)\n", - "show_details(details_clip)" + "required = {\"alpha\", \"raw_alpha\", \"num_fingers\"}\n", + "missing = required.difference(details.columns)\n", + "if missing:\n", + " raise RuntimeError(f\"one-pass WeightWatcher result is missing: {sorted(missing)}\")\n", + "if not details[\"weightwatcher_analysis_calls\"].eq(1).all():\n", + " raise RuntimeError(\"WeightWatcher must be called exactly once\")\n", + "show_details(details)" ] }, { @@ -228,7 +228,7 @@ "id": "compare-note", "metadata": {}, "source": [ - "## Standard versus clipped fit\n" + "## Raw versus clipped alpha from the same call\n" ] }, { @@ -238,40 +238,18 @@ "metadata": {}, "outputs": [], "source": [ - "metrics_to_compare = [\n", - " \"alpha\", \"raw_alpha\", \"D\", \"xmin\", \"xmax\", \"num_fingers\", \"num_pl_spikes\"\n", + "comparison_columns = [\n", + " \"matrix_name\", \"alpha_raw\", \"alpha_clip_xmax\", \"alpha_reduction\",\n", + " \"D\", \"xmin\", \"xmax\", \"num_fingers\", \"num_pl_spikes\",\n", "]\n", - "std_cols = [\"matrix_name\"] + [c for c in metrics_to_compare if c in details_standard.columns]\n", - "clip_cols = [\"matrix_name\"] + [c for c in metrics_to_compare if c in details_clip.columns]\n", - "\n", - "comparison = (\n", - " details_standard[std_cols]\n", - " .rename(columns={c: f\"{c}_standard\" for c in std_cols if c != \"matrix_name\"})\n", - " .merge(\n", - " details_clip[clip_cols].rename(\n", - " columns={c: f\"{c}_clip_xmax\" for c in clip_cols if c != \"matrix_name\"}\n", - " ),\n", - " on=\"matrix_name\",\n", - " how=\"inner\",\n", - " )\n", - ")\n", - "\n", - "if {\"alpha_standard\", \"alpha_clip_xmax\"} <= set(comparison.columns):\n", - " comparison[\"delta_alpha_clip_minus_standard\"] = (\n", - " comparison[\"alpha_clip_xmax\"] - comparison[\"alpha_standard\"]\n", - " )\n", - " comparison[\"alpha_reduction\"] = (\n", - " comparison[\"alpha_standard\"] - comparison[\"alpha_clip_xmax\"]\n", - " )\n", - "\n", - "comparison = comparison.sort_values(\"matrix_name\")\n", - "comparison.to_csv(OUTPUT_DIR / \"standard_vs_clip_xmax.csv\", index=False)\n", + "comparison = details[[c for c in comparison_columns if c in details.columns]].sort_values(\"matrix_name\")\n", + "comparison.to_csv(OUTPUT_DIR / \"raw_vs_clip_xmax_one_pass.csv\", index=False)\n", "display(comparison)\n", "\n", - "if \"alpha_standard\" in comparison and \"alpha_clip_xmax\" in comparison:\n", - " print(\"median alpha, standard :\", float(comparison[\"alpha_standard\"].median()))\n", + "if \"alpha_raw\" in comparison and \"alpha_clip_xmax\" in comparison:\n", + " print(\"median alpha, raw :\", float(comparison[\"alpha_raw\"].median()))\n", " print(\"median alpha, clip_xmax:\", float(comparison[\"alpha_clip_xmax\"].median()))\n", - " print(\"mean alpha, standard :\", float(comparison[\"alpha_standard\"].mean()))\n", + " print(\"mean alpha, raw :\", float(comparison[\"alpha_raw\"].mean()))\n", " print(\"mean alpha, clip_xmax :\", float(comparison[\"alpha_clip_xmax\"].mean()))" ] }, @@ -282,7 +260,7 @@ "source": [ "## All six ESDs\n", "\n", - "The empirical ESD is unchanged by `clip_xmax`; the fit boundaries change. The vertical lines show the standard `xmin`, clipped `xmin`, and clipped `xmax`.\n" + "The empirical ESD is unchanged by `clip_xmax`. The vertical lines show the corrected fit's `xmin` and `xmax`; a separate raw-fit boundary is not fabricated.\n" ] }, { @@ -296,11 +274,11 @@ "axes = axes.ravel()\n", "\n", "for ax, (_, row) in zip(\n", - " axes, details_standard.sort_values(\"matrix_name\").iterrows()\n", + " axes, details.sort_values(\"matrix_name\").iterrows()\n", "):\n", " matrix_name = str(row[\"matrix_name\"])\n", " layer_id = int(row[\"layer_id\"])\n", - " esd = np.asarray(watcher_standard.get_ESD(layer=layer_id), dtype=float)\n", + " esd = np.asarray(watcher.get_ESD(layer=layer_id), dtype=float)\n", " esd = esd[np.isfinite(esd) & (esd > 0)]\n", "\n", " if esd.size == 0:\n", @@ -318,19 +296,17 @@ " valid = np.isfinite(hist) & (hist > 0) & np.isfinite(centers) & (centers > 0)\n", " ax.loglog(centers[valid], hist[valid], marker=\"o\", linestyle=\"none\", label=\"ESD\")\n", "\n", - " clip_row = details_clip.loc[details_clip[\"matrix_name\"] == matrix_name].iloc[0]\n", " for value, label, style in [\n", - " (row.get(\"xmin\", np.nan), \"xmin standard\", \"--\"),\n", - " (clip_row.get(\"xmin\", np.nan), \"xmin clip_xmax\", \":\"),\n", - " (clip_row.get(\"xmax\", np.nan), \"xmax clip_xmax\", \"-.\"),\n", + " (row.get(\"xmin\", np.nan), \"xmin clip_xmax\", \":\"),\n", + " (row.get(\"xmax\", np.nan), \"xmax clip_xmax\", \"-.\"),\n", " ]:\n", " value = pd.to_numeric(pd.Series([value]), errors=\"coerce\").iloc[0]\n", " if np.isfinite(value) and value > 0:\n", " ax.axvline(float(value), linestyle=style, label=label)\n", "\n", - " alpha_std = float(row[\"alpha\"]) if pd.notna(row.get(\"alpha\")) else float(\"nan\")\n", - " alpha_clip = float(clip_row[\"alpha\"]) if pd.notna(clip_row.get(\"alpha\")) else float(\"nan\")\n", - " ax.set_title(f\"{matrix_name} alpha: {alpha_std:.3f} -> {alpha_clip:.3f}\")\n", + " alpha_raw = float(row[\"raw_alpha\"]) if pd.notna(row.get(\"raw_alpha\")) else float(\"nan\")\n", + " alpha_clip = float(row[\"alpha\"]) if pd.notna(row.get(\"alpha\")) else float(\"nan\")\n", + " ax.set_title(f\"{matrix_name} alpha: raw {alpha_raw:.3f} -> clip {alpha_clip:.3f}\")\n", " ax.set_xlabel(\"eigenvalue of X = W^T W\")\n", " ax.set_ylabel(\"density\")\n", " ax.grid(True, which=\"both\", alpha=0.2)\n", @@ -340,7 +316,7 @@ " f\"MuonClip + RMS, epoch {NOMINAL_EPOCH:.2f}, step {STEP}: ESD + clip_xmax\",\n", " fontsize=14,\n", ")\n", - "figure_path = OUTPUT_DIR / \"all_layer_esds_standard_vs_clip_xmax.png\"\n", + "figure_path = OUTPUT_DIR / \"all_layer_esds_raw_vs_clip_xmax_one_pass.png\"\n", "fig.savefig(figure_path, dpi=180, bbox_inches=\"tight\")\n", "plt.show()\n", "print(\"saved:\", figure_path)" @@ -353,7 +329,7 @@ "source": [ "## Interpretation\n", "\n", - "- Large `alpha_reduction` with nonzero `num_fingers`: the standard alpha was strongly influenced by top-of-spectrum finite-size fingers.\n", + "- Large `alpha_reduction` with nonzero `num_fingers`: the raw alpha was strongly influenced by top-of-spectrum finite-size fingers.\n", "- Little alpha reduction: the high alpha is not explained by those fingers.\n", "- Always inspect `D` together with alpha.\n", "- Outputs are saved below `$RUN_DIR/diagnostics/esd_clip_xmax_step_XXXXXXX/`.\n" @@ -373,4 +349,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/baseline/nanogpt_one_head/pyproject.toml b/baseline/nanogpt_one_head/pyproject.toml index 325163bf..2dae8847 100644 --- a/baseline/nanogpt_one_head/pyproject.toml +++ b/baseline/nanogpt_one_head/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "rg-nanogpt-one-head" -version = "0.3.4" +version = "0.5.1" description = "Matched one-head nanoGPT optimizer baselines on pinned FineWeb-Edu" requires-python = ">=3.10" dependencies = [ @@ -22,6 +22,7 @@ dependencies = [ "jupyter>=1.0", "ipykernel>=6.29", "papermill>=2.6,<3", + "packaging>=23", ] [project.optional-dependencies] diff --git a/baseline/nanogpt_one_head/requirements.txt b/baseline/nanogpt_one_head/requirements.txt index 0190ae5e..2c15ac71 100644 --- a/baseline/nanogpt_one_head/requirements.txt +++ b/baseline/nanogpt_one_head/requirements.txt @@ -7,9 +7,11 @@ pyyaml>=6.0 datasets>=2.19 tiktoken>=0.7 sacrebleu>=2.4 -weightwatcher>=0.7.7 +weightwatcher==0.7.7 powerlaw>=2.0.0,<3 jupyter>=1.0 ipykernel>=6.29 +papermill>=2.6,<3 +packaging>=23 pytest>=8.0 nbformat>=5.10 diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/analysis.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/analysis.py index 621c4b4a..fd3d427c 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/analysis.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/analysis.py @@ -14,6 +14,7 @@ OPTIMIZER_LABELS = { "sgd_momentum": "SGD + Nesterov", + "adam": "Adam", "adamw": "AdamW", "muon": "Muon + auxiliary AdamW", } @@ -21,6 +22,7 @@ # Okabe-Ito color-blind-safe optimizer palette. OPTIMIZER_COLORS = { "sgd_momentum": "#0072B2", + "adam": "#E69F00", "adamw": "#D55E00", "muon": "#009E73", } @@ -93,7 +95,7 @@ def run_status_table( results_root: str | Path, *, optimizers: Sequence[str] = SUPPORTED_OPTIMIZERS, - seeds: Sequence[int] = (1337, 2027, 4099), + seeds: Sequence[int] = (1337, 2027, 4099, 31415, 271828), ) -> pd.DataFrame: rows = [] for optimizer in optimizers: @@ -160,7 +162,7 @@ def load_metrics( results_root: str | Path, *, optimizers: Sequence[str] = SUPPORTED_OPTIMIZERS, - seeds: Sequence[int] = (1337, 2027, 4099), + seeds: Sequence[int] = (1337, 2027, 4099, 31415, 271828), require_complete: bool = True, ) -> pd.DataFrame: frame = _load_csvs( @@ -181,7 +183,7 @@ def load_epoch_metrics( results_root: str | Path, *, optimizers: Sequence[str] = SUPPORTED_OPTIMIZERS, - seeds: Sequence[int] = (1337, 2027, 4099), + seeds: Sequence[int] = (1337, 2027, 4099, 31415, 271828), require_complete: bool = True, ) -> pd.DataFrame: frame = _load_csvs( @@ -202,7 +204,7 @@ def load_layer_metrics( results_root: str | Path, *, optimizers: Sequence[str] = SUPPORTED_OPTIMIZERS, - seeds: Sequence[int] = (1337, 2027, 4099), + seeds: Sequence[int] = (1337, 2027, 4099, 31415, 271828), require_complete: bool = True, ) -> pd.DataFrame: frame = _load_csvs( @@ -223,7 +225,7 @@ def load_spectral_summary( results_root: str | Path, *, optimizers: Sequence[str] = SUPPORTED_OPTIMIZERS, - seeds: Sequence[int] = (1337, 2027, 4099), + seeds: Sequence[int] = (1337, 2027, 4099, 31415, 271828), require_complete: bool = True, ) -> pd.DataFrame: frame = _load_csvs( @@ -244,7 +246,7 @@ def load_test_results( results_root: str | Path, *, optimizers: Sequence[str] = SUPPORTED_OPTIMIZERS, - seeds: Sequence[int] = (1337, 2027, 4099), + seeds: Sequence[int] = (1337, 2027, 4099, 31415, 271828), ) -> pd.DataFrame: rows = [] for optimizer in optimizers: @@ -263,8 +265,22 @@ def load_test_results( "step": int(values["step"]), "test_loss": float(values["loss"]), "test_perplexity": float(values["perplexity"]), + "test_bits_per_token": float( + values.get("bits_per_token", np.nan) + ), "test_accuracy": float(values["accuracy"]), + "test_top5_accuracy": float( + values.get("top5_accuracy", np.nan) + ), "test_bleu": float(values["bleu"]), + "test_continuation_token_accuracy": float( + values.get( + "continuation_token_accuracy", np.nan + ) + ), + "test_continuation_exact_match": float( + values.get("continuation_exact_match", np.nan) + ), } ) return pd.DataFrame(rows) @@ -414,7 +430,16 @@ def plot_spectral_optimizer_summary( def final_test_summary(test_results: pd.DataFrame) -> pd.DataFrame: rows = [] for (optimizer, checkpoint), group in test_results.groupby(["optimizer", "checkpoint"]): - for metric in ("test_loss", "test_perplexity", "test_accuracy", "test_bleu"): + for metric in ( + "test_loss", + "test_perplexity", + "test_bits_per_token", + "test_accuracy", + "test_top5_accuracy", + "test_bleu", + "test_continuation_token_accuracy", + "test_continuation_exact_match", + ): rows.append( { "optimizer": optimizer, diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py index 1349ca80..8818f251 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py @@ -1,5 +1,8 @@ from __future__ import annotations +import hashlib +import json +import math from pathlib import Path import random from typing import Any @@ -21,6 +24,101 @@ ) +def model_state_sha256(state: dict[str, torch.Tensor]) -> str: + """Hash model tensor names, shapes, dtypes, and exact bytes.""" + + digest = hashlib.sha256() + if not state: + raise ValueError("model state is empty") + for name in sorted(state): + value = state[name] + if not torch.is_tensor(value): + raise TypeError(f"model state entry is not a tensor: {name}") + tensor = value.detach().to("cpu").contiguous() + metadata = json.dumps( + { + "name": str(name), + "shape": list(tensor.shape), + "dtype": str(tensor.dtype), + }, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + digest.update(len(metadata).to_bytes(8, "big")) + digest.update(metadata) + raw = tensor.reshape(-1).view(torch.uint8).numpy().tobytes() + digest.update(len(raw).to_bytes(8, "big")) + digest.update(raw) + return digest.hexdigest() + + +def _update_state_digest(digest: Any, value: Any) -> None: + if torch.is_tensor(value): + tensor = value.detach().to("cpu").contiguous() + metadata = json.dumps( + {"shape": list(tensor.shape), "dtype": str(tensor.dtype)}, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + digest.update(b"tensor\0") + digest.update(len(metadata).to_bytes(8, "big")) + digest.update(metadata) + raw = tensor.reshape(-1).view(torch.uint8).numpy().tobytes() + digest.update(len(raw).to_bytes(8, "big")) + digest.update(raw) + return + if isinstance(value, dict): + digest.update(b"dict\0") + items = sorted( + value.items(), + key=lambda item: (type(item[0]).__name__, repr(item[0])), + ) + digest.update(len(items).to_bytes(8, "big")) + for key, item in items: + _update_state_digest(digest, key) + _update_state_digest(digest, item) + return + if isinstance(value, (list, tuple)): + digest.update(b"list\0" if isinstance(value, list) else b"tuple\0") + digest.update(len(value).to_bytes(8, "big")) + for item in value: + _update_state_digest(digest, item) + return + if value is None: + digest.update(b"none\0") + return + if isinstance(value, bool): + digest.update(b"bool\0" + (b"1" if value else b"0")) + return + if isinstance(value, int): + digest.update(b"int\0" + str(value).encode("ascii") + b"\0") + return + if isinstance(value, float): + digest.update(b"float\0" + value.hex().encode("ascii") + b"\0") + return + if isinstance(value, str): + encoded = value.encode("utf-8") + digest.update(b"str\0" + len(encoded).to_bytes(8, "big") + encoded) + return + if isinstance(value, bytes): + digest.update(b"bytes\0" + len(value).to_bytes(8, "big") + value) + return + raise TypeError( + "unsupported optimizer-state value for integrity hashing: " + f"{type(value).__name__}" + ) + + +def optimizer_state_sha256(states: list[dict[str, Any]]) -> str: + """Hash the complete optimizer-state structure and exact tensor bytes.""" + + if not states: + raise ValueError("optimizer state inventory is empty") + digest = hashlib.sha256() + _update_state_digest(digest, states) + return digest.hexdigest() + + def _nonfinite_tensor_paths(value: Any, path: str) -> list[str]: bad: list[str] = [] if torch.is_tensor(value): @@ -28,6 +126,9 @@ def _nonfinite_tensor_paths(value: Any, path: str) -> list[str]: if not bool(torch.isfinite(value).all()): bad.append(path) return bad + if isinstance(value, float) and not math.isfinite(value): + bad.append(path) + return bad if isinstance(value, dict): for key, item in value.items(): bad.extend( @@ -48,7 +149,7 @@ def _nonfinite_tensor_paths(value: Any, path: str) -> list[str]: return bad -def _require_finite_checkpoint_state( +def require_finite_checkpoint_state( *, model_state: dict[str, Any], optimizer_states: list[dict[str, Any]] | None, @@ -93,6 +194,7 @@ def save_training_checkpoint( optimizer_name: str, seed: int, train_generator: torch.Generator, + resume_diagnostics: dict[str, Any] | None = None, ) -> Path: device = model_device(model) synchronize(device) @@ -103,18 +205,72 @@ def save_training_checkpoint( # gate ensures checkpoint_latest.pt always remains the last verified state. model_state = tree_to_cpu(model.state_dict()) optimizer_states = tree_to_cpu(optimizer_state_dict(handles)) - _require_finite_checkpoint_state( + require_finite_checkpoint_state( model_state=model_state, optimizer_states=optimizer_states, step=step, ) + persisted_diagnostics: dict[str, Any] | None = None + if resume_diagnostics is not None: + previous_snapshot = tree_to_cpu( + resume_diagnostics.get("previous_eval_snapshot") + ) + if not isinstance(previous_snapshot, list) or not previous_snapshot: + raise ValueError( + "resume diagnostics must contain a non-empty " + "previous_eval_snapshot list" + ) + expected_parameters = list(model.parameters()) + if len(previous_snapshot) != len(expected_parameters): + raise ValueError( + "resume diagnostic parameter inventory does not match the model" + ) + for index, (snapshot, parameter) in enumerate( + zip(previous_snapshot, expected_parameters, strict=True) + ): + if not torch.is_tensor(snapshot): + raise TypeError( + "resume diagnostic snapshot entry is not a tensor: " + f"index={index}" + ) + if tuple(snapshot.shape) != tuple(parameter.shape): + raise ValueError( + "resume diagnostic snapshot shape does not match the model: " + f"index={index}, snapshot={tuple(snapshot.shape)}, " + f"parameter={tuple(parameter.shape)}" + ) + bad_diagnostics = _nonfinite_tensor_paths( + previous_snapshot, + "resume_diagnostics.previous_eval_snapshot", + ) + if bad_diagnostics: + raise FloatingPointError( + "refusing to write contaminated resume diagnostics: " + + ", ".join(bad_diagnostics[:12]) + ) + last_grad_pre = float(resume_diagnostics["last_grad_pre"]) + last_grad_post = float(resume_diagnostics["last_grad_post"]) + if not math.isfinite(last_grad_pre) or not math.isfinite(last_grad_post): + raise FloatingPointError( + "refusing to write non-finite resume diagnostic gradients" + ) + persisted_diagnostics = { + "schema_version": 1, + "previous_eval_snapshot": previous_snapshot, + "last_grad_pre": last_grad_pre, + "last_grad_post": last_grad_post, + "last_clipped": bool(resume_diagnostics["last_clipped"]), + } + payload: dict[str, Any] = { - "schema_version": 4, + "schema_version": 5, # CPU tensors keep checkpoints portable between MPS, CUDA, TPU/XLA, # and CPU environments. "model": model_state, "optimizers": optimizer_states, + "model_state_sha256": model_state_sha256(model_state), + "optimizer_state_sha256": optimizer_state_sha256(optimizer_states), "step": int(step), "best_validation_loss": float(best_validation_loss), "best_validation_step": int(best_validation_step), @@ -127,26 +283,33 @@ def save_training_checkpoint( "numpy_random_state": np.random.get_state(), "torch_random_state": torch.random.get_rng_state(), "train_generator_state": train_generator.get_state(), + "resume_diagnostics": persisted_diagnostics, **capture_accelerator_rng_state(device), } return _atomic_torch_save(payload, Path(path)) -def load_training_checkpoint( +def _load_training_checkpoint( path: str | Path, *, model, handles: list[OptimizerHandle], expected_fingerprint: str, train_generator: torch.Generator, -) -> tuple[int, float, int, float]: +) -> tuple[int, float, int, float, dict[str, Any] | None]: path = Path(path) payload = torch.load(path, map_location="cpu", weights_only=False) if str(payload.get("fingerprint")) != str(expected_fingerprint): raise RuntimeError( "checkpoint protocol fingerprint does not match the requested run" ) - _require_finite_checkpoint_state( + model_hash = model_state_sha256(payload["model"]) + if str(payload.get("model_state_sha256", "")) != model_hash: + raise RuntimeError("checkpoint model-state SHA-256 does not match") + optimizer_hash = optimizer_state_sha256(payload["optimizers"]) + if str(payload.get("optimizer_state_sha256", "")) != optimizer_hash: + raise RuntimeError("checkpoint optimizer-state SHA-256 does not match") + require_finite_checkpoint_state( model_state=payload["model"], optimizer_states=payload["optimizers"], step=int(payload.get("step", -1)), @@ -158,11 +321,86 @@ def load_training_checkpoint( torch.random.set_rng_state(payload["torch_random_state"]) train_generator.set_state(payload["train_generator_state"]) restore_accelerator_rng_state(payload, model_device(model)) + diagnostics = payload.get("resume_diagnostics") + if diagnostics is not None: + if not isinstance(diagnostics, dict): + raise RuntimeError("checkpoint resume_diagnostics is not a mapping") + if int(diagnostics.get("schema_version", -1)) != 1: + raise RuntimeError( + "checkpoint resume_diagnostics has an unsupported schema" + ) + previous_snapshot = diagnostics.get("previous_eval_snapshot") + expected_parameters = list(model.parameters()) + if ( + not isinstance(previous_snapshot, list) + or len(previous_snapshot) != len(expected_parameters) + ): + raise RuntimeError( + "checkpoint resume diagnostic parameter inventory is invalid" + ) + for index, (snapshot, parameter) in enumerate( + zip(previous_snapshot, expected_parameters, strict=True) + ): + if ( + not torch.is_tensor(snapshot) + or tuple(snapshot.shape) != tuple(parameter.shape) + or not bool(torch.isfinite(snapshot).all()) + ): + raise RuntimeError( + "checkpoint resume diagnostic snapshot is invalid at " + f"parameter index {index}" + ) + diagnostics = { + "previous_eval_snapshot": tree_to_cpu(previous_snapshot), + "last_grad_pre": float(diagnostics["last_grad_pre"]), + "last_grad_post": float(diagnostics["last_grad_post"]), + "last_clipped": bool(diagnostics["last_clipped"]), + } return ( int(payload["step"]), float(payload["best_validation_loss"]), int(payload["best_validation_step"]), float(payload["elapsed_seconds"]), + diagnostics, + ) + + +def load_training_checkpoint( + path: str | Path, + *, + model, + handles: list[OptimizerHandle], + expected_fingerprint: str, + train_generator: torch.Generator, +) -> tuple[int, float, int, float]: + """Load a checkpoint while preserving the historical four-value API.""" + + loaded = _load_training_checkpoint( + path, + model=model, + handles=handles, + expected_fingerprint=expected_fingerprint, + train_generator=train_generator, + ) + return loaded[:4] + + +def load_training_checkpoint_for_resume( + path: str | Path, + *, + model, + handles: list[OptimizerHandle], + expected_fingerprint: str, + train_generator: torch.Generator, +) -> tuple[int, float, int, float, dict[str, Any] | None]: + """Load training state plus deterministic monitoring diagnostics.""" + + return _load_training_checkpoint( + path, + model=model, + handles=handles, + expected_fingerprint=expected_fingerprint, + train_generator=train_generator, ) @@ -187,7 +425,7 @@ def save_epoch_model_checkpoint( device = model_device(model) synchronize(device) model_state = tree_to_cpu(model.state_dict()) - _require_finite_checkpoint_state( + require_finite_checkpoint_state( model_state=model_state, optimizer_states=None, step=step, @@ -195,6 +433,7 @@ def save_epoch_model_checkpoint( payload = { "schema_version": 3, "model": model_state, + "model_state_sha256": model_state_sha256(model_state), "step": int(step), "nominal_epoch": float(nominal_epoch), "actual_epoch": float(actual_epoch), diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/completion.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/completion.py index fa62bc03..247ab679 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/completion.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/completion.py @@ -2,6 +2,7 @@ from __future__ import annotations +import hashlib import json import math from pathlib import Path @@ -11,11 +12,18 @@ import pandas as pd import torch +from .checkpoints import ( + model_state_sha256, + optimizer_state_sha256, + require_finite_checkpoint_state, +) + _REQUIRED_FILES = ( "run_complete.json", "manifest.json", "metrics.csv", "epoch_metrics.csv", + "checkpoint_initial.pt", "checkpoint_latest.pt", "checkpoint_best.pt", "checkpoint_final.pt", @@ -23,6 +31,29 @@ "spectral/layers.csv", "spectral/summary.csv", ) +_TEST_METRIC_COMPLETION_KEYS = { + "loss": "final_test_loss", + "perplexity": "final_test_perplexity", + "bits_per_token": "final_test_bits_per_token", + "accuracy": "final_test_accuracy", + "top5_accuracy": "final_test_top5_accuracy", + "bleu": "final_test_bleu", + "continuation_token_accuracy": ( + "final_test_continuation_token_accuracy" + ), + "continuation_exact_match": "final_test_continuation_exact_match", +} +_HELD_OUT_CURVE_COLUMNS = ( + "test_loss", + "test_perplexity", + "test_bits_per_token", + "test_accuracy", + "test_top5_accuracy", + "test_bleu", + "test_continuation_token_accuracy", + "test_continuation_exact_match", + "test_generalization_gap", +) class CompletedRunValidationError(RuntimeError): @@ -57,6 +88,14 @@ def _read_csv(path: Path) -> pd.DataFrame: return frame +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(1024 * 1024): + digest.update(block) + return digest.hexdigest() + + def _as_int(value: Any, label: str) -> int: try: return int(value) @@ -64,6 +103,16 @@ def _as_int(value: Any, label: str) -> int: _fail(f"{label} is not an integer: {value!r}") +def _as_finite_float(value: Any, label: str) -> float: + try: + result = float(value) + except (TypeError, ValueError): + _fail(f"{label} is not numeric: {value!r}") + if not math.isfinite(result): + _fail(f"{label} is non-finite: {value!r}") + return result + + def _expect(observed: Any, expected: Any, label: str) -> None: if observed != expected: _fail( @@ -96,6 +145,170 @@ def _load_checkpoint(path: Path) -> dict[str, Any]: return payload +def _validate_checkpoint_identity( + payload: dict[str, Any], + *, + path: Path, + fingerprint: str, + optimizer: str, + seed: int, + step: int, + schema_version: int, +) -> None: + _expect( + _as_int(payload.get("schema_version"), f"{path.name} schema_version"), + schema_version, + f"{path.name} schema_version", + ) + _expect( + str(payload.get("fingerprint", "")), + fingerprint, + f"{path.name} fingerprint", + ) + _expect( + str(payload.get("optimizer_name", "")), + optimizer, + f"{path.name} optimizer", + ) + _expect( + _as_int(payload.get("seed"), f"{path.name} seed"), + seed, + f"{path.name} seed", + ) + _expect( + _as_int(payload.get("step"), f"{path.name} step"), + step, + f"{path.name} step", + ) + model_state = payload.get("model") + if not isinstance(model_state, dict) or not model_state: + _fail(f"{path.name} has no non-empty model state") + + +def _validate_model_tensor_inventory( + state: dict[str, torch.Tensor], + *, + path: Path, + expected: dict[str, tuple[tuple[int, ...], str]] | None = None, +) -> dict[str, tuple[tuple[int, ...], str]]: + inventory: dict[str, tuple[tuple[int, ...], str]] = {} + for name, value in state.items(): + if not torch.is_tensor(value): + _fail(f"{path.name} model entry is not a tensor: {name}") + if (value.is_floating_point() or value.is_complex()) and not bool( + torch.isfinite(value).all() + ): + _fail(f"{path.name} contains a non-finite model tensor: {name}") + inventory[str(name)] = (tuple(value.shape), str(value.dtype)) + if expected is not None and inventory != expected: + _fail( + f"{path.name} model key/shape/dtype inventory differs from " + "checkpoint_initial.pt" + ) + return inventory + + +def _validate_muonclip_qk( + root: Path, + *, + manifest: dict[str, Any], + total_steps: int, +) -> None: + path = root / "muonclip_qk.csv" + if not path.is_file() or path.stat().st_size == 0: + _fail(f"MuonClip run lacks required QK diagnostics: {path}") + frame = _read_csv(path) + required = ( + "step", + "threshold", + "steps_in_interval", + "head_observations", + "active_heads", + "active_fraction", + "mean_max_logit", + "max_logit", + "mean_gamma", + "min_gamma", + ) + missing = set(required).difference(frame.columns) + if missing: + _fail( + "MuonClip QK diagnostics lack columns: " + + ", ".join(sorted(missing)) + ) + numeric = frame[list(required)].apply(pd.to_numeric, errors="coerce") + if not np.isfinite(numeric.to_numpy(dtype=float)).all(): + _fail("MuonClip QK diagnostics contain non-finite values") + + profile = manifest.get("optimizer_profile") + if not isinstance(profile, dict): + _fail("manifest optimizer_profile is not a mapping") + diagnostic_interval = _as_int( + profile.get("qk_diagnostics_interval"), + "manifest MuonClip QK diagnostic interval", + ) + if diagnostic_interval <= 0: + _fail("manifest has a non-positive MuonClip QK diagnostic interval") + expected_steps = list( + range(diagnostic_interval, total_steps + 1, diagnostic_interval) + ) + if not expected_steps or expected_steps[-1] != total_steps: + expected_steps.append(total_steps) + observed_steps = numeric["step"].to_numpy(dtype=float) + if not np.allclose(observed_steps, np.rint(observed_steps)): + _fail("MuonClip QK diagnostic steps are not integers") + if tuple(int(value) for value in observed_steps) != tuple(expected_steps): + _fail("MuonClip QK diagnostics do not cover the exact interval grid") + + expected_intervals = np.diff(np.asarray([0, *expected_steps], dtype=int)) + observed_intervals = numeric["steps_in_interval"].to_numpy(dtype=float) + if not np.array_equal(observed_intervals, expected_intervals.astype(float)): + _fail("MuonClip QK interval lengths do not match their recorded steps") + if int(observed_intervals.sum()) != total_steps: + _fail("MuonClip QK intervals do not cover the full training horizon") + + model = manifest.get("model") + if not isinstance(model, dict): + _fail("manifest model configuration is not a mapping") + heads_per_step = _as_int(model.get("n_layer"), "manifest n_layer") * _as_int( + model.get("n_head"), "manifest n_head" + ) + if heads_per_step <= 0: + _fail("manifest has a non-positive layer/head inventory") + expected_observations = observed_intervals * heads_per_step + observations = numeric["head_observations"].to_numpy(dtype=float) + active_heads = numeric["active_heads"].to_numpy(dtype=float) + active_fraction = numeric["active_fraction"].to_numpy(dtype=float) + if not np.array_equal(observations, expected_observations): + _fail("MuonClip QK head observations do not match one observation/head/step") + if ( + (active_heads < 0.0).any() + or (active_heads > observations).any() + or not np.allclose( + active_fraction, + active_heads / observations, + rtol=1e-12, + atol=1e-12, + ) + ): + _fail("MuonClip QK active-head counts/fractions are inconsistent") + + threshold = _as_finite_float( + profile.get("qk_clip_threshold"), + "manifest MuonClip QK threshold", + ) + if not numeric["threshold"].eq(threshold).all(): + _fail("MuonClip QK thresholds differ from the optimizer profile") + if ( + not numeric["active_fraction"].between(0.0, 1.0).all() + or not numeric["mean_gamma"].between(0.0, 1.0).all() + or not numeric["min_gamma"].between(0.0, 1.0).all() + or (numeric["min_gamma"] > numeric["mean_gamma"]).any() + or (numeric["mean_max_logit"] > numeric["max_logit"]).any() + ): + _fail("MuonClip QK diagnostics violate registered bounds") + + def validate_completed_run( run_dir: str | Path, *, @@ -176,11 +389,93 @@ def validate_completed_run( total_steps, "manifest max_steps", ) + initial_model_hash = str(manifest.get("initial_model_sha256", "")) + if ( + len(initial_model_hash) != 64 + or any(character not in "0123456789abcdef" for character in initial_model_hash) + ): + _fail("manifest has no valid initial-model tensor hash") final_test = test_results.get("final") selected_test = test_results.get("validation_selected") + test_policy = str(test_results.get("policy", "")).lower() + if ( + "held out" not in test_policy + or "validation" not in test_policy + or "never" not in test_policy + ): + _fail("test_results.json does not declare the held-out test policy") if not isinstance(final_test, dict) or not isinstance(selected_test, dict): _fail("test_results.json lacks final or validation_selected results") + parsed_test_metrics: dict[str, dict[str, float]] = {} + for label, values in ( + ("final", final_test), + ("validation_selected", selected_test), + ): + missing_metrics = set(_TEST_METRIC_COMPLETION_KEYS).difference(values) + if missing_metrics: + _fail( + f"test_results.json {label} lacks metrics: " + + ", ".join(sorted(missing_metrics)) + ) + parsed = { + metric: _as_finite_float( + values[metric], f"{label} test {metric}" + ) + for metric in _TEST_METRIC_COMPLETION_KEYS + } + for metric in ( + "accuracy", + "top5_accuracy", + "continuation_token_accuracy", + "continuation_exact_match", + ): + if not 0.0 <= parsed[metric] <= 1.0: + _fail(f"{label} test {metric} is outside [0, 1]") + if parsed["top5_accuracy"] < parsed["accuracy"]: + _fail(f"{label} test top5_accuracy is below top-1 accuracy") + if ( + parsed["continuation_exact_match"] + > parsed["continuation_token_accuracy"] + 1e-12 + ): + _fail( + f"{label} continuation exact-match exceeds token accuracy" + ) + if parsed["loss"] < 0.0 or parsed["bits_per_token"] < 0.0: + _fail(f"{label} test loss/bits_per_token must be nonnegative") + if parsed["perplexity"] <= 0.0: + _fail(f"{label} test perplexity must be positive") + if not 0.0 <= parsed["bleu"] <= 100.0: + _fail(f"{label} test BLEU is outside [0, 100]") + if not math.isclose( + math.log(parsed["perplexity"]), + parsed["loss"], + rel_tol=1e-10, + abs_tol=1e-10, + ): + _fail(f"{label} test perplexity is inconsistent with loss") + if not math.isclose( + parsed["bits_per_token"] * math.log(2.0), + parsed["loss"], + rel_tol=1e-10, + abs_tol=1e-10, + ): + _fail(f"{label} test bits_per_token is inconsistent with loss") + parsed_test_metrics[label] = parsed + + for metric, completion_key in _TEST_METRIC_COMPLETION_KEYS.items(): + completion_value = _as_finite_float( + completion.get(completion_key), f"completion {completion_key}" + ) + if not math.isclose( + completion_value, + parsed_test_metrics["final"][metric], + rel_tol=1e-12, + abs_tol=1e-12, + ): + _fail( + f"completion {completion_key} does not match final test {metric}" + ) _expect( _as_int(final_test.get("step"), "final test step"), total_steps, @@ -199,6 +494,10 @@ def validate_completed_run( metric_steps = _step_tuple(metrics, "metrics.csv") epoch_steps = _step_tuple(epoch_metrics, "epoch_metrics.csv") summary_steps = _step_tuple(summary, "spectral/summary.csv") + weightwatcher = manifest.get("weightwatcher", {}) + if not isinstance(weightwatcher, dict): + _fail("manifest weightwatcher configuration is not a mapping") + clip_xmax = weightwatcher.get("fix_fingers") == "clip_xmax" for label, steps in ( ("metrics.csv", metric_steps), @@ -207,11 +506,39 @@ def validate_completed_run( if 0 not in steps or total_steps not in steps or max(steps) != total_steps: _fail(f"{label} does not span step zero through {total_steps}") + for label, frame in ( + ("metrics.csv", metrics), + ("epoch_metrics.csv", epoch_metrics), + ): + missing_held_out = set(_HELD_OUT_CURVE_COLUMNS).difference( + frame.columns + ) + if missing_held_out: + _fail( + f"{label} is missing held-out placeholder columns " + + ", ".join(sorted(missing_held_out)) + ) + leaked = [ + column + for column in _HELD_OUT_CURVE_COLUMNS + if not frame[column].isna().all() + ] + if leaked: + _fail( + f"{label} leaks held-out test outcomes into training curves: " + + ", ".join(leaked) + ) + if "test_monitoring_only" not in epoch_metrics.columns: - _fail("epoch_metrics.csv has no test_monitoring_only column") + _fail("epoch_metrics.csv has no legacy test policy flag") policy = pd.to_numeric(epoch_metrics["test_monitoring_only"], errors="coerce") if policy.isna().any() or not policy.astype(int).eq(1).all(): - _fail("epoch_metrics.csv violates the monitoring-only test policy") + _fail("epoch_metrics.csv violates the held-out test policy") + if "test_held_out" not in epoch_metrics.columns: + _fail("epoch_metrics.csv has no test_held_out column") + held_out = pd.to_numeric(epoch_metrics["test_held_out"], errors="coerce") + if held_out.isna().any() or not held_out.astype(int).eq(1).all(): + _fail("epoch_metrics.csv does not mark every test curve as held out") required_layer_columns = { "step", @@ -219,7 +546,24 @@ def validate_completed_run( "alpha", "ERG_gap", "num_traps", + "run_seed", + "diagnostic_seed", + "protocol_fingerprint", + "model_state_sha256", } + if clip_xmax: + required_layer_columns.update( + { + "raw_alpha", + "alpha_raw", + "alpha_clip_xmax", + "alpha_delta", + "num_fingers", + "finger_policy", + "primary_alpha_variant", + "weightwatcher_analysis_calls", + } + ) missing_columns = required_layer_columns.difference(layers.columns) if missing_columns: _fail( @@ -236,6 +580,188 @@ def validate_completed_run( _fail("spectral steps do not match epoch_metrics.csv") if not layers.groupby("step")["matrix_name"].nunique().eq(6).all(): _fail("spectral/layers.csv does not contain six matrices per epoch") + if not layers["protocol_fingerprint"].astype(str).eq(fingerprint).all(): + _fail("spectral/layers.csv fingerprint does not match the run") + layer_seeds = pd.to_numeric(layers["run_seed"], errors="coerce") + if layer_seeds.isna().any() or not layer_seeds.astype(int).eq(seed).all(): + _fail("spectral/layers.csv seed does not match the run") + layer_diagnostic_seeds = pd.to_numeric( + layers["diagnostic_seed"], errors="coerce" + ) + expected_diagnostic_seeds = ( + pd.to_numeric(layers["step"], errors="coerce").astype(int) + + int(seed) + + 1_000_003 + ) + if ( + layer_diagnostic_seeds.isna().any() + or not layer_diagnostic_seeds.astype(int).equals( + expected_diagnostic_seeds + ) + ): + _fail("spectral/layers.csv diagnostic seed binding is invalid") + spectral_model_hash_by_step: dict[int, str] = {} + for step_value, group in layers.groupby("step"): + hashes = group["model_state_sha256"].astype(str).unique().tolist() + if ( + len(hashes) != 1 + or len(hashes[0]) != 64 + or any(character not in "0123456789abcdef" for character in hashes[0].lower()) + ): + _fail("spectral/layers.csv has an invalid model-state hash") + spectral_model_hash_by_step[int(step_value)] = hashes[0] + + raw_path = ( + root + / "spectral" + / "raw" + / f"weightwatcher_step_{int(step_value):07d}.csv" + ) + status_path = ( + root / "spectral" / f"status_step_{int(step_value):07d}.json" + ) + if not raw_path.is_file() or not status_path.is_file(): + _fail( + f"spectral step {int(step_value)} lacks raw CSV/integrity status" + ) + status = _read_json(status_path) + expected_diagnostic_seed = int(seed) + 1_000_003 + int(step_value) + if ( + status.get("completed") is not True + or str(status.get("raw_csv_sha256", "")) != _file_sha256(raw_path) + or str(status.get("protocol_fingerprint", "")) != fingerprint + or str(status.get("model_state_sha256", "")) != hashes[0] + or _as_int(status.get("run_seed"), "spectral status seed") != seed + or _as_int( + status.get("diagnostic_seed"), + "spectral status diagnostic seed", + ) + != expected_diagnostic_seed + ): + _fail(f"spectral step {int(step_value)} integrity status is invalid") + raw = _read_csv(raw_path) + if ( + "matrix_name" not in raw.columns + or len(raw) != 6 + or raw["matrix_name"].astype(str).nunique() != 6 + ): + _fail(f"spectral raw step {int(step_value)} lacks six matrices") + for identity, expected in ( + ("protocol_fingerprint", fingerprint), + ("model_state_sha256", hashes[0]), + ("run_seed", seed), + ("diagnostic_seed", expected_diagnostic_seed), + ): + if identity not in raw.columns or not raw[identity].astype(str).eq( + str(expected) + ).all(): + _fail( + f"spectral raw step {int(step_value)} has invalid {identity}" + ) + raw_sorted = raw.sort_values("matrix_name").reset_index(drop=True) + layer_sorted = group.sort_values("matrix_name").reset_index(drop=True) + if not raw_sorted["matrix_name"].astype(str).equals( + layer_sorted["matrix_name"].astype(str) + ): + _fail( + f"spectral raw/layer matrix identity differs at step {int(step_value)}" + ) + for metric in ("alpha", "ERG_gap", "num_traps", "rand_distance"): + if metric not in raw_sorted.columns or metric not in layer_sorted.columns: + _fail(f"spectral raw/layer data lacks {metric}") + raw_values = pd.to_numeric(raw_sorted[metric], errors="coerce") + layer_values = pd.to_numeric(layer_sorted[metric], errors="coerce") + if not np.allclose( + raw_values.to_numpy(dtype=float), + layer_values.to_numpy(dtype=float), + rtol=1e-12, + atol=1e-12, + equal_nan=True, + ): + _fail( + f"spectral raw/layer {metric} differs at step {int(step_value)}" + ) + if clip_xmax: + numeric_columns = ( + "raw_alpha", + "alpha_raw", + "alpha_clip_xmax", + "alpha_delta", + "num_fingers", + "weightwatcher_analysis_calls", + ) + numeric = layers[list(numeric_columns)].apply( + pd.to_numeric, errors="coerce" + ) + if not np.isfinite(numeric.to_numpy(dtype=float)).all(): + _fail("clip-Xmax/raw-alpha spectral values are non-finite") + alpha = pd.to_numeric(layers["alpha"], errors="coerce").to_numpy( + dtype=float + ) + if not np.allclose( + alpha, + numeric["alpha_clip_xmax"].to_numpy(dtype=float), + rtol=0.0, + atol=0.0, + ): + _fail("alpha and alpha_clip_xmax aliases disagree") + if not np.allclose( + numeric["raw_alpha"].to_numpy(dtype=float), + numeric["alpha_raw"].to_numpy(dtype=float), + rtol=0.0, + atol=0.0, + ): + _fail("raw_alpha and alpha_raw aliases disagree") + if not np.allclose( + numeric["alpha_delta"].to_numpy(dtype=float), + numeric["alpha_raw"].to_numpy(dtype=float) + - numeric["alpha_clip_xmax"].to_numpy(dtype=float), + rtol=1e-12, + atol=1e-12, + ): + _fail("alpha_delta does not equal raw minus clip-Xmax alpha") + if not numeric["weightwatcher_analysis_calls"].eq(1).all(): + _fail("WeightWatcher was not called exactly once per checkpoint") + if not layers["finger_policy"].astype(str).eq( + "fix_fingers=clip_xmax" + ).all(): + _fail("spectral rows do not declare fix_fingers=clip_xmax") + if not layers["primary_alpha_variant"].astype(str).eq( + "clip_xmax" + ).all(): + _fail("spectral rows do not declare clipped alpha as primary") + if len(epoch_steps) < 10: + _fail("clip-Xmax campaign has fewer than ten permanent states") + for column in ( + "alpha_raw_n", + "alpha_raw_median", + "alpha_clip_xmax_n", + "alpha_clip_xmax_median", + ): + if column not in summary.columns: + _fail(f"spectral/summary.csv has no {column} column") + status_paths = [ + root / "spectral" / f"status_step_{step:07d}.json" + for step in epoch_steps + ] + missing_status = [str(path) for path in status_paths if not path.is_file()] + if missing_status: + _fail( + "missing WeightWatcher completion records: " + + ", ".join(missing_status) + ) + for path in status_paths: + status = _read_json(path) + if status.get("completed") is not True: + _fail(f"WeightWatcher status is incomplete: {path}") + _expect( + _as_int( + status.get("weightwatcher_analysis_calls"), + f"{path.name} analysis call count", + ), + 1, + f"{path.name} analysis call count", + ) if "n_matrices" not in summary.columns: _fail("spectral/summary.csv has no n_matrices column") matrix_counts = pd.to_numeric(summary["n_matrices"], errors="coerce") @@ -263,7 +789,15 @@ def validate_completed_run( f"missing epoch checkpoint recorded by " f"epoch_metrics.csv: {recorded}" ) - resolved_checkpoint_paths.append(candidate.resolve()) + resolved = candidate.resolve() + try: + resolved.relative_to(root.resolve()) + except ValueError: + _fail( + "epoch_metrics.csv references a checkpoint outside the run " + f"directory: {recorded}" + ) + resolved_checkpoint_paths.append(resolved) if len(resolved_checkpoint_paths) != len( set(resolved_checkpoint_paths) ): @@ -272,37 +806,79 @@ def validate_completed_run( "checkpoints" ) + if optimizer == "muon_clip": + _validate_muonclip_qk( + root, + manifest=manifest, + total_steps=total_steps, + ) + if verify_checkpoints: try: best_loss = float(completion.get("best_validation_loss")) except (TypeError, ValueError): _fail("run_complete.json has invalid best_validation_loss") + initial_inventory: dict[str, tuple[tuple[int, ...], str]] | None = None for filename, expected_step in ( + ("checkpoint_initial.pt", 0), ("checkpoint_latest.pt", total_steps), ("checkpoint_final.pt", total_steps), ("checkpoint_best.pt", best_step), ): payload = _load_checkpoint(root / filename) - _expect( - str(payload.get("fingerprint", "")), - fingerprint, - f"{filename} fingerprint", + _validate_checkpoint_identity( + payload, + path=root / filename, + fingerprint=fingerprint, + optimizer=optimizer, + seed=seed, + step=expected_step, + schema_version=5, ) - _expect( - str(payload.get("optimizer_name", "")), - optimizer, - f"{filename} optimizer", + computed_model_hash = model_state_sha256(payload["model"]) + if str(payload.get("model_state_sha256", "")) != computed_model_hash: + _fail(f"{filename} model-state SHA-256 does not match") + optimizer_states = payload.get("optimizers") + expected_optimizer_count = ( + 1 if optimizer in {"adam", "adamw"} else 2 ) - _expect( - _as_int(payload.get("seed"), f"{filename} seed"), - seed, - f"{filename} seed", - ) - _expect( - _as_int(payload.get("step"), f"{filename} step"), - expected_step, - f"{filename} step", + if ( + not isinstance(optimizer_states, list) + or len(optimizer_states) != expected_optimizer_count + or not all(isinstance(state, dict) for state in optimizer_states) + ): + _fail(f"{filename} optimizer-state inventory is invalid") + try: + computed_optimizer_hash = optimizer_state_sha256( + optimizer_states + ) + except (TypeError, ValueError) as exc: + _fail(f"{filename} optimizer-state hashing failed: {exc}") + if str(payload.get("optimizer_state_sha256", "")) != ( + computed_optimizer_hash + ): + _fail(f"{filename} optimizer-state SHA-256 does not match") + try: + require_finite_checkpoint_state( + model_state=payload["model"], + optimizer_states=optimizer_states, + step=expected_step, + ) + except FloatingPointError as exc: + _fail(f"{filename} contains non-finite checkpoint state: {exc}") + observed_inventory = _validate_model_tensor_inventory( + payload["model"], + path=root / filename, + expected=initial_inventory, ) + if filename == "checkpoint_initial.pt": + initial_inventory = observed_inventory + if model_state_sha256(payload["model"]) != initial_model_hash: + _fail( + "checkpoint_initial.pt model tensors do not match the " + "manifest initial-model hash" + ) + continue _expect( _as_int( payload.get("best_validation_step"), @@ -320,4 +896,77 @@ def validate_completed_run( ): _fail(f"{filename} best_validation_loss does not match completion") + required_epoch_columns = {"step", "epoch", "nominal_epoch"} + missing_epoch_columns = required_epoch_columns.difference( + epoch_metrics.columns + ) + if missing_epoch_columns: + _fail( + "epoch_metrics.csv lacks checkpoint identity columns: " + + ", ".join(sorted(missing_epoch_columns)) + ) + for (_, row), path in zip( + epoch_metrics.iterrows(), + resolved_checkpoint_paths, + strict=True, + ): + expected_step = _as_int(row["step"], "epoch checkpoint step") + nominal_epoch = _as_finite_float( + row["nominal_epoch"], "epoch checkpoint nominal_epoch" + ) + actual_epoch = _as_finite_float( + row["epoch"], "epoch checkpoint actual_epoch" + ) + payload = _load_checkpoint(path) + _validate_checkpoint_identity( + payload, + path=path, + fingerprint=fingerprint, + optimizer=optimizer, + seed=seed, + step=expected_step, + schema_version=3, + ) + if str(payload.get("model_state_sha256", "")) != ( + model_state_sha256(payload["model"]) + ): + _fail(f"{path.name} model-state SHA-256 does not match") + if str(payload.get("model_state_sha256", "")) != ( + spectral_model_hash_by_step.get(expected_step) + ): + _fail( + f"{path.name} model tensors do not match WeightWatcher state" + ) + if initial_inventory is None: # pragma: no cover - initial loads first + _fail("initial model tensor inventory is unavailable") + _validate_model_tensor_inventory( + payload["model"], + path=path, + expected=initial_inventory, + ) + if payload.get("purpose") != ( + "per_epoch_model_only_analysis_checkpoint" + ): + _fail(f"{path.name} has the wrong checkpoint purpose") + stored_nominal = _as_finite_float( + payload.get("nominal_epoch"), f"{path.name} nominal_epoch" + ) + stored_actual = _as_finite_float( + payload.get("actual_epoch"), f"{path.name} actual_epoch" + ) + if not math.isclose( + stored_nominal, + nominal_epoch, + rel_tol=0.0, + abs_tol=1e-12, + ): + _fail(f"{path.name} nominal_epoch does not match epoch_metrics.csv") + if not math.isclose( + stored_actual, + actual_epoch, + rel_tol=0.0, + abs_tol=1e-12, + ): + _fail(f"{path.name} actual_epoch does not match epoch_metrics.csv") + return completion diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py index d5841e0d..29f9643b 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py @@ -10,10 +10,18 @@ import yaml +from .provenance import ( + scientific_dependency_versions, + source_fingerprint_payload, +) from .runtime import is_tpu_environment BASELINE_OPTIMIZERS = ("sgd_momentum", "adamw", "muon") -SUPPORTED_OPTIMIZERS = BASELINE_OPTIMIZERS +# ``adam`` is optional in the historical reference YAMLs, but is a first-class +# arm in the dated long-horizon campaign. Keeping it out of +# BASELINE_OPTIMIZERS preserves compatibility with old configs while exposing +# it whenever a profile is declared. +SUPPORTED_OPTIMIZERS = ("sgd_momentum", "adam", "adamw", "muon") DEFAULT_ROOT = Path("/tmp/rg-nanogpt-one-head") TPU_ROOT_ENV = "RG_NANOGPT_ONE_HEAD_TPU_ROOT" TPU_PERSISTENT_ENV = "RG_TPU_PERSISTENT_ROOT" @@ -301,11 +309,42 @@ def validate_config(cfg: dict[str, Any]) -> None: raise ValueError("WeightWatcher randomize must be enabled") if int(ww["min_evals"]) < 5: raise ValueError("weightwatcher.min_evals must be at least 5") + finger_policy = ww.get("fix_fingers", False) + if finger_policy not in (False, "clip_xmax"): + raise ValueError( + "weightwatcher.fix_fingers must be false or 'clip_xmax'" + ) + if finger_policy == "clip_xmax": + if int(ww.get("max_fingers", 0)) < 1: + raise ValueError( + "weightwatcher.max_fingers must be positive when " + "fix_fingers='clip_xmax'" + ) + if not bool(ww.get("require_raw_alpha", True)): + raise ValueError( + "clip_xmax monitoring must retain WeightWatcher's raw_alpha" + ) + + runtime = cfg["runtime"] + if str(runtime.get("matmul_precision", "high")) not in { + "highest", + "high", + "medium", + }: + raise ValueError( + "runtime.matmul_precision must be highest, high, or medium" + ) + if bool(runtime.get("allow_tf32", False)) and str( + runtime.get("matmul_precision", "high") + ) == "highest": + raise ValueError( + "runtime.allow_tf32 cannot be true when matmul_precision=highest" + ) def validate_optimizer_profile(profile: dict[str, Any]) -> None: family = str(profile.get("family", "")) - if family not in {"sgd", "adamw", "muon"}: + if family not in {"sgd", "adam", "adamw", "muon"}: raise ValueError(f"unsupported optimizer family: {family}") warmup_fraction = float(profile.get("warmup_fraction", -1.0)) if not 0.0 <= warmup_fraction < 1.0: @@ -315,7 +354,7 @@ def validate_optimizer_profile(profile: dict[str, Any]) -> None: if "lr_schedule_epochs" in profile and float(profile["lr_schedule_epochs"]) <= 0: raise ValueError("lr_schedule_epochs must be positive") - if family in {"sgd", "adamw"}: + if family in {"sgd", "adam", "adamw"}: peak = float(profile["learning_rate"]) floor = float(profile["min_learning_rate"]) if peak <= 0 or floor < 0 or floor > peak: @@ -441,6 +480,9 @@ def protocol_fingerprint( "optimizer_profile": optimizer_profile(cfg, optimizer), "evaluation": cfg["evaluation"], "weightwatcher": cfg["weightwatcher"], + "runtime": cfg["runtime"], + "source": source_fingerprint_payload(), + "scientific_dependencies": scientific_dependency_versions(), "optimizer": str(optimizer), "seed": int(seed), "data_metadata": data_metadata, diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/data.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/data.py index 63015c48..8c7a6b68 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/data.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/data.py @@ -209,6 +209,8 @@ def validate_prepared_data( ) metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + if int(metadata.get("schema_version", -1)) != 2: + raise RuntimeError("prepared corpus metadata schema must be version 2") expected_splits = { "train": int(cfg["dataset"]["train_tokens"]), "val": int(cfg["dataset"]["val_tokens"]), @@ -225,8 +227,21 @@ def validate_prepared_data( raise RuntimeError("prepared dataset configuration does not match config") if metadata.get("dataset_revision") != cfg["dataset"]["revision"]: raise RuntimeError("prepared dataset revision does not match config") - if metadata.get("tokenizer") != "gpt2": - raise RuntimeError("prepared tokenizer must be GPT-2 BPE") + if metadata.get("dataset_split") != cfg["dataset"].get("split", "train"): + raise RuntimeError("prepared dataset split does not match config") + expected_tokenizer = str(cfg["dataset"].get("tokenizer", "gpt2")) + if metadata.get("tokenizer") != expected_tokenizer: + raise RuntimeError("prepared tokenizer does not match config") + if expected_tokenizer != "gpt2": + raise RuntimeError("this baseline requires the GPT-2 BPE tokenizer") + if int(metadata.get("vocab_size", -1)) != int(cfg["model"]["vocab_size"]): + raise RuntimeError("prepared tokenizer vocabulary does not match model") + eot_token = int(metadata.get("eot_token", -1)) + if int(cfg["model"]["vocab_size"]) == 50_257: + if eot_token != 50_256: + raise RuntimeError("prepared GPT-2 end-of-text token must be 50256") + elif not 0 <= eot_token < int(cfg["model"]["vocab_size"]): + raise RuntimeError("prepared end-of-text token is outside the vocabulary") if metadata.get("dtype") != TOKEN_DTYPE.name: raise RuntimeError("prepared token dtype must be uint16") if metadata.get("document_disjoint_splits") is not True: diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/doctor_smoke.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/doctor_smoke.py new file mode 100644 index 00000000..3b4e1fac --- /dev/null +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/doctor_smoke.py @@ -0,0 +1,554 @@ +from __future__ import annotations + +"""Executable accelerator/optimizer/checkpoint/WeightWatcher smoke test. + +This module deliberately uses a tiny model and synthetic tokens. It is a +backend gate, not an experiment run, and writes only beneath explicitly +provided temporary paths. +""" + +import argparse +import csv +from copy import deepcopy +from datetime import datetime, timezone +import json +import math +import os +from pathlib import Path +import tempfile +import time +from typing import Any, Sequence + + +_OPTIMIZERS = ("adamw", "muon_clip") +_CACHE_ENVIRONMENTS = { + "HOME": "home", + "HF_HOME": "huggingface", + "HF_DATASETS_CACHE": "huggingface/datasets", + "HUGGINGFACE_HUB_CACHE": "huggingface/hub", + "HF_HUB_CACHE": "huggingface/hub", + "HF_ASSETS_CACHE": "huggingface/assets", + "HF_MODULES_CACHE": "huggingface/modules", + "TRANSFORMERS_CACHE": "huggingface/transformers", + "TIKTOKEN_CACHE_DIR": "tiktoken", + "MPLCONFIGDIR": "matplotlib", + "XDG_CACHE_HOME": "xdg/cache", + "XDG_CONFIG_HOME": "xdg/config", + "XDG_DATA_HOME": "xdg/data", + "XDG_STATE_HOME": "xdg/state", + "TORCH_HOME": "torch", + "TORCH_EXTENSIONS_DIR": "torch_extensions", + "TORCHINDUCTOR_CACHE_DIR": "torchinductor", + "CUDA_CACHE_PATH": "cuda", + "TRITON_CACHE_DIR": "triton", + "CUPY_CACHE_DIR": "cupy", + "XLA_PERSISTENT_CACHE_PATH": "xla", + "NUMBA_CACHE_DIR": "numba", + "JOBLIB_TEMP_FOLDER": "joblib", + "PYTHONPYCACHEPREFIX": "pycache", +} + + +def _is_below_temporary_root(path: Path) -> bool: + resolved = path.resolve(strict=False) + candidates = { + Path("/tmp").resolve(strict=False), + Path("/private/tmp").resolve(strict=False), + Path(tempfile.gettempdir()).resolve(strict=False), + } + for candidate in candidates: + try: + resolved.relative_to(candidate) + return resolved != candidate + except ValueError: + continue + return False + + +def _require_temporary_path(value: str | Path, *, label: str) -> Path: + path = Path(value).expanduser().resolve(strict=False) + if not _is_below_temporary_root(path): + raise ValueError( + f"{label} must be a descendant of /tmp (or macOS /private/tmp): " + f"{path}" + ) + return path + + +def _configure_private_caches(session_dir: Path) -> dict[str, str]: + cache_root = session_dir / "cache" + configured: dict[str, str] = {} + for name, relative in _CACHE_ENVIRONMENTS.items(): + path = cache_root / relative + path.mkdir(parents=True, exist_ok=True) + os.environ[name] = str(path) + configured[name] = str(path) + temporary = session_dir / "tmp" + temporary.mkdir(parents=True, exist_ok=True) + for name in ("TMPDIR", "TMP", "TEMP"): + os.environ[name] = str(temporary) + configured[name] = str(temporary) + os.environ["MPLBACKEND"] = "Agg" + configured["MPLBACKEND"] = "Agg" + return configured + + +def _atomic_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + with temporary.open("w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, sort_keys=True, allow_nan=False) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + temporary.replace(path) + + +def _finite_float(value: Any, *, label: str) -> float: + result = float(value) + if not math.isfinite(result): + raise FloatingPointError(f"{label} is not finite: {result}") + return result + + +def _state_is_finite(state: dict[str, Any]) -> bool: + import torch + + for value in state.values(): + if torch.is_tensor(value) and ( + value.is_floating_point() or value.is_complex() + ): + if not bool(torch.isfinite(value).all()): + return False + return True + + +def _gradient_norm(model: Any) -> float: + import torch + + squares = [ + parameter.grad.detach().float().square().sum() + for parameter in model.parameters() + if parameter.grad is not None + ] + if not squares: + raise RuntimeError("the smoke backward pass produced no gradients") + total = torch.stack(squares).sum().sqrt() + return _finite_float(total.detach().cpu(), label="gradient norm") + + +def _snapshot_equal(expected: Sequence[Any], observed: Sequence[Any]) -> bool: + import torch + + return len(expected) == len(observed) and all( + torch.equal(left, right) + for left, right in zip(expected, observed, strict=True) + ) + + +def _optimizer_smoke( + *, + optimizer_name: str, + cfg: dict[str, Any], + device: Any, + session_dir: Path, + seed: int, +) -> tuple[dict[str, Any], Any]: + import torch + + from rg_nanogpt_one_head.checkpoints import ( + load_training_checkpoint_for_resume, + save_training_checkpoint, + ) + from rg_nanogpt_one_head.model import GPT, GPTConfig + from rg_nanogpt_one_head.optimizers import ( + make_optimizer_handles, + optimizer_step, + zero_grad, + ) + from rg_nanogpt_one_head.run_utils import model_state_sha256 + from rg_nanogpt_one_head.runtime import ( + mark_step, + seed_everything, + synchronize, + ) + + tiny_config = GPTConfig( + vocab_size=128, + block_size=8, + n_layer=1, + n_head=1, + n_embd=32, + dropout=0.0, + bias=False, + tie_weights=True, + ) + seed_everything(seed, device) + model = GPT(tiny_config).to(device) + profile = deepcopy(cfg["optimizer_profiles"][optimizer_name]) + handles = make_optimizer_handles(model, profile) + + token_generator = torch.Generator(device="cpu").manual_seed(seed + 17) + inputs = torch.randint( + 0, + tiny_config.vocab_size, + (2, tiny_config.block_size), + generator=token_generator, + dtype=torch.long, + ).to(device) + targets = torch.randint( + 0, + tiny_config.vocab_size, + (2, tiny_config.block_size), + generator=token_generator, + dtype=torch.long, + ).to(device) + + model.train() + zero_grad(handles) + _, loss = model(inputs, targets) + if loss is None: + raise RuntimeError("the smoke forward pass did not return a loss") + loss.backward() + gradient_norm = _gradient_norm(model) + optimizer_step(handles) + mark_step(device) + synchronize(device) + loss_value = _finite_float(loss.detach().cpu(), label=f"{optimizer_name} loss") + + state = { + name: value.detach().cpu().contiguous() + for name, value in model.state_dict().items() + } + if not _state_is_finite(state): + raise FloatingPointError( + f"{optimizer_name} produced non-finite model parameters" + ) + state_hash = model_state_sha256(state) + previous_eval_snapshot = [ + parameter.detach().cpu().clone() for parameter in model.parameters() + ] + diagnostics = { + "previous_eval_snapshot": previous_eval_snapshot, + "last_grad_pre": gradient_norm, + "last_grad_post": gradient_norm, + "last_clipped": False, + } + train_generator = torch.Generator(device="cpu").manual_seed(seed + 29) + # Advance the stream so this verifies restoration of non-initial state. + torch.randint(0, 2**16, (19,), generator=train_generator) + expected_generator_state = train_generator.get_state().clone() + checkpoint_path = session_dir / "checkpoints" / optimizer_name / "step_1.pt" + fingerprint = f"doctor-smoke-v1:{optimizer_name}:{seed}" + save_training_checkpoint( + checkpoint_path, + model=model, + handles=handles, + step=1, + best_validation_loss=loss_value, + best_validation_step=1, + elapsed_seconds=0.0, + fingerprint=fingerprint, + cfg=cfg, + optimizer_name=optimizer_name, + seed=seed, + train_generator=train_generator, + resume_diagnostics=diagnostics, + ) + raw_checkpoint = torch.load( + checkpoint_path, + map_location="cpu", + weights_only=False, + ) + if int(raw_checkpoint.get("schema_version", -1)) != 5: + raise RuntimeError( + f"{optimizer_name} smoke checkpoint is not schema version 5" + ) + if len(raw_checkpoint.get("optimizers", ())) != len(handles): + raise RuntimeError( + f"{optimizer_name} smoke checkpoint optimizer inventory changed" + ) + + seed_everything(seed + 1, device) + resumed_model = GPT(tiny_config).to(device) + resumed_handles = make_optimizer_handles(resumed_model, profile) + resumed_generator = torch.Generator(device="cpu").manual_seed(seed + 1) + loaded = load_training_checkpoint_for_resume( + checkpoint_path, + model=resumed_model, + handles=resumed_handles, + expected_fingerprint=fingerprint, + train_generator=resumed_generator, + ) + step, best_loss, best_step, elapsed, loaded_diagnostics = loaded + if (step, best_step, elapsed) != (1, 1, 0.0): + raise RuntimeError( + f"{optimizer_name} checkpoint metadata changed during round-trip" + ) + if best_loss != loss_value: + raise RuntimeError( + f"{optimizer_name} validation loss changed during round-trip" + ) + if loaded_diagnostics is None: + raise RuntimeError( + f"{optimizer_name} resume diagnostics were not restored" + ) + if not _snapshot_equal( + previous_eval_snapshot, + loaded_diagnostics["previous_eval_snapshot"], + ): + raise RuntimeError( + f"{optimizer_name} resume snapshot changed during round-trip" + ) + if ( + loaded_diagnostics["last_grad_pre"] != gradient_norm + or loaded_diagnostics["last_grad_post"] != gradient_norm + or loaded_diagnostics["last_clipped"] is not False + ): + raise RuntimeError( + f"{optimizer_name} scalar resume diagnostics changed" + ) + if not torch.equal(resumed_generator.get_state(), expected_generator_state): + raise RuntimeError( + f"{optimizer_name} train generator state was not restored" + ) + resumed_state = { + name: value.detach().cpu().contiguous() + for name, value in resumed_model.state_dict().items() + } + resumed_hash = model_state_sha256(resumed_state) + if resumed_hash != state_hash: + raise RuntimeError( + f"{optimizer_name} model state changed during checkpoint round-trip" + ) + + # Optimizer state device/dtype mistakes (especially on XLA and MPS) can + # remain latent until the first update after load_state_dict. Exercise that + # boundary before the backend is approved for a multi-day run. + resumed_model.train() + zero_grad(resumed_handles) + _, resumed_loss = resumed_model(inputs, targets) + if resumed_loss is None: + raise RuntimeError( + f"{optimizer_name} resumed forward pass returned no loss" + ) + resumed_loss.backward() + resumed_gradient_norm = _gradient_norm(resumed_model) + optimizer_step(resumed_handles) + mark_step(device) + synchronize(device) + resumed_loss_value = _finite_float( + resumed_loss.detach().cpu(), + label=f"{optimizer_name} resumed loss", + ) + resumed_updated_state = { + name: value.detach().cpu().contiguous() + for name, value in resumed_model.state_dict().items() + } + if not _state_is_finite(resumed_updated_state): + raise FloatingPointError( + f"{optimizer_name} produced non-finite state after resume" + ) + + return ( + { + "optimizer": optimizer_name, + "loss": loss_value, + "gradient_norm": gradient_norm, + "optimizer_handles": len(handles), + "checkpoint_schema_version": 5, + "checkpoint_path": str(checkpoint_path), + "model_state_sha256": state_hash, + "checkpoint_roundtrip": True, + "resume_diagnostics_roundtrip": True, + "train_generator_roundtrip": True, + "resumed_optimizer_step": True, + "resumed_loss": resumed_loss_value, + "resumed_gradient_norm": resumed_gradient_norm, + }, + resumed_model, + ) + + +def _weightwatcher_smoke( + *, + model: Any, + cfg: dict[str, Any], + session_dir: Path, + seed: int, +) -> dict[str, Any]: + from rg_nanogpt_one_head.model import transformer_matrix_items + from rg_nanogpt_one_head.spectral import run_weightwatcher + + ww_config = deepcopy(cfg["weightwatcher"]) + if ww_config.get("fix_fingers") != "clip_xmax": + raise RuntimeError( + "the backend smoke requires weightwatcher.fix_fingers=clip_xmax" + ) + expected_names = { + name for name, _, _, _ in transformer_matrix_items(model) + } + if len(expected_names) != 6: + raise RuntimeError("the tiny model does not expose exactly six matrices") + + run_dir = session_dir / "weightwatcher" + summary = run_weightwatcher( + model, + run_dir, + step=1, + tokens_seen=16, + train_tokens=int(cfg["dataset"]["train_tokens"]), + config=ww_config, + seed=seed, + fingerprint=f"doctor-smoke-v1:weightwatcher:{seed}", + ) + raw_path = run_dir / "spectral" / "raw" / "weightwatcher_step_0000001.csv" + with raw_path.open("r", newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + observed_names = {str(row["matrix_name"]) for row in rows} + if len(rows) != 6 or observed_names != expected_names: + raise RuntimeError( + "WeightWatcher did not return exactly the six transformer matrices" + ) + for index, row in enumerate(rows): + if int(float(row["weightwatcher_analysis_calls"])) != 1: + raise RuntimeError( + "WeightWatcher smoke result does not prove a one-pass analysis" + ) + _finite_float(row["raw_alpha"], label=f"raw_alpha row {index}") + _finite_float(row["alpha_raw"], label=f"alpha_raw row {index}") + _finite_float(row["alpha"], label=f"alpha row {index}") + _finite_float( + row["alpha_clip_xmax"], + label=f"alpha_clip_xmax row {index}", + ) + if row["finger_policy"] != "fix_fingers=clip_xmax": + raise RuntimeError("WeightWatcher finger policy metadata changed") + + return { + "analysis_calls": 1, + "matrix_count": 6, + "matrix_names": sorted(observed_names), + "finger_policy": "fix_fingers=clip_xmax", + "raw_alpha_count": int(summary["alpha_raw_n"]), + "clipped_alpha_count": int(summary["alpha_clip_xmax_n"]), + "raw_csv": str(raw_path), + } + + +def run_smoke( + *, + config_path: str | Path, + work_dir: str | Path, + summary_path: str | Path, + device_request: str, + seed: int, +) -> dict[str, Any]: + work_root = _require_temporary_path(work_dir, label="work directory") + output_path = _require_temporary_path(summary_path, label="summary path") + work_root.mkdir(parents=True, exist_ok=True) + session_dir = work_root / ( + f"doctor-smoke-{time.time_ns()}-pid-{os.getpid()}" + ) + session_dir.mkdir(parents=False, exist_ok=False) + cache_environment = _configure_private_caches(session_dir) + + # MuonClip extends both profile validation and optimizer construction, so + # it must be installed before the frozen YAML is loaded and validated. + from rg_nanogpt_one_head.muonclip import install_muonclip_extension + + install_muonclip_extension() + + from rg_nanogpt_one_head.config import load_config + from rg_nanogpt_one_head.runtime import ( + accelerator_name, + choose_device, + configure_runtime, + runtime_metadata, + ) + + config = Path(config_path).expanduser().resolve(strict=True) + cfg = load_config(config) + device = choose_device(device_request) + configure_runtime(device, cfg) + + optimizer_results: list[dict[str, Any]] = [] + spectral_model = None + for optimizer_name in _OPTIMIZERS: + result, spectral_model = _optimizer_smoke( + optimizer_name=optimizer_name, + cfg=cfg, + device=device, + session_dir=session_dir, + seed=int(seed), + ) + optimizer_results.append(result) + if spectral_model is None: # pragma: no cover - fixed non-empty inventory + raise AssertionError("no optimizer smoke model was produced") + weightwatcher = _weightwatcher_smoke( + model=spectral_model, + cfg=cfg, + session_dir=session_dir, + seed=int(seed), + ) + + payload = { + "schema_version": 1, + "completed": True, + "completed_at_utc": datetime.now(timezone.utc).isoformat(), + "config": str(config), + "device_request": str(device_request), + "resolved_device": str(device), + "accelerator": accelerator_name(device), + "seed": int(seed), + "session_dir": str(session_dir), + "cache_environment": cache_environment, + "runtime": runtime_metadata(device), + "tiny_model": { + "vocab_size": 128, + "block_size": 8, + "n_layer": 1, + "n_head": 1, + "n_embd": 32, + "batch_size": 2, + }, + "optimizers": optimizer_results, + "weightwatcher": weightwatcher, + } + _atomic_json(output_path, payload) + return payload + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Exercise the selected backend with all campaign optimizers, " + "schema-v5 checkpoint resume, and one-pass WeightWatcher." + ) + ) + parser.add_argument("--config", required=True, type=Path) + parser.add_argument("--work-dir", required=True, type=Path) + parser.add_argument("--summary", required=True, type=Path) + parser.add_argument( + "--device", + default="auto", + choices=("auto", "cpu", "cuda", "mps", "tpu", "xla"), + ) + parser.add_argument("--seed", type=int, default=24_681_357) + return parser + + +def main(argv: Sequence[str] | None = None) -> None: + args = _parser().parse_args(argv) + payload = run_smoke( + config_path=args.config, + work_dir=args.work_dir, + summary_path=args.summary, + device_request=args.device, + seed=args.seed, + ) + print(json.dumps(payload, indent=2, sort_keys=True, allow_nan=False)) + + +if __name__ == "__main__": + main() diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py index de679ee5..51d57f51 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py @@ -1,6 +1,7 @@ from __future__ import annotations import csv +from datetime import datetime, timezone import json from pathlib import Path import shutil @@ -9,7 +10,7 @@ from .completion import validate_completed_run from .checkpoints import ( - load_training_checkpoint, + load_training_checkpoint_for_resume, save_training_checkpoint, ) from .config import ( @@ -32,7 +33,9 @@ checkpoint_eval, prepare_csv, run_directory, + truncate_muonclip_qk_after, truncate_spectral_after, + validate_existing_manifest_runtime, write_manifest, ) from .runtime import ( @@ -76,6 +79,13 @@ def run_one( ) run_dir.mkdir(parents=True, exist_ok=True) + # Runtime identity is part of the run identity. Resolve and validate it + # before a completed run is reused or a partial run is loaded/truncated. + resolved_device = choose_device(device) + configure_runtime(resolved_device, cfg) + current_runtime = runtime_metadata(resolved_device) + validate_existing_manifest_runtime(run_dir, current_runtime) + data_metadata, arrays = load_memmaps(data_root, cfg) train_tokens = int(data_metadata["splits"]["train"]) total_steps = max_steps(cfg, train_tokens) @@ -104,11 +114,8 @@ def run_one( ) return run_dir - resolved_device = choose_device(device) - configure_runtime(resolved_device, cfg) seed_everything(int(seed), resolved_device) if progress: - metadata = runtime_metadata(resolved_device) print( "[one-head-env] " f"requested={device} " @@ -116,7 +123,7 @@ def run_one( f"device={resolved_device} " f"data_root={data_root} " f"results_root={results_root} " - f"runtime={json.dumps(metadata, sort_keys=True)}", + f"runtime={json.dumps(current_runtime, sort_keys=True)}", flush=True, ) @@ -168,30 +175,47 @@ def run_one( best_validation_loss = float("inf") best_validation_step = 0 elapsed_offset = 0.0 + resume_diagnostics = None + resumed_from_checkpoint = False initial_checkpoint = run_dir / "checkpoint_initial.pt" latest_checkpoint = run_dir / "checkpoint_latest.pt" best_checkpoint = run_dir / "checkpoint_best.pt" final_checkpoint = run_dir / "checkpoint_final.pt" - if resume and latest_checkpoint.is_file(): + resume_checkpoint = ( + latest_checkpoint + if latest_checkpoint.is_file() + else initial_checkpoint + if initial_checkpoint.is_file() + else None + ) + if resume and resume_checkpoint is not None: ( start_step, best_validation_loss, best_validation_step, elapsed_offset, - ) = load_training_checkpoint( - latest_checkpoint, + resume_diagnostics, + ) = load_training_checkpoint_for_resume( + resume_checkpoint, model=model, handles=handles, expected_fingerprint=fingerprint, train_generator=train_generator, ) + if resume_diagnostics is None and start_step > 0: + raise RuntimeError( + "checkpoint predates deterministic resume diagnostics; use a " + "new results directory or rerun with explicit overwrite" + ) + resumed_from_checkpoint = True model.to(resolved_device) synchronize(resolved_device) truncate_spectral_after(run_dir, start_step) + truncate_muonclip_qk_after(run_dir, start_step) if progress: print( f"[one-head-train] resume {optimizer_name} " - f"seed={seed} step={start_step}" + f"seed={seed} step={start_step} checkpoint={resume_checkpoint.name}" ) elif run_dir.exists() and any(run_dir.iterdir()) and resume: # Opt-in diagnostics may register an external append-only artifact @@ -267,12 +291,12 @@ def run_one( prepare_csv( metrics_path, METRIC_FIELDS, - start_step if start_step else None, + start_step if resumed_from_checkpoint else None, ) prepare_csv( epoch_metrics_path, EPOCH_FIELDS, - start_step if start_step else None, + start_step if resumed_from_checkpoint else None, ) with ( metrics_path.open( @@ -286,7 +310,12 @@ def run_one( encoding="utf-8", ) as epoch_handle, ): - best_validation_loss, best_validation_step, elapsed_total = ( + ( + best_validation_loss, + best_validation_step, + elapsed_total, + final_resume_diagnostics, + ) = ( execute_training_loop( cfg=cfg, model=model, @@ -294,8 +323,6 @@ def run_one( arrays=arrays, train_probe=train_probe, val_probe=val_probe, - test_probe=test_probe, - bleu_probe=bleu_probe, device=resolved_device, optimizer_name=optimizer_name, seed=int(seed), @@ -307,6 +334,7 @@ def run_one( best_validation_loss=best_validation_loss, best_validation_step=best_validation_step, elapsed_offset=elapsed_offset, + resume_diagnostics=resume_diagnostics, fingerprint=fingerprint, train_generator=train_generator, epoch_steps=epoch_steps, @@ -327,6 +355,15 @@ def run_one( ) ) + for handle in handles: + flush = getattr( + handle.optimizer, + "flush_pending_diagnostics", + None, + ) + if callable(flush): + flush() + for checkpoint in (final_checkpoint, latest_checkpoint): save_training_checkpoint( checkpoint, @@ -341,6 +378,7 @@ def run_one( optimizer_name=optimizer_name, seed=int(seed), train_generator=train_generator, + resume_diagnostics=final_resume_diagnostics, ) final_state = torch.load( @@ -370,8 +408,8 @@ def run_one( test_results = { "policy": ( - "test is monitoring-only; validation loss selects " - "checkpoint_best.pt" + "test is held out until post-training audit; validation loss " + "selects checkpoint_best.pt and test never tunes the protocol" ), "final": final_test, "validation_selected": best_test, @@ -386,6 +424,7 @@ def run_one( ) completion = { "completed": True, + "completed_at_utc": datetime.now(timezone.utc).isoformat(), "optimizer": optimizer_name, "seed": int(seed), "optimizer_steps": int(total_steps), @@ -397,8 +436,16 @@ def run_one( "best_validation_loss": float(best_validation_loss), "final_test_loss": float(final_test["loss"]), "final_test_perplexity": float(final_test["perplexity"]), + "final_test_bits_per_token": float(final_test["bits_per_token"]), "final_test_accuracy": float(final_test["accuracy"]), + "final_test_top5_accuracy": float(final_test["top5_accuracy"]), "final_test_bleu": float(final_test["bleu"]), + "final_test_continuation_token_accuracy": float( + final_test["continuation_token_accuracy"] + ), + "final_test_continuation_exact_match": float( + final_test["continuation_exact_match"] + ), "fingerprint": fingerprint, } temporary = run_dir / "run_complete.json.tmp" diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py index 8c51746c..a5adc95b 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py @@ -142,6 +142,7 @@ def evaluate_probe( # validation-checkpoint selection comparable to existing runs. losses: list[float] = [] correct = 0 + top5_correct = 0 total = 0 for x_cpu, y_cpu in probe: x = x_cpu.to(device) @@ -155,17 +156,31 @@ def evaluate_probe( correct += int( (logits.argmax(dim=-1) == y).sum().detach().cpu() ) + top_k = min(5, int(logits.shape[-1])) + top5_correct += int( + ( + logits.topk(top_k, dim=-1).indices + == y.unsqueeze(-1) + ) + .any(dim=-1) + .sum() + .detach() + .cpu() + ) total += int(y.numel()) model.train(was_training) mean_loss = float(np.mean(losses)) return { "loss": mean_loss, - "perplexity": float(math.exp(min(20.0, mean_loss))), + "perplexity": float(math.exp(mean_loss)), + "bits_per_token": float(mean_loss / math.log(2.0)), "accuracy": correct / max(1, total), + "top5_accuracy": top5_correct / max(1, total), } loss_sum = torch.zeros((), dtype=torch.float32, device=device) correct = torch.zeros((), dtype=torch.int64, device=device) + top5_correct = torch.zeros((), dtype=torch.int64, device=device) total = 0 batches = 0 for x_cpu, y_cpu in probe: @@ -176,6 +191,10 @@ def evaluate_probe( raise RuntimeError("evaluation forward pass did not return loss") loss_sum = loss_sum + loss.detach().float() correct = correct + (logits.argmax(dim=-1) == y).sum() + top_k = min(5, int(logits.shape[-1])) + top5_correct = top5_correct + ( + logits.topk(top_k, dim=-1).indices == y.unsqueeze(-1) + ).any(dim=-1).sum() total += int(y.numel()) batches += 1 # XLA is lazy. Execute each fixed-shape batch without transferring @@ -186,11 +205,14 @@ def evaluate_probe( synchronize(device) mean_loss = float((loss_sum / batches).detach().cpu()) correct_value = int(correct.detach().cpu()) + top5_correct_value = int(top5_correct.detach().cpu()) model.train(was_training) return { "loss": mean_loss, - "perplexity": float(math.exp(min(20.0, mean_loss))), + "perplexity": float(math.exp(mean_loss)), + "bits_per_token": float(mean_loss / math.log(2.0)), "accuracy": correct_value / max(1, total), + "top5_accuracy": top5_correct_value / max(1, total), } @@ -198,9 +220,10 @@ def _cpu_bleu_model(model) -> nn.Module: """Build a CPU copy for BLEU when the live model is on TPU/XLA. Greedy decoding changes sequence length at every token and would otherwise - trigger a series of XLA compilations. BLEU is monitoring-only, so the small - CPU copy avoids that accelerator-specific overhead without affecting - training, checkpoint selection, or WeightWatcher measurements. + trigger a series of XLA compilations. BLEU is a post-training secondary + audit, so the small CPU copy avoids that accelerator-specific overhead + without affecting training, checkpoint selection, or WeightWatcher + measurements. """ synchronize(model.lm_head.weight.device) @@ -221,9 +244,10 @@ def evaluate_bleu( This is not a translation benchmark. It measures exact lexical overlap between deterministic model continuations and the held-out continuation. - On TPU/XLA, decoding is intentionally performed on a CPU snapshot because - it is monitoring-only and its changing sequence lengths are a poor fit for - repeated XLA compilation. + It is evaluated only after training for the final and validation-selected + checkpoints. On TPU/XLA, decoding is intentionally performed on a CPU + snapshot because changing sequence lengths are a poor fit for repeated XLA + compilation. """ try: @@ -246,6 +270,9 @@ def evaluate_bleu( encoder = tiktoken.get_encoding("gpt2") hypotheses: list[str] = [] references: list[str] = [] + continuation_correct = 0 + continuation_total = 0 + continuation_exact = 0 for start in range(0, len(probe.prompts), int(batch_size)): prompts = probe.prompts[ start : start + int(batch_size) @@ -261,6 +288,13 @@ def evaluate_bleu( reference_batch = probe.references[ start : start + int(batch_size) ] + continuation_correct += int( + (continuation == reference_batch).sum().item() + ) + continuation_total += int(reference_batch.numel()) + continuation_exact += int( + (continuation == reference_batch).all(dim=-1).sum().item() + ) for predicted_tokens, reference_tokens in zip( continuation, reference_batch, @@ -277,4 +311,10 @@ def evaluate_bleu( "bleu_examples": float(len(hypotheses)), "bleu_sys_len": float(score.sys_len), "bleu_ref_len": float(score.ref_len), + "continuation_token_accuracy": ( + continuation_correct / max(1, continuation_total) + ), + "continuation_exact_match": ( + continuation_exact / max(1, len(hypotheses)) + ), } diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/monitor.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/monitor.py index 004b13aa..79ef9641 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/monitor.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/monitor.py @@ -21,6 +21,8 @@ "step", "epoch", "alpha", + "alpha_raw", + "alpha_clip_xmax", "D", "rand_distance", "ERG_gap", @@ -168,7 +170,11 @@ def format_monitor_snapshot( latest_step = int(layers["step"].max()) latest = layers[layers["step"] == latest_step].copy() latest_epoch = float(latest["epoch"].iloc[0]) - table = latest[list(_LAYER_COLUMNS)].sort_values("matrix_name") + table_columns = list(_LAYER_COLUMNS) + for column in ("alpha_raw", "alpha_clip_xmax", "num_fingers"): + if column in latest.columns: + table_columns.append(column) + table = latest[table_columns].sort_values("matrix_name") lines.extend( [ @@ -180,21 +186,39 @@ def format_monitor_snapshot( "", _format_table(table), "", - "ALPHA: " + _finite_summary(latest["alpha"]), + "ALPHA PRIMARY: " + _finite_summary(latest["alpha"]), "RAND_DISTANCE: " + _finite_summary(latest["rand_distance"]), ] ) + if "alpha_raw" in latest.columns: + lines.append( + "ALPHA RAW: " + _finite_summary(latest["alpha_raw"]) + ) + if "alpha_clip_xmax" in latest.columns: + lines.append( + "ALPHA CLIPPED: " + + _finite_summary(latest["alpha_clip_xmax"]) + ) + + aggregations = { + "alpha_median": ("alpha", "median"), + "rand_distance_median": ("rand_distance", "median"), + "D_median": ("D", "median"), + "ERG_gap_median": ("ERG_gap", "median"), + "num_traps_mean": ("num_traps", "mean"), + } + if "alpha_raw" in layers.columns: + aggregations["alpha_raw_median"] = ("alpha_raw", "median") + if "alpha_clip_xmax" in layers.columns: + aggregations["alpha_clip_xmax_median"] = ( + "alpha_clip_xmax", + "median", + ) recent_frame = ( layers.groupby(["step", "epoch"], as_index=False) - .agg( - alpha_median=("alpha", "median"), - rand_distance_median=("rand_distance", "median"), - D_median=("D", "median"), - ERG_gap_median=("ERG_gap", "median"), - num_traps_mean=("num_traps", "mean"), - ) + .agg(**aggregations) .sort_values("step") .tail(max(1, int(recent))) ) diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py index df7b6403..83278195 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py @@ -11,6 +11,7 @@ from copy import deepcopy import importlib import math +import os from pathlib import Path from typing import Any, Iterable @@ -281,12 +282,42 @@ def _write_diagnostics(self, values: dict[str, float]) -> None: "mean_gamma", "min_gamma", ] - write_header = not path.is_file() or path.stat().st_size == 0 - with path.open("a", newline="", encoding="utf-8") as handle: + rows: list[dict[str, Any]] = [] + if path.is_file() and path.stat().st_size: + with path.open("r", newline="", encoding="utf-8") as handle: + reader = csv.DictReader(handle) + if reader.fieldnames != fields: + raise RuntimeError( + "MuonClip QK diagnostic schema changed during a run" + ) + for row in reader: + if None in row or any(value is None for value in row.values()): + raise RuntimeError( + "MuonClip QK diagnostic row is incomplete" + ) + try: + step = int(float(row["step"])) + except (TypeError, ValueError) as exc: + raise RuntimeError( + "MuonClip QK diagnostic row has an invalid step" + ) from exc + if step != int(self.step_index): + rows.append(dict(row)) + rows.append({field: values[field] for field in fields}) + rows.sort(key=lambda row: int(float(row["step"]))) + temporary = path.with_suffix(path.suffix + ".tmp") + with temporary.open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter(handle, fieldnames=fields) - if write_header: - writer.writeheader() - writer.writerow(values) + writer.writeheader() + writer.writerows(rows) + handle.flush() + os.fsync(handle.fileno()) + temporary.replace(path) + + def flush_pending_diagnostics(self) -> None: + """Persist the final partial diagnostic interval, if one exists.""" + + self._flush_diagnostics() @torch.no_grad() def step(self, closure=None): diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/optimizers.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/optimizers.py index 102e4d29..b1a97ad6 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/optimizers.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/optimizers.py @@ -209,8 +209,11 @@ def make_optimizer_handles( ) ] - if family == "adamw": - optimizer = torch.optim.AdamW( + if family in {"adam", "adamw"}: + optimizer_class = ( + torch.optim.Adam if family == "adam" else torch.optim.AdamW + ) + optimizer = optimizer_class( _decay_groups(named, float(profile["weight_decay"])), lr=float(profile["learning_rate"]), betas=( diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/provenance.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/provenance.py new file mode 100644 index 00000000..4f6b14df --- /dev/null +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/provenance.py @@ -0,0 +1,183 @@ +from __future__ import annotations + +"""Source-tree provenance used by manifests and restart fingerprints.""" + +import hashlib +from importlib import metadata as importlib_metadata +from pathlib import Path +import platform +import re +import subprocess +from typing import Any + + +_SCIENTIFIC_DISTRIBUTIONS = { + "torch": "torch", + "torch-xla": "torch-xla", + "numpy": "numpy", + "pandas": "pandas", + "scipy": "scipy", + "PyYAML": "PyYAML", + "weightwatcher": "weightwatcher", + "powerlaw": "powerlaw", + "datasets": "datasets", + "tiktoken": "tiktoken", + "sacrebleu": "sacrebleu", + "matplotlib": "matplotlib", + "jupyter": "jupyter", + "ipykernel": "ipykernel", + "nbformat": "nbformat", + "papermill": "papermill", + "packaging": "packaging", +} + + +def _git(root: Path, *arguments: str) -> str: + completed = subprocess.run( + ["git", "-C", str(root), *arguments], + check=True, + capture_output=True, + text=True, + timeout=10, + ) + return completed.stdout.strip() + + +def _git_optional(root: Path, *arguments: str) -> str: + try: + return _git(root, *arguments) + except (OSError, subprocess.SubprocessError): + return "unknown" + + +def repository_provenance() -> dict[str, Any]: + """Return a deterministic source identity without assuming a clone path. + + A clean commit is the normal production identity. For an explicitly + allowed development run, the SHA-256 of the tracked diff is also included + so restart compatibility cannot silently survive a code change. + """ + + start = Path(__file__).resolve().parent + try: + root = Path(_git(start, "rev-parse", "--show-toplevel")) + commit = _git(root, "rev-parse", "HEAD") + branch = _git(root, "rev-parse", "--abbrev-ref", "HEAD") + describe = _git(root, "describe", "--tags", "--always", "--dirty") + tags = tuple( + value + for value in _git_optional( + root, "tag", "--points-at", "HEAD" + ).splitlines() + if value and value != "unknown" + ) + status = _git(root, "status", "--porcelain=v1", "--untracked-files=all") + tracked_diff = _git(root, "diff", "--binary", "HEAD", "--") + return { + "available": True, + "repository_name": root.name, + "commit": commit, + "branch": branch, + "describe": describe, + "tags_at_commit": list(tags), + "tag_status": ",".join(tags) if tags else "untagged", + "origin_url": _git_optional(root, "remote", "get-url", "origin"), + "dirty": bool(status), + "tracked_diff_sha256": hashlib.sha256( + tracked_diff.encode("utf-8") + ).hexdigest(), + "untracked_file_count": sum( + line.startswith("??") for line in status.splitlines() + ), + } + except (OSError, subprocess.SubprocessError, ValueError): + return { + "available": False, + "repository_name": "unknown", + "commit": "unknown", + "branch": "unknown", + "describe": "unknown", + "tags_at_commit": [], + "tag_status": "unknown", + "origin_url": "unknown", + "dirty": None, + "tracked_diff_sha256": "unknown", + "untracked_file_count": None, + } + + +def source_fingerprint_payload() -> dict[str, Any]: + """Return only source fields that change executable semantics.""" + + source = repository_provenance() + return { + "available": source["available"], + "commit": source["commit"], + "dirty": source["dirty"], + "tracked_diff_sha256": source["tracked_diff_sha256"], + "untracked_file_count": source["untracked_file_count"], + } + + +def scientific_dependency_versions() -> dict[str, str]: + """Return exact direct and transitive campaign dependency identities. + + Including the installed dependency closure in every protocol fingerprint + prevents seeds run weeks apart from being pooled after a silent transitive + library upgrade. Optional requirements are included only when installed on + the current hardware block. + """ + + try: + from packaging.requirements import InvalidRequirement, Requirement + except ImportError as exc: + raise RuntimeError( + "packaging is required to inventory the dependency closure" + ) from exc + + versions = {"python": platform.python_version()} + pending: list[str] = [] + for name, distribution in _SCIENTIFIC_DISTRIBUTIONS.items(): + try: + versions[name] = importlib_metadata.version(distribution) + pending.append(distribution) + except importlib_metadata.PackageNotFoundError: + versions[name] = "not-installed" + + normalized = lambda value: re.sub( # noqa: E731 - compact local invariant + r"[-_.]+", "-", str(value).strip() + ).lower() + visited: set[str] = set() + while pending: + requested = pending.pop() + requested_key = normalized(requested) + if requested_key in visited: + continue + visited.add(requested_key) + try: + distribution = importlib_metadata.distribution(requested) + except importlib_metadata.PackageNotFoundError: + continue + actual_name = str(distribution.metadata.get("Name", requested)).strip() + if not actual_name: + raise RuntimeError( + f"installed dependency {requested!r} has no distribution name" + ) + versions.setdefault(actual_name, str(distribution.version)) + for requirement_text in distribution.requires or (): + try: + requirement = Requirement(requirement_text) + except InvalidRequirement as exc: + raise RuntimeError( + f"installed dependency {actual_name} has an invalid " + f"requirement: {requirement_text!r}" + ) from exc + dependency_key = normalized(requirement.name) + if dependency_key in visited: + continue + try: + importlib_metadata.version(requirement.name) + except importlib_metadata.PackageNotFoundError: + continue + pending.append(requirement.name) + return versions diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py index 977719b6..79905bba 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py @@ -1,8 +1,13 @@ from __future__ import annotations import csv +from datetime import datetime, timezone +import hashlib +from importlib import metadata as importlib_metadata import json +import os from pathlib import Path +import sys import pandas as pd import torch @@ -11,16 +16,25 @@ CompletedRunValidationError, validate_completed_run, ) +from .checkpoints import model_state_sha256 from .config import tokens_per_step from .evaluation import evaluate_bleu, evaluate_probe from .model import GPT +from .provenance import ( + repository_provenance, + scientific_dependency_versions, +) from .runtime import runtime_metadata METRIC_FIELDS = [ "step", "tokens_seen", "epoch", "elapsed_sec", "tokens_per_sec", "primary_lr", "auxiliary_lr", "train_loss", "train_perplexity", - "train_accuracy", "val_loss", "val_perplexity", "val_accuracy", - "test_loss", "test_perplexity", "test_accuracy", "test_bleu", + "train_bits_per_token", "train_accuracy", "train_top5_accuracy", + "val_loss", "val_perplexity", "val_bits_per_token", "val_accuracy", + "val_top5_accuracy", "test_loss", "test_perplexity", + "test_bits_per_token", "test_accuracy", "test_top5_accuracy", + "test_bleu", "test_continuation_token_accuracy", + "test_continuation_exact_match", "val_generalization_gap", "test_generalization_gap", "grad_norm_pre_clip", "grad_norm_post_clip", "gradient_clipped", "weight_norm", "update_norm_since_eval", "update_to_weight_ratio", @@ -31,9 +45,148 @@ "nominal_epoch", "checkpoint_path", "test_monitoring_only", + "test_held_out", ] +def _package_versions() -> dict[str, str]: + versions = scientific_dependency_versions() + packages = ( + "rg-nanogpt-one-head", + "papermill", + ) + for package in packages: + try: + versions[package] = importlib_metadata.version(package) + except importlib_metadata.PackageNotFoundError: + versions[package] = "not-installed" + return versions + + +_COMMON_RUNTIME_IDENTITY_FIELDS = ( + "platform", + "machine", + "python_version", + "accelerator", + "device", + "torch_version", + "float32_matmul_precision", + "deterministic_algorithms", + "deterministic_warn_only", + "hardware_block_id", + "hardware_block_id_source", +) +_ACCELERATOR_RUNTIME_IDENTITY_FIELDS = { + "cuda": ( + "cuda_version", + "cudnn_version", + "cuda_device_name", + "cuda_device_capability", + "cuda_device_count", + "cuda_device_uuid", + "cuda_driver_version", + "cuda_device_total_memory_bytes", + "cuda_multi_processor_count", + "cuda_nvidia_smi_memory_mib", + "cuda_matmul_allow_tf32", + "cudnn_allow_tf32", + ), + "mps": ( + "mps_built", + "mps_available", + "mac_hardware_model", + "mac_cpu_brand", + "mac_memory_bytes", + ), + "tpu": ( + "torch_xla_version", + "pjrt_device", + "tpu_accelerator_type", + "xla_process_count", + "xla_process_index", + "xla_addressable_device_count", + ), +} + + +def _read_existing_manifest(path: Path) -> dict | None: + if not path.exists(): + return None + if not path.is_file() or path.stat().st_size == 0: + raise RuntimeError(f"existing manifest is missing or empty: {path}") + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError(f"existing manifest is unreadable: {path}: {exc}") from exc + if not isinstance(payload, dict): + raise RuntimeError(f"existing manifest is not a JSON object: {path}") + return payload + + +def runtime_identity_payload(metadata: dict) -> dict: + accelerator = str(metadata.get("accelerator", "")) + fields = ( + *_COMMON_RUNTIME_IDENTITY_FIELDS, + *_ACCELERATOR_RUNTIME_IDENTITY_FIELDS.get(accelerator, ()), + ) + missing = [field for field in fields if field not in metadata] + if missing: + raise RuntimeError( + "runtime metadata lacks identity fields: " + ", ".join(missing) + ) + return {field: metadata[field] for field in fields} + + +def validate_existing_manifest_runtime( + run_dir: str | Path, + current_runtime: dict, +) -> dict | None: + """Reject corrupt or cross-runtime reuse before any run artifacts mutate.""" + + manifest_path = Path(run_dir) / "manifest.json" + previous = _read_existing_manifest(manifest_path) + if previous is None: + ignored_pre_manifest_names = {"muonclip_walk_location.json"} + substantive = [ + path + for path in Path(run_dir).iterdir() + if path.name not in ignored_pre_manifest_names + ] + if substantive: + raise RuntimeError( + "run artifacts exist without manifest.json; refusing to infer " + "runtime/source provenance before reuse or resume: " + + ", ".join(str(path) for path in substantive[:12]) + ) + return None + previous_runtime = previous.get("runtime_environment") + if not isinstance(previous_runtime, dict): + raise RuntimeError( + f"existing manifest has no runtime_environment mapping: {manifest_path}" + ) + previous_identity = runtime_identity_payload(previous_runtime) + current_identity = runtime_identity_payload(current_runtime) + mismatches = { + field: ( + previous_identity.get(field, ""), + current_identity.get(field, ""), + ) + for field in sorted(set(previous_identity) | set(current_identity)) + if previous_identity.get(field, "") + != current_identity.get(field, "") + } + if mismatches: + detail = "; ".join( + f"{field}: existing={old!r}, current={new!r}" + for field, (old, new) in mismatches.items() + ) + raise RuntimeError( + "Refusing cross-runtime reuse/resume before modifying artifacts: " + + detail + ) + return previous + + def run_directory( results_root: str | Path, optimizer: str, @@ -122,6 +275,50 @@ def truncate_spectral_after( path.unlink(missing_ok=True) +def truncate_muonclip_qk_after( + run_dir: Path, + resume_step: int, +) -> None: + """Discard QK diagnostics newer than the verified restart checkpoint. + + MuonClip flushes a diagnostic interval immediately after the optimizer + update that closes it. A process can therefore stop after the CSV row is + durable but before ``checkpoint_latest.pt`` is replaced. Retaining only + rows at or before the checkpoint step makes a fresh-process resume + idempotent; a row at the checkpoint itself is already represented by the + checkpoint's reset diagnostic accumulator and remains valid. + """ + + path = Path(run_dir) / "muonclip_qk.csv" + if not path.is_file(): + return + try: + frame = pd.read_csv(path) + except Exception as exc: + raise RuntimeError( + f"could not validate MuonClip QK diagnostics before resume: {path}" + ) from exc + if "step" not in frame.columns: + raise RuntimeError( + f"MuonClip QK diagnostics have no step column: {path}" + ) + steps = pd.to_numeric(frame["step"], errors="coerce") + if steps.isna().any() or not steps.mod(1).eq(0).all(): + raise RuntimeError( + f"MuonClip QK diagnostics contain invalid steps: {path}" + ) + frame = frame.loc[steps <= int(resume_step)].copy() + retained_steps = pd.to_numeric(frame["step"], errors="raise") + if retained_steps.duplicated().any(): + raise RuntimeError( + f"MuonClip QK diagnostics contain duplicate verified steps: {path}" + ) + frame = frame.sort_values("step") + temporary = path.with_suffix(path.suffix + ".tmp") + frame.to_csv(temporary, index=False) + temporary.replace(path) + + def write_manifest( run_dir: Path, *, @@ -139,20 +336,58 @@ def write_manifest( fingerprint: str, model: GPT, ) -> None: + now = datetime.now(timezone.utc).isoformat() + manifest_path = run_dir / "manifest.json" + started_at = now + resume_count = 0 + current_runtime = runtime_metadata(device) + previous = validate_existing_manifest_runtime(run_dir, current_runtime) + current_model_hash = model_state_sha256(model.state_dict()) + if previous is not None: + started_at = str(previous.get("run_started_at_utc", now)) + resume_count = int(previous.get("resume_count", 0)) + 1 + initial_model_hash = str(previous.get("initial_model_sha256", "")) + if len(initial_model_hash) != 64: + raise RuntimeError( + "existing manifest has no valid initial-model tensor hash" + ) + else: + initial_model_hash = current_model_hash + canonical_config = json.dumps( + cfg, + sort_keys=True, + separators=(",", ":"), + default=str, + ) payload = { - "schema_version": 2, + "schema_version": 3, + "run_started_at_utc": started_at, + "manifest_updated_at_utc": now, + "resume_count": resume_count, + "config_sha256": hashlib.sha256( + canonical_config.encode("utf-8") + ).hexdigest(), + "invocation": { + "argv": [str(value) for value in sys.argv], + "campaign_command": str( + os.environ.get("RG_NANOGPT_CAMPAIGN_COMMAND", "") + ), + }, "protocol": cfg["protocol"], "optimizer": optimizer_name, "optimizer_profile": profile, "seed": int(seed), "device": str(device), - "runtime_environment": runtime_metadata(device), + "runtime_environment": current_runtime, + "source_repository": repository_provenance(), + "initial_model_sha256": initial_model_hash, "storage": { "data_root": str(Path(data_root)), "results_root": str(Path(results_root)), "run_dir": str(Path(run_dir)), }, "torch_version": torch.__version__, + "package_versions": _package_versions(), "model": cfg["model"], "parameter_count": model.parameter_count(), "data_metadata": data_metadata, @@ -168,8 +403,8 @@ def write_manifest( ), "protocol_fingerprint": fingerprint, "test_policy": ( - "fixed test probes are monitoring-only and never select " - "checkpoints or tune schedules" + "test is held out until post-training; validation selects the " + "best checkpoint and test never tunes the protocol" ), "bleu_policy": ( "fixed greedy held-out continuation BLEU; secondary diagnostic, " @@ -186,7 +421,7 @@ def write_manifest( ), encoding="utf-8", ) - temporary.replace(run_dir / "manifest.json") + temporary.replace(manifest_path) def checkpoint_eval( @@ -215,6 +450,14 @@ def checkpoint_eval( "step": int(payload["step"]), "loss": float(metrics["loss"]), "perplexity": float(metrics["perplexity"]), + "bits_per_token": float(metrics["bits_per_token"]), "accuracy": float(metrics["accuracy"]), + "top5_accuracy": float(metrics["top5_accuracy"]), "bleu": float(bleu["bleu"]), + "continuation_token_accuracy": float( + bleu.get("continuation_token_accuracy", float("nan")) + ), + "continuation_exact_match": float( + bleu.get("continuation_exact_match", float("nan")) + ), } diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/runtime.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/runtime.py index ea1e6aa9..2cea820e 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/runtime.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/runtime.py @@ -1,6 +1,7 @@ from __future__ import annotations import argparse +import hashlib import importlib.util import json import math @@ -8,6 +9,9 @@ import platform import random import re +import shutil +import subprocess +import sys from pathlib import Path from typing import Any, Iterable @@ -201,6 +205,13 @@ def configure_runtime(device: torch.device, cfg: dict) -> None: torch.set_float32_matmul_precision( str(cfg["runtime"].get("matmul_precision", "high")) ) + if device.type == "cuda": + allow_tf32 = bool(cfg["runtime"].get("allow_tf32", False)) + torch.backends.cuda.matmul.allow_tf32 = allow_tf32 + torch.backends.cudnn.allow_tf32 = allow_tf32 + torch.backends.cudnn.benchmark = bool( + cfg["runtime"].get("cudnn_benchmark", False) + ) if device.type == "mps" and bool(cfg["runtime"].get("mps_fallback", True)): os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") if device.type == "xla": @@ -242,8 +253,17 @@ def configure_runtime(device: torch.device, cfg: dict) -> None: f"Detected process_count={count}. Run the ordinary single-process " "launcher or implement an explicitly distributed protocol." ) - if bool(cfg["runtime"].get("deterministic_algorithms", False)): - torch.use_deterministic_algorithms(True, warn_only=True) + deterministic = bool( + cfg["runtime"].get("deterministic_algorithms", False) + ) + torch.use_deterministic_algorithms( + deterministic, + warn_only=bool( + cfg["runtime"].get("deterministic_warn_only", True) + ), + ) + if device.type == "cuda": + torch.backends.cudnn.deterministic = deterministic def mark_step(device: torch.device) -> None: @@ -361,14 +381,158 @@ def restore_accelerator_rng_state( xm.set_rng_state(int(payload["xla_random_state"]), device=device) +def _command_output(arguments: list[str]) -> str: + try: + return subprocess.run( + arguments, + check=True, + capture_output=True, + text=True, + timeout=10, + ).stdout.strip() + except (OSError, subprocess.SubprocessError): + return "" + + +def _hardware_block_identity(metadata: dict[str, Any]) -> tuple[str, str]: + override = str(os.environ.get("RG_NANOGPT_HARDWARE_BLOCK_ID", "")).strip() + if override: + if len(override) > 200 or "\n" in override: + raise RuntimeError("RG_NANOGPT_HARDWARE_BLOCK_ID is malformed") + return override, "user" + + accelerator = str(metadata["accelerator"]) + if accelerator == "cuda": + required = ( + "cuda_device_name", + "cuda_device_capability", + "cuda_device_uuid", + "cuda_device_total_memory_bytes", + "cuda_driver_version", + ) + elif accelerator == "mps": + required = ("mac_hardware_model", "mac_cpu_brand", "mac_memory_bytes") + elif accelerator == "tpu": + required = ("tpu_accelerator_type",) + else: + required = ("platform", "machine", "processor") + missing = [ + key + for key in required + if metadata.get(key) in (None, "", "unknown", 0) + ] + if missing and accelerator != "cpu": + raise RuntimeError( + "could not determine a complete hardware block identity (missing " + + ", ".join(missing) + + "); set RG_NANOGPT_HARDWARE_BLOCK_ID to a stable collaborator-" + "chosen identifier for this homogeneous device block" + ) + payload = {key: metadata.get(key, "unknown") for key in required} + digest = hashlib.sha256( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode( + "utf-8" + ) + ).hexdigest()[:20] + return f"auto-{accelerator}-{digest}", "auto" + + def runtime_metadata(device: torch.device) -> dict[str, Any]: metadata: dict[str, Any] = { "platform": platform.platform(), + "machine": platform.machine(), + "processor": platform.processor(), + "python_version": sys.version, + "python_executable": sys.executable, "accelerator": accelerator_name(device), "device": str(device), "torch_version": torch.__version__, + "float32_matmul_precision": torch.get_float32_matmul_precision(), + "deterministic_algorithms": bool( + torch.are_deterministic_algorithms_enabled() + ), + "deterministic_warn_only": bool( + getattr( + torch, + "is_deterministic_algorithms_warn_only_enabled", + lambda: False, + )() + ), } - if device.type == "xla": + if device.type == "cuda": + index = device.index if device.index is not None else 0 + properties = torch.cuda.get_device_properties(index) + smi_values: list[str] = [] + nvidia_smi = shutil.which("nvidia-smi") + if nvidia_smi is not None: + output = _command_output( + [ + nvidia_smi, + f"--id={index}", + "--query-gpu=uuid,driver_version,memory.total", + "--format=csv,noheader,nounits", + ] + ) + if output: + smi_values = [ + value.strip() + for value in output.splitlines()[0].split(",") + ] + property_uuid = str(getattr(properties, "uuid", "")).strip() + metadata.update( + { + "cuda_version": torch.version.cuda, + "cudnn_version": torch.backends.cudnn.version(), + "cuda_device_name": torch.cuda.get_device_name(index), + "cuda_device_capability": list( + torch.cuda.get_device_capability(index) + ), + "cuda_device_count": torch.cuda.device_count(), + "cuda_device_uuid": ( + smi_values[0] + if len(smi_values) >= 1 and smi_values[0] + else property_uuid or "unknown" + ), + "cuda_driver_version": ( + smi_values[1] + if len(smi_values) >= 2 and smi_values[1] + else "unknown" + ), + "cuda_device_total_memory_bytes": int( + properties.total_memory + ), + "cuda_multi_processor_count": int( + properties.multi_processor_count + ), + "cuda_nvidia_smi_memory_mib": ( + float(smi_values[2]) + if len(smi_values) >= 3 and smi_values[2] + else None + ), + "cuda_matmul_allow_tf32": bool( + torch.backends.cuda.matmul.allow_tf32 + ), + "cudnn_allow_tf32": bool(torch.backends.cudnn.allow_tf32), + } + ) + elif device.type == "mps": + mac_model = _command_output(["sysctl", "-n", "hw.model"]) + mac_brand = _command_output( + ["sysctl", "-n", "machdep.cpu.brand_string"] + ) + mac_memory = _command_output(["sysctl", "-n", "hw.memsize"]) + metadata.update( + { + "mps_built": bool(torch.backends.mps.is_built()), + "mps_available": bool(torch.backends.mps.is_available()), + "mac_hardware_model": mac_model or "unknown", + "mac_cpu_brand": mac_brand or "unknown", + "mac_memory_bytes": ( + int(mac_memory) if mac_memory.isdigit() else 0 + ), + } + ) + elif device.type == "xla": modules = _load_xla(required=True) assert modules is not None torch_xla, xr, _ = modules @@ -378,6 +542,9 @@ def runtime_metadata(device: torch.device) -> dict[str, Any]: torch_xla, "__version__", "unknown" ), "pjrt_device": _xla_device_type(xr), + "tpu_accelerator_type": str( + os.environ.get("TPU_ACCELERATOR_TYPE", "unknown") + ), "xla_process_count": _xla_process_count(xr), "xla_process_index": _xla_process_index(xr), "xla_addressable_device_count": int( @@ -385,6 +552,9 @@ def runtime_metadata(device: torch.device) -> dict[str, Any]: ), } ) + block_id, block_source = _hardware_block_identity(metadata) + metadata["hardware_block_id"] = block_id + metadata["hardware_block_id_source"] = block_source return metadata diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/spectral.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/spectral.py index adec193f..1260d56f 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/spectral.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/spectral.py @@ -1,5 +1,6 @@ from __future__ import annotations +import hashlib import json from pathlib import Path import random @@ -11,6 +12,7 @@ import torch.nn as nn from .model import GPT, transformer_matrix_items +from .checkpoints import model_state_sha256 from .runtime import ( capture_accelerator_rng_state, model_device, @@ -20,6 +22,9 @@ SPECTRAL_METRICS = ( "alpha", + "alpha_raw", + "alpha_clip_xmax", + "alpha_delta", "alpha_weighted", "ERG_gap", "num_traps", @@ -135,6 +140,14 @@ def _atomic_csv(path: Path, frame: pd.DataFrame) -> None: temporary.replace(path) +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(1024 * 1024): + digest.update(block) + return digest.hexdigest() + + def _append_deduplicated( path: Path, frame: pd.DataFrame, @@ -156,6 +169,175 @@ def _append_deduplicated( _atomic_csv(path, combined) +def _validate_weightwatcher_frame( + frame: pd.DataFrame, + *, + finger_policy: str | bool, +) -> None: + """Reject incomplete results and stale pre-finger-policy caches.""" + + required = { + "matrix_name", + "alpha", + "alpha_raw", + "ERG_gap", + "num_traps", + "rand_distance", + "finger_policy", + "primary_alpha_variant", + "weightwatcher_analysis_calls", + "run_seed", + "diagnostic_seed", + "protocol_fingerprint", + "model_state_sha256", + } + if finger_policy == "clip_xmax": + required.update( + { + "raw_alpha", + "alpha_clip_xmax", + "alpha_delta", + "num_fingers", + } + ) + missing = sorted(required.difference(frame.columns)) + if missing: + raise RuntimeError( + "WeightWatcher result is missing the required one-pass schema: " + + ", ".join(missing) + + ". Remove the stale run directory before restarting." + ) + if len(frame) != 6 or frame["matrix_name"].nunique() != 6: + raise RuntimeError( + "WeightWatcher must return exactly the six declared transformer " + "matrices" + ) + + numeric = [ + "alpha", + "alpha_raw", + "ERG_gap", + "num_traps", + "rand_distance", + "weightwatcher_analysis_calls", + ] + expected_policy = "none" + expected_variant = "raw" + if finger_policy == "clip_xmax": + numeric.extend( + ["raw_alpha", "alpha_clip_xmax", "alpha_delta", "num_fingers"] + ) + expected_policy = "fix_fingers=clip_xmax" + expected_variant = "clip_xmax" + values = frame[numeric].apply(pd.to_numeric, errors="coerce") + if not np.isfinite(values.to_numpy(dtype=float)).all(): + raise RuntimeError( + "WeightWatcher required one-pass spectral values contain NaN or " + "infinity" + ) + calls = pd.to_numeric( + frame["weightwatcher_analysis_calls"], errors="coerce" + ) + if not calls.eq(1).all(): + raise RuntimeError( + "WeightWatcher analysis must be called exactly once per checkpoint" + ) + if not frame["finger_policy"].astype(str).eq(expected_policy).all(): + raise RuntimeError("WeightWatcher finger-policy metadata is inconsistent") + if not frame["primary_alpha_variant"].astype(str).eq( + expected_variant + ).all(): + raise RuntimeError("WeightWatcher primary-alpha metadata is inconsistent") + if finger_policy == "clip_xmax": + alpha = values["alpha"].to_numpy(dtype=float) + raw = values["raw_alpha"].to_numpy(dtype=float) + canonical_raw = values["alpha_raw"].to_numpy(dtype=float) + clipped = values["alpha_clip_xmax"].to_numpy(dtype=float) + delta = values["alpha_delta"].to_numpy(dtype=float) + if not np.allclose(alpha, clipped, rtol=0.0, atol=0.0): + raise RuntimeError("WeightWatcher clipped-alpha aliases disagree") + if not np.allclose(raw, canonical_raw, rtol=0.0, atol=0.0): + raise RuntimeError("WeightWatcher raw-alpha aliases disagree") + if not np.allclose( + delta, + canonical_raw - clipped, + rtol=1e-12, + atol=1e-12, + ): + raise RuntimeError("WeightWatcher alpha-delta values are inconsistent") + if (values["num_fingers"] < 0).any(): + raise RuntimeError("WeightWatcher num_fingers must be nonnegative") + + +def _record_successful_frame( + frame: pd.DataFrame, + *, + spectral_root: Path, + raw_path: Path, + step: int, + tokens_seen: int, + train_tokens: int, + run_seed: int, + diagnostic_seed: int, + protocol_fingerprint: str, + model_hash: str, +) -> dict[str, Any]: + epoch = tokens_seen / max(1, int(train_tokens)) + for column, expected in (("step", step), ("tokens_seen", tokens_seen)): + if column not in frame.columns: + raise RuntimeError( + f"WeightWatcher cached result has no {column} column" + ) + observed = pd.to_numeric(frame[column], errors="coerce") + if observed.isna().any() or not observed.eq(int(expected)).all(): + raise RuntimeError( + f"WeightWatcher cached {column} does not match this checkpoint" + ) + _append_deduplicated( + spectral_root / "layers.csv", + frame, + keys=["step", "matrix_name"], + ) + summary = summarize_spectral_frame( + frame, + step=step, + tokens_seen=tokens_seen, + epoch=epoch, + ) + _append_deduplicated( + spectral_root / "summary.csv", + pd.DataFrame([summary]), + keys=["step"], + ) + status = { + "step": int(step), + "tokens_seen": int(tokens_seen), + "epoch": float(epoch), + "completed": True, + "raw_path": str(raw_path), + "raw_csv_sha256": _sha256(raw_path), + "run_seed": int(run_seed), + "diagnostic_seed": int(diagnostic_seed), + "protocol_fingerprint": str(protocol_fingerprint), + "model_state_sha256": str(model_hash), + "weightwatcher_analysis_calls": 1, + "finger_policy": str(frame["finger_policy"].iloc[0]), + "alpha_valid_matrices": int(summary["alpha_n"]), + "alpha_raw_valid_matrices": int(summary["alpha_raw_n"]), + "alpha_clip_xmax_valid_matrices": int( + summary["alpha_clip_xmax_n"] + ), + "ERG_gap_valid_matrices": int(summary["ERG_gap_n"]), + "num_traps_valid_matrices": int(summary["num_traps_n"]), + "rand_distance_valid_matrices": int(summary["rand_distance_n"]), + } + (spectral_root / f"status_step_{int(step):07d}.json").write_text( + json.dumps(status, indent=2, sort_keys=True), + encoding="utf-8", + ) + return summary + + def summarize_spectral_frame( frame: pd.DataFrame, *, @@ -169,6 +351,19 @@ def summarize_spectral_frame( "epoch": float(epoch), "n_matrices": int(len(frame)), } + for identity in ( + "run_seed", + "diagnostic_seed", + "protocol_fingerprint", + "model_state_sha256", + ): + if identity in frame.columns: + values = frame[identity].drop_duplicates() + if len(values) != 1: + raise RuntimeError( + f"WeightWatcher frame has multiple {identity} values" + ) + summary[identity] = values.iloc[0] for metric in SPECTRAL_METRICS: values = ( pd.to_numeric( @@ -222,16 +417,21 @@ def run_weightwatcher( train_tokens: int, config: dict[str, Any], seed: int, + fingerprint: str, ) -> dict[str, Any]: - """Run WeightWatcher exactly with ERG=True and randomize=True. + """Run one WeightWatcher analysis with ERG and randomization enabled. - `alpha`, `ERG_gap`, `num_traps`, and `rand_distance` are retained directly - from WeightWatcher. `rand_distance` is the Jensen-Shannon distance between - the empirical ESD and the entry-wise randomized ESD. No fallback alpha, - proxy trap count, synthesized ERG gap, or replacement random-distance - statistic is permitted. Every CPU and accelerator RNG stream is restored - after the randomized diagnostic so measurement cannot change the - subsequent training path. + With ``fix_fingers='clip_xmax'``, WeightWatcher returns the corrected + exponent in ``alpha`` and the uncorrected exponent from the same call in + ``raw_alpha``. The baseline stores those values canonically as + ``alpha_clip_xmax`` and ``alpha_raw``. It does not run WeightWatcher a + second time. `ERG_gap`, `num_traps`, and `rand_distance` are retained + directly from that same call. `rand_distance` is the Jensen-Shannon + distance between the empirical ESD and the entry-wise randomized ESD. No + fallback alpha, proxy trap count, synthesized ERG gap, or replacement + random-distance statistic is permitted. Every CPU and accelerator RNG + stream is restored after the randomized diagnostic so measurement cannot + change the subsequent training path. """ try: @@ -250,22 +450,69 @@ def run_weightwatcher( raw_root / f"weightwatcher_step_{int(step):07d}.csv" ) + if not str(fingerprint).strip(): + raise ValueError("WeightWatcher requires a non-empty protocol fingerprint") + device = model_device(model) + synchronize(device) + current_model_hash = model_state_sha256(model.state_dict()) + diagnostic_seed = int(seed) + 1_000_003 + int(step) if raw_path.is_file(): frame = pd.read_csv(raw_path) - return summarize_spectral_frame( + finger_policy = config.get("fix_fingers", False) + _validate_weightwatcher_frame( + frame, + finger_policy=finger_policy, + ) + expected_identities = { + "run_seed": int(seed), + "diagnostic_seed": int(diagnostic_seed), + "protocol_fingerprint": str(fingerprint), + "model_state_sha256": str(current_model_hash), + } + for column, expected in expected_identities.items(): + if column not in frame.columns or not frame[column].astype(str).eq( + str(expected) + ).all(): + raise RuntimeError( + f"cached WeightWatcher {column} does not match the " + "current checkpoint" + ) + status_path = spectral_root / f"status_step_{int(step):07d}.json" + try: + status = json.loads(status_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError( + "cached WeightWatcher raw CSV has no valid integrity status" + ) from exc + if ( + not isinstance(status, dict) + or status.get("completed") is not True + or str(status.get("raw_csv_sha256", "")) != _sha256(raw_path) + or any( + str(status.get(column, "")) != str(expected) + for column, expected in expected_identities.items() + ) + ): + raise RuntimeError( + "cached WeightWatcher raw CSV/status integrity binding failed" + ) + return _record_successful_frame( frame, + spectral_root=spectral_root, + raw_path=raw_path, step=step, tokens_seen=tokens_seen, - epoch=tokens_seen / max(1, int(train_tokens)), + train_tokens=train_tokens, + run_seed=seed, + diagnostic_seed=diagnostic_seed, + protocol_fingerprint=fingerprint, + model_hash=current_model_hash, ) - device = model_device(model) - synchronize(device) py_state = random.getstate() np_state = np.random.get_state() torch_state = torch.random.get_rng_state() accelerator_state = capture_accelerator_rng_state(device) - diagnostic_seed = int(seed) + 1_000_003 + int(step) random.seed(diagnostic_seed) np.random.seed(diagnostic_seed % (2**32 - 1)) torch.manual_seed(diagnostic_seed) @@ -276,11 +523,20 @@ def run_weightwatcher( # host and never exposes the live accelerator model to NumPy. holder = WeightMatrixHolder(model) watcher = ww.WeightWatcher(model=holder) + analysis_kwargs: dict[str, Any] = { + "ERG": True, + "randomize": True, + "plot": False, + "min_evals": int(config.get("min_evals", 20)), + } + finger_policy = config.get("fix_fingers", False) + if finger_policy: + analysis_kwargs["fix_fingers"] = finger_policy + analysis_kwargs["max_fingers"] = int( + config.get("max_fingers", 10) + ) details = watcher.analyze( - ERG=True, - randomize=True, - plot=False, - min_evals=int(config.get("min_evals", 20)), + **analysis_kwargs, ) if details is None or len(details) == 0: raise RuntimeError( @@ -290,28 +546,37 @@ def run_weightwatcher( pd.DataFrame(details), holder.matrix_metadata, ) - required_columns = ( - "alpha", - "ERG_gap", - "num_traps", - "rand_distance", - ) - missing = [ - column - for column in required_columns - if column not in frame.columns - ] - if missing: - raise RuntimeError( - "WeightWatcher did not return required " - "ERG/randomization columns: " - + ", ".join(missing) + if finger_policy == "clip_xmax": + if "raw_alpha" not in frame.columns: + raise RuntimeError( + "WeightWatcher fix_fingers='clip_xmax' did not return " + "the required raw_alpha column" + ) + if "num_fingers" not in frame.columns: + raise RuntimeError( + "WeightWatcher fix_fingers='clip_xmax' did not return " + "the required num_fingers column" + ) + frame["alpha_raw"] = pd.to_numeric( + frame["raw_alpha"], errors="coerce" ) - if frame[list(required_columns)].isna().any().any(): - raise RuntimeError( - "WeightWatcher required alpha/ERG_gap/num_traps/" - "rand_distance values contain NaN" + frame["alpha_clip_xmax"] = pd.to_numeric( + frame["alpha"], errors="coerce" + ) + frame["alpha_delta"] = ( + frame["alpha_raw"] - frame["alpha_clip_xmax"] ) + frame["finger_policy"] = "fix_fingers=clip_xmax" + frame["primary_alpha_variant"] = "clip_xmax" + else: + frame["alpha_raw"] = pd.to_numeric( + frame["alpha"], errors="coerce" + ) + frame["alpha_clip_xmax"] = np.nan + frame["alpha_delta"] = np.nan + frame["finger_policy"] = "none" + frame["primary_alpha_variant"] = "raw" + frame["weightwatcher_analysis_calls"] = 1 epoch = tokens_seen / max(1, int(train_tokens)) frame.insert(0, "step", int(step)) frame.insert(1, "tokens_seen", int(tokens_seen)) @@ -321,54 +586,26 @@ def run_weightwatcher( "diagnostic_seed", int(diagnostic_seed), ) - _atomic_csv(raw_path, frame) - _append_deduplicated( - spectral_root / "layers.csv", + frame.insert(4, "run_seed", int(seed)) + frame.insert(5, "protocol_fingerprint", str(fingerprint)) + frame.insert(6, "model_state_sha256", str(current_model_hash)) + _validate_weightwatcher_frame( frame, - keys=["step", "matrix_name"], + finger_policy=finger_policy, ) - summary = summarize_spectral_frame( + _atomic_csv(raw_path, frame) + return _record_successful_frame( frame, + spectral_root=spectral_root, + raw_path=raw_path, step=step, tokens_seen=tokens_seen, - epoch=epoch, - ) - _append_deduplicated( - spectral_root / "summary.csv", - pd.DataFrame([summary]), - keys=["step"], - ) - status = { - "step": int(step), - "tokens_seen": int(tokens_seen), - "epoch": float(epoch), - "completed": True, - "raw_path": str(raw_path), - "alpha_valid_matrices": int( - summary["alpha_n"] - ), - "ERG_gap_valid_matrices": int( - summary["ERG_gap_n"] - ), - "num_traps_valid_matrices": int( - summary["num_traps_n"] - ), - "rand_distance_valid_matrices": int( - summary["rand_distance_n"] - ), - } - ( - spectral_root - / f"status_step_{int(step):07d}.json" - ).write_text( - json.dumps( - status, - indent=2, - sort_keys=True, - ), - encoding="utf-8", + train_tokens=train_tokens, + run_seed=seed, + diagnostic_seed=diagnostic_seed, + protocol_fingerprint=fingerprint, + model_hash=current_model_hash, ) - return summary except Exception as exc: status = { "step": int(step), diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py index f60e4232..77de9871 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py @@ -11,7 +11,7 @@ save_epoch_model_checkpoint, save_training_checkpoint, ) -from .evaluation import evaluate_bleu, evaluate_probe, random_batch +from .evaluation import evaluate_probe, random_batch from .optimizers import optimizer_step, set_learning_rates, zero_grad from .runtime import ( empty_mps_cache, @@ -35,10 +35,14 @@ def _require_finite_metrics( values = { "train_loss": float(train_metrics["loss"]), "train_perplexity": float(train_metrics["perplexity"]), + "train_bits_per_token": float(train_metrics["bits_per_token"]), "train_accuracy": float(train_metrics["accuracy"]), + "train_top5_accuracy": float(train_metrics["top5_accuracy"]), "val_loss": float(val_metrics["loss"]), "val_perplexity": float(val_metrics["perplexity"]), + "val_bits_per_token": float(val_metrics["bits_per_token"]), "val_accuracy": float(val_metrics["accuracy"]), + "val_top5_accuracy": float(val_metrics["top5_accuracy"]), } bad = [ name @@ -93,6 +97,21 @@ def _evaluation_due( ) +def _resume_diagnostics( + previous_eval_snapshot: list[torch.Tensor], + *, + last_grad_pre: float, + last_grad_post: float, + last_clipped: bool, +) -> dict: + return { + "previous_eval_snapshot": previous_eval_snapshot, + "last_grad_pre": float(last_grad_pre), + "last_grad_post": float(last_grad_post), + "last_clipped": bool(last_clipped), + } + + def execute_training_loop( *, cfg: dict, @@ -101,8 +120,6 @@ def execute_training_loop( arrays: dict, train_probe, val_probe, - test_probe, - bleu_probe, device: torch.device, optimizer_name: str, seed: int, @@ -114,6 +131,7 @@ def execute_training_loop( best_validation_loss: float, best_validation_step: int, elapsed_offset: float, + resume_diagnostics: dict | None, fingerprint: str, train_generator: torch.Generator, epoch_steps: dict[int, float], @@ -125,12 +143,11 @@ def execute_training_loop( latest_checkpoint: Path, best_checkpoint: Path, progress: bool, -) -> tuple[float, int, float]: +) -> tuple[float, int, float, dict]: batch_size = int(cfg["training"]["batch_size"]) grad_accum = int(cfg["training"]["grad_accum_steps"]) block_size = int(cfg["model"]["block_size"]) step_tokens = batch_size * grad_accum * block_size - eval_cfg = cfg["evaluation"] if not 1 <= schedule_steps <= total_steps: raise ValueError( @@ -142,11 +159,23 @@ def execute_training_loop( f"schedule_steps={schedule_steps}" ) - previous_snapshot = parameter_snapshot(model) - last_grad_pre = float("nan") - last_grad_post = float("nan") - last_clipped = False + if resume_diagnostics is None: + if start_step > 0: + raise RuntimeError( + "checkpoint has no deterministic resume diagnostics; refusing " + "to rewrite monitoring rows with reset gradient/update values" + ) + previous_snapshot = parameter_snapshot(model) + last_grad_pre = float("nan") + last_grad_post = float("nan") + last_clipped = False + else: + previous_snapshot = resume_diagnostics["previous_eval_snapshot"] + last_grad_pre = float(resume_diagnostics["last_grad_pre"]) + last_grad_post = float(resume_diagnostics["last_grad_post"]) + last_clipped = bool(resume_diagnostics["last_clipped"]) started = time.time() + final_resume_diagnostics: dict | None = None # CSV rows describe the model state at `completed_steps`, so the recorded # LR must be the LR used by the update that produced that state. At step @@ -181,6 +210,16 @@ def execute_training_loop( ) if evaluation_due: + # These are the diagnostics that an uninterrupted run uses for the + # row at this exact model state. Keep them even after advancing the + # in-memory evaluation snapshot so a crash after evaluation can + # reproduce the row byte-for-byte (apart from wall-clock fields). + current_state_resume_diagnostics = _resume_diagnostics( + previous_snapshot, + last_grad_pre=last_grad_pre, + last_grad_post=last_grad_post, + last_clipped=last_clipped, + ) synchronize(device) train_metrics = evaluate_probe(model, train_probe, device) val_metrics = evaluate_probe(model, val_probe, device) @@ -215,21 +254,18 @@ def execute_training_loop( test_metrics = { "loss": float("nan"), "perplexity": float("nan"), + "bits_per_token": float("nan"), "accuracy": float("nan"), + "top5_accuracy": float("nan"), } - bleu_metrics = {"bleu": float("nan")} - if epoch_due or completed_steps == total_steps: - test_metrics = evaluate_probe( - model, - test_probe, - device, - ) - bleu_metrics = evaluate_bleu( - model, - bleu_probe, - device=device, - batch_size=int(eval_cfg["bleu_batch_size"]), - ) + bleu_metrics = { + "bleu": float("nan"), + "continuation_token_accuracy": float("nan"), + "continuation_exact_match": float("nan"), + } + # Keep the test split genuinely held out during optimization. + # Final and validation-selected checkpoints are evaluated once, + # after training, by engine.checkpoint_eval. tokens_seen = int(completed_steps * step_tokens) actual_epoch = tokens_seen / max(1, train_tokens) @@ -263,24 +299,52 @@ def execute_training_loop( "train_perplexity": float( train_metrics["perplexity"] ), + "train_bits_per_token": float( + train_metrics["bits_per_token"] + ), "train_accuracy": float( train_metrics["accuracy"] ), + "train_top5_accuracy": float( + train_metrics["top5_accuracy"] + ), "val_loss": float(val_metrics["loss"]), "val_perplexity": float( val_metrics["perplexity"] ), + "val_bits_per_token": float( + val_metrics["bits_per_token"] + ), "val_accuracy": float( val_metrics["accuracy"] ), + "val_top5_accuracy": float( + val_metrics["top5_accuracy"] + ), "test_loss": float(test_metrics["loss"]), "test_perplexity": float( test_metrics["perplexity"] ), + "test_bits_per_token": float( + test_metrics["bits_per_token"] + ), "test_accuracy": float( test_metrics["accuracy"] ), + "test_top5_accuracy": float( + test_metrics["top5_accuracy"] + ), "test_bleu": float(bleu_metrics["bleu"]), + "test_continuation_token_accuracy": float( + bleu_metrics.get( + "continuation_token_accuracy", float("nan") + ) + ), + "test_continuation_exact_match": float( + bleu_metrics.get( + "continuation_exact_match", float("nan") + ) + ), "val_generalization_gap": float( val_metrics["loss"] - train_metrics["loss"] ), @@ -323,6 +387,7 @@ def execute_training_loop( "nominal_epoch": nominal_epoch, "checkpoint_path": str(checkpoint_path), "test_monitoring_only": 1, + "test_held_out": 1, } ) epoch_handle.flush() @@ -335,13 +400,15 @@ def execute_training_loop( train_tokens=train_tokens, config=cfg["weightwatcher"], seed=int(seed), + fingerprint=fingerprint, ) if progress: print( "[one-head-ww] " f"optimizer={optimizer_name} seed={seed} " f"epoch={nominal_epoch:.2f} " - f"alpha={ww_summary.get('alpha_median', float('nan')):.3f} " + f"alpha_clip={ww_summary.get('alpha_clip_xmax_median', ww_summary.get('alpha_median', float('nan'))):.3f} " + f"alpha_raw={ww_summary.get('alpha_raw_median', float('nan')):.3f} " f"ERG_gap={ww_summary.get('ERG_gap_median', float('nan')):.3f} " f"num_traps={ww_summary.get('num_traps_mean', float('nan')):.2f}", flush=True, @@ -382,6 +449,9 @@ def execute_training_loop( flush=True, ) + if completed_steps == total_steps: + final_resume_diagnostics = current_state_resume_diagnostics + if completed_steps == total_steps: break @@ -465,11 +535,23 @@ def execute_training_loop( optimizer_name=optimizer_name, seed=int(seed), train_generator=train_generator, + resume_diagnostics=_resume_diagnostics( + previous_snapshot, + last_grad_pre=last_grad_pre, + last_grad_post=last_grad_post, + last_clipped=last_clipped, + ), ) synchronize(device) + if final_resume_diagnostics is None: + raise RuntimeError( + "final model state was not evaluated; deterministic resume " + "diagnostics are unavailable" + ) return ( float(best_validation_loss), int(best_validation_step), float(elapsed_offset + time.time() - started), + final_resume_diagnostics, ) diff --git a/baseline/nanogpt_one_head/tests/test_20260821_campaign.py b/baseline/nanogpt_one_head/tests/test_20260821_campaign.py new file mode 100644 index 00000000..cd77a46e --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_20260821_campaign.py @@ -0,0 +1,909 @@ +from __future__ import annotations + +import ast +import copy +import hashlib +import importlib.util +import json +from pathlib import Path +import re +import subprocess +import sys +import tempfile +from types import ModuleType, SimpleNamespace + +import nbformat +import numpy as np +import pandas as pd +import pytest +import torch +import yaml + + +PACKAGE_ROOT = Path(__file__).resolve().parents[1] +BASELINE_ROOT = PACKAGE_ROOT.parent +EXPERIMENT_ROOT = ( + BASELINE_ROOT + / "experiments" + / "nanogpt_one_head_2026_08_21_baseline" +) +CONFIG_PATH = EXPERIMENT_ROOT / "configs" / "baseline.yaml" +CAMPAIGN_PATH = EXPERIMENT_ROOT / "campaign.yaml" +RUNNER_PATH = EXPERIMENT_ROOT / "scripts" / "run_experiment.py" +REPORT_PATH = EXPERIMENT_ROOT / "scripts" / "build_report.py" +NOTEBOOK_PATH = ( + EXPERIMENT_ROOT + / "notebooks" + / "01_Performance_and_Spectra.ipynb" +) +DOCTOR_SMOKE_PATH = ( + PACKAGE_ROOT / "src" / "rg_nanogpt_one_head" / "doctor_smoke.py" +) + +sys.path.insert(0, str(PACKAGE_ROOT / "src")) + +from rg_nanogpt_one_head.config import ( # noqa: E402 + epoch_step_map, + max_steps, + optimizer_profile, + tokens_per_step, +) +from rg_nanogpt_one_head.model import GPT, GPTConfig # noqa: E402 +from rg_nanogpt_one_head.optimizers import make_optimizer_handles # noqa: E402 +from rg_nanogpt_one_head.spectral import run_weightwatcher # noqa: E402 + + +EXPECTED_ARMS = ["adamw", "muon_clip"] +EXPECTED_SEEDS = [1337, 2027, 4099, 31415, 271828] + + +def _load_runner() -> ModuleType: + spec = importlib.util.spec_from_file_location( + "campaign_run_experiment_20260821", + RUNNER_PATH, + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _load_report() -> ModuleType: + spec = importlib.util.spec_from_file_location( + "campaign_build_report_20260821", + REPORT_PATH, + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _campaign_document() -> dict: + document = yaml.safe_load(CAMPAIGN_PATH.read_text(encoding="utf-8")) + assert isinstance(document, dict) + return document + + +def _baseline_config_document() -> dict: + document = yaml.safe_load(CONFIG_PATH.read_text(encoding="utf-8")) + assert isinstance(document, dict) + return document + + +def test_dated_campaign_has_the_frozen_two_by_five_design_and_grid(): + runner = _load_runner() + report = _load_report() + assert list(runner.CANONICAL_OPTIMIZERS) == EXPECTED_ARMS + assert list(runner.CANONICAL_SEEDS) == EXPECTED_SEEDS + assert runner.EXPECTED_REPLICATES == 10 + assert list(report.OPTIMIZERS) == EXPECTED_ARMS + assert list(report.SEEDS) == EXPECTED_SEEDS + + campaign_document = _campaign_document() + campaign = campaign_document["campaign"] + assert campaign["id"] == "nanogpt_one_head_2026_08_21_baseline_v3" + assert str(campaign["prepared_date"]) == "2026-08-22" + assert campaign["config"] == "configs/baseline.yaml" + assert campaign["optimizers"] == EXPECTED_ARMS + assert campaign["seeds"] == EXPECTED_SEEDS + assert campaign["require_complete_replicates"] == 10 + assert campaign["require_clean_git"] is True + assert campaign["require_tmp_root"] is True + assert campaign["primary_checkpoint_policy"] == ( + "minimum_validation_probe_nll" + ) + assert campaign["protected_test_policy"] == ( + "held_out_posthoc_never_selects" + ) + + cfg = _baseline_config_document() + assert cfg["training"]["seeds"] == EXPECTED_SEEDS + assert cfg["dataset"] == { + "name": "HuggingFaceFW/fineweb-edu", + "config": "sample-10BT", + "split": "train", + "revision": "593b3a867298afb8ce42625a270ef20ddcad28f9", + "tokenizer": "gpt2", + "train_tokens": 80_000_000, + "val_tokens": 1_000_000, + "test_tokens": 1_000_000, + } + assert cfg["model"] == { + "vocab_size": 50_257, + "block_size": 256, + "n_layer": 1, + "n_head": 1, + "n_embd": 128, + "dropout": 0.0, + "bias": False, + "tie_weights": True, + } + assert cfg["training"] == { + "seeds": EXPECTED_SEEDS, + "batch_size": 4, + "grad_accum_steps": 8, + "target_epochs": 4.0, + "epoch_interval": 0.25, + "eval_interval_steps": 500, + "eval_batches": 64, + "checkpoint_interval_steps": 500, + "grad_clip": 1.0, + } + + expected_profile_fields = { + "adamw": { + "family": "adamw", + "learning_rate": 0.0006, + "min_learning_rate": 0.00006, + "warmup_fraction": 0.01, + "lr_schedule_epochs": 1.0, + "weight_decay": 0.10, + }, + "muon_clip": { + "family": "muon_clip", + "learning_rate": 0.0002, + "min_learning_rate": 0.00002, + "warmup_fraction": 0.0512, + "lr_schedule_epochs": 1.0, + "newton_schulz_steps": 5, + "update_rms_scale": 0.20, + "qk_clip_threshold": 100.0, + }, + } + profiles = cfg["optimizer_profiles"] + for arm, expected in expected_profile_fields.items(): + assert arm in profiles + assert {key: profiles[arm][key] for key in expected} == expected + + assert tokens_per_step(cfg) == 8_192 + assert max_steps(cfg) == 39_063 + permanent_grid = epoch_step_map(cfg) + assert len(permanent_grid) == 17 + assert len(permanent_grid) >= 10 + assert list(permanent_grid) == sorted(set(permanent_grid)) + assert next(iter(permanent_grid)) == 0 + assert next(reversed(permanent_grid)) == max_steps(cfg) + assert list(permanent_grid.values()) == pytest.approx( + [index * 0.25 for index in range(17)] + ) + + assert cfg["weightwatcher"] == { + "enabled": True, + "ERG": True, + "randomize": True, + "strict": True, + "min_evals": 20, + "fix_fingers": "clip_xmax", + "max_fingers": 10, + "require_raw_alpha": True, + } + + statistics = report._stats([1.0, 2.0, 3.0, 4.0, 5.0]) + assert statistics["n"] == 5 + assert statistics["mean"] == pytest.approx(3.0) + assert statistics["ci95_half_width"] == pytest.approx( + report.T_975_DF4 / np.sqrt(5) * np.std([1, 2, 3, 4, 5], ddof=1) + ) + + +def test_dated_campaign_contains_reproduction_and_report_entrypoints(): + required = ( + EXPERIMENT_ROOT / "README.md", + EXPERIMENT_ROOT / "RESULTS.md", + CAMPAIGN_PATH, + CONFIG_PATH, + RUNNER_PATH, + REPORT_PATH, + NOTEBOOK_PATH, + DOCTOR_SMOKE_PATH, + ) + assert all(path.is_file() for path in required) + assert "--require-complete" in REPORT_PATH.read_text(encoding="utf-8") + assert "erg_gap_num_traps" in REPORT_PATH.read_text(encoding="utf-8") + assert "erg_gap_num_traps" in RUNNER_PATH.read_text(encoding="utf-8") + + +def test_adam_and_adamw_are_distinct_optimizer_arms(): + cfg = _baseline_config_document() + model_config = GPTConfig( + vocab_size=64, + block_size=8, + n_layer=1, + n_head=1, + n_embd=16, + ) + adam_model = GPT(model_config) + adamw_model = GPT(model_config) + adam_handles = make_optimizer_handles( + adam_model, + optimizer_profile(cfg, "adam"), + ) + adamw_handles = make_optimizer_handles( + adamw_model, + optimizer_profile(cfg, "adamw"), + ) + assert len(adam_handles) == len(adamw_handles) == 1 + assert type(adam_handles[0].optimizer) is torch.optim.Adam + assert type(adamw_handles[0].optimizer) is torch.optim.AdamW + assert all( + float(group["weight_decay"]) == 0.0 + for group in adam_handles[0].optimizer.param_groups + ) + assert any( + float(group["weight_decay"]) == pytest.approx(0.1) + for group in adamw_handles[0].optimizer.param_groups + ) + + +def test_campaign_root_is_explicit_resolved_and_strictly_below_tmp(): + runner = _load_runner() + resolve_experiment_root = runner.resolve_experiment_root + + with tempfile.TemporaryDirectory( + prefix="rg-campaign-root-test-", + dir="/tmp", + ) as temporary: + temporary_root = Path(temporary) + requested = temporary_root / "campaign-output" + assert resolve_experiment_root( + {"RG_NANOGPT_EXPERIMENT_ROOT": str(requested)} + ) == requested.resolve() + assert not requested.exists(), "path validation must not create output" + + escape = temporary_root / "escape" + escape.symlink_to(Path("/"), target_is_directory=True) + with pytest.raises(ValueError): + resolve_experiment_root( + {"RG_NANOGPT_EXPERIMENT_ROOT": str(escape / "campaign")} + ) + + invalid_environments = ( + {}, + {"RG_NANOGPT_EXPERIMENT_ROOT": ""}, + {"RG_NANOGPT_EXPERIMENT_ROOT": "relative/output"}, + {"RG_NANOGPT_EXPERIMENT_ROOT": "~/campaign-output"}, + {"RG_NANOGPT_EXPERIMENT_ROOT": "/tmp"}, + {"RG_NANOGPT_EXPERIMENT_ROOT": "/var/tmp/campaign-output"}, + { + "RG_NANOGPT_EXPERIMENT_ROOT": "/tmp/pretend-home/campaign", + "HOME": "/tmp/pretend-home", + }, + ) + for environment in invalid_environments: + with pytest.raises(ValueError): + resolve_experiment_root(environment) + + +def test_weightwatcher_uses_one_clip_xmax_call_for_both_alphas( + tmp_path, + monkeypatch, +): + analyze_calls: list[dict] = [] + + class FakeWeightWatcher: + def __init__(self, *, model): + self.model = model + + def analyze(self, **kwargs): + analyze_calls.append(dict(kwargs)) + names = [ + str(item["matrix_name"]) + for item in self.model.matrix_metadata + ] + count = len(names) + return pd.DataFrame( + { + "layer_id": range(1, count + 1), + "name": [f"holder.{name}" for name in names], + "alpha": [2.0 + index / 10 for index in range(count)], + "raw_alpha": [ + 2.5 + index / 10 for index in range(count) + ], + "num_fingers": [2] * count, + "ERG_gap": [0.25] * count, + "num_traps": [1] * count, + "rand_distance": [0.10] * count, + } + ) + + monkeypatch.setitem( + sys.modules, + "weightwatcher", + SimpleNamespace(WeightWatcher=FakeWeightWatcher), + ) + model = GPT( + GPTConfig( + vocab_size=64, + block_size=8, + n_layer=1, + n_head=1, + n_embd=16, + ) + ) + ww_config = _baseline_config_document()["weightwatcher"] + summary = run_weightwatcher( + model, + tmp_path, + step=7, + tokens_seen=4_096, + train_tokens=8_192, + config=ww_config, + seed=1337, + fingerprint="unit-weightwatcher-fingerprint", + ) + + assert analyze_calls == [ + { + "ERG": True, + "randomize": True, + "plot": False, + "min_evals": 20, + "fix_fingers": "clip_xmax", + "max_fingers": 10, + } + ] + assert summary["n_matrices"] == 6 + assert summary["alpha_raw_n"] == 6 + assert summary["alpha_clip_xmax_n"] == 6 + + raw = pd.read_csv( + tmp_path + / "spectral" + / "raw" + / "weightwatcher_step_0000007.csv" + ) + layers = pd.read_csv(tmp_path / "spectral" / "layers.csv") + for frame in (raw, layers): + assert len(frame) == 6 + assert frame["alpha"].equals(frame["alpha_clip_xmax"]) + assert frame["raw_alpha"].equals(frame["alpha_raw"]) + assert frame["alpha_delta"].tolist() == pytest.approx( + ( + frame["alpha_raw"] - frame["alpha_clip_xmax"] + ).tolist() + ) + assert (frame["weightwatcher_analysis_calls"] == 1).all() + assert set(frame["primary_alpha_variant"]) == {"clip_xmax"} + + +def test_source_notebook_is_valid_report_only_papermill_document(): + raw_document = json.loads(NOTEBOOK_PATH.read_text(encoding="utf-8")) + notebook = nbformat.from_dict(raw_document) + nbformat.validate(notebook) + + parameter_cells = [ + cell + for cell in notebook.cells + if cell.cell_type == "code" + and "parameters" in cell.metadata.get("tags", []) + ] + assert len(parameter_cells) == 1 + parameter_source = parameter_cells[0].source + parameter_tree = ast.parse( + "".join(parameter_source) + if isinstance(parameter_source, list) + else parameter_source + ) + parameters: dict[str, object] = {} + for statement in parameter_tree.body: + if ( + isinstance(statement, ast.Assign) + and len(statement.targets) == 1 + and isinstance(statement.targets[0], ast.Name) + and isinstance(statement.value, ast.Constant) + ): + parameters[statement.targets[0].id] = statement.value.value + assert parameters == { + "RESULTS_ROOT": "", + "OUTPUT_ROOT": "", + "REQUIRE_COMPLETE": True, + } + + trees = [ + ast.parse( + "".join(cell.source) + if isinstance(cell.source, list) + else cell.source + ) + for cell in notebook.cells + if cell.cell_type == "code" + ] + subprocess_calls = [ + node + for tree in trees + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "subprocess" + and node.func.attr == "run" + ] + assert len(subprocess_calls) == 1 + keywords = { + keyword.arg: keyword.value + for keyword in subprocess_calls[0].keywords + if keyword.arg is not None + } + assert isinstance(keywords.get("check"), ast.Constant) + assert keywords["check"].value is True + assert "shell" not in keywords + + source = "\n".join( + "".join(cell.source) + if isinstance(cell.source, list) + else cell.source + for cell in notebook.cells + if cell.cell_type == "code" + ) + for required_text in ( + "sys.executable", + "build_report.py", + "--results-root", + "--output-root", + "--require-complete", + "--allow-incomplete", + "SUMMARY.md", + "*.csv", + "*.png", + ): + assert required_text in source + for forbidden_training_entrypoint in ( + "rg-onehead-train", + "run_one(", + "run_optimizer_replicates", + "rg_nanogpt_one_head.training", + ): + assert forbidden_training_entrypoint not in source + + +def test_dated_launch_materials_do_not_expand_home_or_tilde_paths(): + launch_suffixes = {".md", ".py", ".sh", ".yaml", ".yml", ".ipynb"} + forbidden_patterns = { + "shell home expansion": re.compile(r"\$(?:\{HOME\}|HOME\b)"), + "tilde path": re.compile(r"(?=1.24"]), + "numpy": Distribution("numpy", "2.1.0", []), + } + monkeypatch.setattr( + runner.importlib.metadata, + "distribution", + lambda name: distributions[str(name).lower()], + ) + lock, requirements = runner._installed_distribution_lock( + {"python": sys.version.split()[0], "torch": "2.8.0"} + ) + assert lock["schema_version"] == 2 + assert set(lock["packages"]) == {"torch", "numpy"} + assert requirements.splitlines() == ["numpy==2.1.0", "torch==2.8.0"] + + +def test_dependency_contract_returns_the_manifest_closure(monkeypatch): + runner = _load_runner() + direct = {name: "test-version" for name in runner.DEPENDENCIES} + direct["rg-nanogpt-one-head"] = runner.PINNED_PACKAGE_VERSION + direct["weightwatcher"] = runner.PINNED_WEIGHTWATCHER + monkeypatch.setattr(runner, "_dependency_versions", lambda: direct) + monkeypatch.setattr( + runner.importlib.metadata, + "version", + lambda name: "not-installed" if name == "torch-xla" else "test-version", + ) + monkeypatch.setattr( + runner, + "_installed_distribution_lock", + lambda scientific: ( + { + "schema_version": 2, + "packages": { + "transitive-lib": { + "name": "transitive-lib", + "version": "9.8.7", + } + }, + "scientific_packages": dict(scientific), + }, + "transitive-lib==9.8.7\n", + ), + ) + contract = runner._require_dependency_contract() + assert contract["transitive-lib"] == "9.8.7" diff --git a/baseline/nanogpt_one_head/tests/test_completion.py b/baseline/nanogpt_one_head/tests/test_completion.py index d61f1cd6..16fd8d96 100644 --- a/baseline/nanogpt_one_head/tests/test_completion.py +++ b/baseline/nanogpt_one_head/tests/test_completion.py @@ -1,17 +1,21 @@ from __future__ import annotations from copy import deepcopy +import hashlib import json +import math from pathlib import Path import pandas as pd import pytest import torch +import rg_nanogpt_one_head.config as config_module from rg_nanogpt_one_head.completion import ( CompletedRunValidationError, validate_completed_run, ) +from rg_nanogpt_one_head.checkpoints import optimizer_state_sha256 from rg_nanogpt_one_head.config import ( load_config, max_steps, @@ -20,7 +24,12 @@ warmup_steps, ) from rg_nanogpt_one_head.engine import run_one -from rg_nanogpt_one_head.run_utils import run_directory, run_is_complete +from rg_nanogpt_one_head.run_utils import ( + model_state_sha256, + run_directory, + run_is_complete, + validate_existing_manifest_runtime, +) EXPERIMENT_ROOT = Path(__file__).resolve().parents[1] @@ -34,6 +43,47 @@ "W_MLP_IN", "W_MLP_OUT", ) +HELD_OUT_CURVE_COLUMNS = ( + "test_loss", + "test_perplexity", + "test_bits_per_token", + "test_accuracy", + "test_top5_accuracy", + "test_bleu", + "test_continuation_token_accuracy", + "test_continuation_exact_match", + "test_generalization_gap", +) + + +def _runtime_identity() -> dict: + return { + "platform": "unit-platform", + "machine": "unit-machine", + "python_version": "unit-python", + "accelerator": "cpu", + "device": "cpu", + "torch_version": torch.__version__, + "float32_matmul_precision": "highest", + "deterministic_algorithms": True, + "deterministic_warn_only": False, + "hardware_block_id": "unit-cpu-block", + "hardware_block_id_source": "test", + } + + +def _test_metrics(loss: float, *, step: int) -> dict: + return { + "step": int(step), + "loss": float(loss), + "perplexity": math.exp(float(loss)), + "bits_per_token": float(loss) / math.log(2.0), + "accuracy": 0.10, + "top5_accuracy": 0.25, + "bleu": 1.5, + "continuation_token_accuracy": 0.05, + "continuation_exact_match": 0.0, + } def _config() -> dict: @@ -111,6 +161,7 @@ def _write_completed_run(results_root: Path, cfg: dict) -> tuple[Path, str, int] metadata = _data_metadata(cfg) total_steps = max_steps(cfg, metadata["splits"]["train"]) profile = optimizer_profile(cfg, OPTIMIZER) + test_model_state = {"weight": torch.tensor([1.0])} fingerprint = protocol_fingerprint( cfg, optimizer=OPTIMIZER, @@ -119,10 +170,13 @@ def _write_completed_run(results_root: Path, cfg: dict) -> tuple[Path, str, int] ) run_dir = run_directory(results_root, OPTIMIZER, SEED) (run_dir / "spectral").mkdir(parents=True, exist_ok=True) + (run_dir / "spectral" / "raw").mkdir(parents=True, exist_ok=True) (run_dir / "epoch_checkpoints").mkdir(parents=True, exist_ok=True) best_step = 0 best_loss = 2.0 + final_test = _test_metrics(1.0, step=total_steps) + selected_test = _test_metrics(2.0, step=best_step) completion = { "completed": True, "optimizer": OPTIMIZER, @@ -130,6 +184,18 @@ def _write_completed_run(results_root: Path, cfg: dict) -> tuple[Path, str, int] "optimizer_steps": total_steps, "best_validation_step": best_step, "best_validation_loss": best_loss, + "final_test_loss": final_test["loss"], + "final_test_perplexity": final_test["perplexity"], + "final_test_bits_per_token": final_test["bits_per_token"], + "final_test_accuracy": final_test["accuracy"], + "final_test_top5_accuracy": final_test["top5_accuracy"], + "final_test_bleu": final_test["bleu"], + "final_test_continuation_token_accuracy": final_test[ + "continuation_token_accuracy" + ], + "final_test_continuation_exact_match": final_test[ + "continuation_exact_match" + ], "fingerprint": fingerprint, } (run_dir / "run_complete.json").write_text( @@ -141,6 +207,8 @@ def _write_completed_run(results_root: Path, cfg: dict) -> tuple[Path, str, int] "max_steps": total_steps, "warmup_steps": warmup_steps(profile, total_steps), "protocol_fingerprint": fingerprint, + "initial_model_sha256": model_state_sha256(test_model_state), + "runtime_environment": _runtime_identity(), } (run_dir / "manifest.json").write_text( json.dumps(manifest), encoding="utf-8" @@ -148,16 +216,21 @@ def _write_completed_run(results_root: Path, cfg: dict) -> tuple[Path, str, int] (run_dir / "test_results.json").write_text( json.dumps( { - "policy": "validation selects; test monitoring-only", - "final": {"step": total_steps, "loss": 1.0}, - "validation_selected": {"step": best_step, "loss": 2.0}, + "policy": "test is held out; validation selects and test never tunes", + "final": final_test, + "validation_selected": selected_test, } ), encoding="utf-8", ) + optimizer_states = [{"state": {}, "param_groups": []}] checkpoint_common = { - "schema_version": 2, + "schema_version": 5, + "model": test_model_state, + "optimizers": optimizer_states, + "model_state_sha256": model_state_sha256(test_model_state), + "optimizer_state_sha256": optimizer_state_sha256(optimizer_states), "fingerprint": fingerprint, "optimizer_name": OPTIMIZER, "seed": SEED, @@ -165,6 +238,7 @@ def _write_completed_run(results_root: Path, cfg: dict) -> tuple[Path, str, int] "best_validation_loss": best_loss, } for filename, step in ( + ("checkpoint_initial.pt", 0), ("checkpoint_latest.pt", total_steps), ("checkpoint_final.pt", total_steps), ("checkpoint_best.pt", best_step), @@ -173,15 +247,24 @@ def _write_completed_run(results_root: Path, cfg: dict) -> tuple[Path, str, int] steps = (0, total_steps) pd.DataFrame( - [{"step": step, "train_loss": 2.0 - 0.1 * index} - for index, step in enumerate(steps)] + [ + { + "step": step, + "train_loss": 2.0 - 0.1 * index, + **{column: float("nan") for column in HELD_OUT_CURVE_COLUMNS}, + } + for index, step in enumerate(steps) + ] ).to_csv(run_dir / "metrics.csv", index=False) pd.DataFrame( [ { "step": step, + "epoch": float(index), "nominal_epoch": float(index), "test_monitoring_only": 1, + "test_held_out": 1, + **{column: float("nan") for column in HELD_OUT_CURVE_COLUMNS}, "checkpoint_path": str( run_dir / "epoch_checkpoints" @@ -193,40 +276,89 @@ def _write_completed_run(results_root: Path, cfg: dict) -> tuple[Path, str, int] ).to_csv(run_dir / "epoch_metrics.csv", index=False) layer_rows = [] + model_hash = model_state_sha256(test_model_state) for step in steps: + diagnostic_seed = SEED + 1_000_003 + step + step_rows = [] for matrix in MATRICES: - layer_rows.append( + row = { + "step": step, + "matrix_name": matrix, + "alpha": 2.0, + "ERG_gap": 0, + "num_traps": 0, + "rand_distance": 0.0, + "run_seed": SEED, + "diagnostic_seed": diagnostic_seed, + "protocol_fingerprint": fingerprint, + "model_state_sha256": model_hash, + } + step_rows.append(row) + layer_rows.append(row) + raw_path = ( + run_dir + / "spectral" + / "raw" + / f"weightwatcher_step_{step:07d}.csv" + ) + pd.DataFrame(step_rows).to_csv(raw_path, index=False) + raw_hash = hashlib.sha256(raw_path.read_bytes()).hexdigest() + (run_dir / "spectral" / f"status_step_{step:07d}.json").write_text( + json.dumps( { + "completed": True, "step": step, - "matrix_name": matrix, - "alpha": 2.0, - "ERG_gap": 0, - "num_traps": 0, + "run_seed": SEED, + "diagnostic_seed": diagnostic_seed, + "protocol_fingerprint": fingerprint, + "model_state_sha256": model_hash, + "raw_csv_sha256": raw_hash, } - ) + ), + encoding="utf-8", + ) pd.DataFrame(layer_rows).to_csv( run_dir / "spectral" / "layers.csv", index=False ) summary_rows = [ - {"step": step, "n_matrices": len(MATRICES)} + { + "step": step, + "n_matrices": len(MATRICES), + "run_seed": SEED, + "diagnostic_seed": SEED + 1_000_003 + step, + "protocol_fingerprint": fingerprint, + "model_state_sha256": model_hash, + } for step in steps ] pd.DataFrame(summary_rows).to_csv( run_dir / "spectral" / "summary.csv", index=False, ) - # Keep the synthetic fixture's checkpoint inventory - # identical to the paths in epoch_metrics.csv. inventory = pd.read_csv(run_dir / "epoch_metrics.csv") - for value in inventory["checkpoint_path"]: - checkpoint_path = Path(str(value)) + for _, row in inventory.iterrows(): + checkpoint_path = Path(str(row["checkpoint_path"])) checkpoint_path.parent.mkdir( parents=True, exist_ok=True ) - if not checkpoint_path.is_file(): - checkpoint_path.write_bytes( - b"synthetic epoch checkpoint" - ) + torch.save( + { + "schema_version": 3, + "model": {"weight": torch.tensor([1.0])}, + "model_state_sha256": model_state_sha256( + {"weight": torch.tensor([1.0])} + ), + "step": int(row["step"]), + "nominal_epoch": float(row["nominal_epoch"]), + "actual_epoch": float(row["epoch"]), + "fingerprint": fingerprint, + "config": cfg, + "optimizer_name": OPTIMIZER, + "seed": SEED, + "purpose": "per_epoch_model_only_analysis_checkpoint", + }, + checkpoint_path, + ) return run_dir, fingerprint, total_steps @@ -257,9 +389,15 @@ def test_engine_validates_before_skipping_completed_run(tmp_path, monkeypatch): ) monkeypatch.setattr( "rg_nanogpt_one_head.engine.choose_device", - lambda *args, **kwargs: (_ for _ in ()).throw( - AssertionError("device setup should not run for a verified result") - ), + lambda *args, **kwargs: torch.device("cpu"), + ) + monkeypatch.setattr( + "rg_nanogpt_one_head.engine.configure_runtime", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "rg_nanogpt_one_head.engine.runtime_metadata", + lambda *args, **kwargs: _runtime_identity(), ) observed = run_one( cfg=cfg, @@ -288,6 +426,18 @@ def test_engine_rejects_completed_run_from_changed_protocol( {"train": object(), "val": object(), "test": object()}, ), ) + monkeypatch.setattr( + "rg_nanogpt_one_head.engine.choose_device", + lambda *args, **kwargs: torch.device("cpu"), + ) + monkeypatch.setattr( + "rg_nanogpt_one_head.engine.configure_runtime", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "rg_nanogpt_one_head.engine.runtime_metadata", + lambda *args, **kwargs: _runtime_identity(), + ) with pytest.raises( CompletedRunValidationError, match="fingerprint|optimizer_steps" ): @@ -335,3 +485,234 @@ def test_checkpoint_with_stale_fingerprint_is_rejected(tmp_path): expected_seed=SEED, expected_total_steps=total_steps, ) + + +def test_finite_checkpoint_model_bit_change_is_rejected(tmp_path): + cfg = _config() + run_dir, fingerprint, total_steps = _write_completed_run(tmp_path, cfg) + path = run_dir / "checkpoint_final.pt" + checkpoint = torch.load(path, map_location="cpu", weights_only=False) + checkpoint["model"]["weight"][0] += 0.125 + torch.save(checkpoint, path) + with pytest.raises(CompletedRunValidationError, match="model-state SHA-256"): + validate_completed_run( + run_dir, + expected_fingerprint=fingerprint, + expected_optimizer=OPTIMIZER, + expected_seed=SEED, + expected_total_steps=total_steps, + ) + + +def test_finite_checkpoint_optimizer_change_is_rejected(tmp_path): + cfg = _config() + run_dir, fingerprint, total_steps = _write_completed_run(tmp_path, cfg) + path = run_dir / "checkpoint_final.pt" + checkpoint = torch.load(path, map_location="cpu", weights_only=False) + checkpoint["optimizers"][0]["param_groups"].append({"lr": 0.5}) + torch.save(checkpoint, path) + with pytest.raises( + CompletedRunValidationError, + match="optimizer-state SHA-256", + ): + validate_completed_run( + run_dir, + expected_fingerprint=fingerprint, + expected_optimizer=OPTIMIZER, + expected_seed=SEED, + expected_total_steps=total_steps, + ) + + +def test_missing_initial_checkpoint_is_not_complete(tmp_path): + cfg = _config() + run_dir, fingerprint, total_steps = _write_completed_run(tmp_path, cfg) + (run_dir / "checkpoint_initial.pt").unlink() + with pytest.raises(CompletedRunValidationError, match="checkpoint_initial"): + validate_completed_run( + run_dir, + expected_fingerprint=fingerprint, + expected_optimizer=OPTIMIZER, + expected_seed=SEED, + expected_total_steps=total_steps, + ) + + +def test_permanent_checkpoint_payload_identity_is_validated(tmp_path): + cfg = _config() + run_dir, fingerprint, total_steps = _write_completed_run(tmp_path, cfg) + epoch_metrics = pd.read_csv(run_dir / "epoch_metrics.csv") + path = Path(str(epoch_metrics.iloc[-1]["checkpoint_path"])) + payload = torch.load(path, map_location="cpu", weights_only=False) + payload["step"] = int(payload["step"]) - 1 + torch.save(payload, path) + with pytest.raises(CompletedRunValidationError, match="step mismatch"): + validate_completed_run( + run_dir, + expected_fingerprint=fingerprint, + expected_optimizer=OPTIMIZER, + expected_seed=SEED, + expected_total_steps=total_steps, + ) + + +def test_nonfinite_test_metric_is_not_complete(tmp_path): + cfg = _config() + run_dir, fingerprint, total_steps = _write_completed_run(tmp_path, cfg) + path = run_dir / "test_results.json" + payload = json.loads(path.read_text(encoding="utf-8")) + payload["final"]["accuracy"] = float("nan") + path.write_text(json.dumps(payload), encoding="utf-8") + with pytest.raises(CompletedRunValidationError, match="non-finite"): + validate_completed_run( + run_dir, + expected_fingerprint=fingerprint, + expected_optimizer=OPTIMIZER, + expected_seed=SEED, + expected_total_steps=total_steps, + ) + + +def test_weightwatcher_raw_csv_hash_is_bound_to_status(tmp_path): + cfg = _config() + run_dir, fingerprint, total_steps = _write_completed_run(tmp_path, cfg) + path = ( + run_dir + / "spectral" + / "raw" + / "weightwatcher_step_0000000.csv" + ) + raw = pd.read_csv(path) + raw.loc[0, "alpha"] = 9.0 + raw.to_csv(path, index=False) + with pytest.raises(CompletedRunValidationError, match="integrity status"): + validate_completed_run( + run_dir, + expected_fingerprint=fingerprint, + expected_optimizer=OPTIMIZER, + expected_seed=SEED, + expected_total_steps=total_steps, + ) + + +def test_completion_test_metric_must_match_test_results(tmp_path): + cfg = _config() + run_dir, fingerprint, total_steps = _write_completed_run(tmp_path, cfg) + path = run_dir / "run_complete.json" + payload = json.loads(path.read_text(encoding="utf-8")) + payload["final_test_loss"] += 0.5 + path.write_text(json.dumps(payload), encoding="utf-8") + with pytest.raises(CompletedRunValidationError, match="does not match"): + validate_completed_run( + run_dir, + expected_fingerprint=fingerprint, + expected_optimizer=OPTIMIZER, + expected_seed=SEED, + expected_total_steps=total_steps, + ) + + +def test_training_curve_cannot_contain_held_out_test_outcome(tmp_path): + cfg = _config() + run_dir, fingerprint, total_steps = _write_completed_run(tmp_path, cfg) + path = run_dir / "metrics.csv" + metrics = pd.read_csv(path) + metrics.loc[metrics.index[-1], "test_loss"] = 1.25 + metrics.to_csv(path, index=False) + with pytest.raises(CompletedRunValidationError, match="leaks held-out"): + validate_completed_run( + run_dir, + expected_fingerprint=fingerprint, + expected_optimizer=OPTIMIZER, + expected_seed=SEED, + expected_total_steps=total_steps, + ) + + +def test_runtime_guard_rejects_hardware_change_before_artifact_mutation( + tmp_path, +): + run_dir = tmp_path / "run" + run_dir.mkdir() + marker = run_dir / "checkpoint_latest.pt" + marker.write_bytes(b"unchanged") + (run_dir / "manifest.json").write_text( + json.dumps({"runtime_environment": _runtime_identity()}), + encoding="utf-8", + ) + changed = {**_runtime_identity(), "accelerator": "cuda", "device": "cuda"} + changed.update( + { + "cuda_version": "12.8", + "cudnn_version": 9000, + "cuda_device_name": "NVIDIA H100", + "cuda_device_capability": [9, 0], + "cuda_device_count": 1, + "cuda_device_uuid": "GPU-test-h100", + "cuda_driver_version": "575.57", + "cuda_device_total_memory_bytes": 85_899_345_920, + "cuda_multi_processor_count": 132, + "cuda_nvidia_smi_memory_mib": 81_920, + "cuda_matmul_allow_tf32": False, + "cudnn_allow_tf32": False, + } + ) + with pytest.raises(RuntimeError, match="cross-runtime"): + validate_existing_manifest_runtime(run_dir, changed) + assert marker.read_bytes() == b"unchanged" + + +@pytest.mark.parametrize("manifest_text", ["{", "[]"]) +def test_runtime_guard_rejects_malformed_existing_manifest( + tmp_path, + manifest_text, +): + run_dir = tmp_path / "run" + run_dir.mkdir() + marker = run_dir / "checkpoint_latest.pt" + marker.write_bytes(b"unchanged") + (run_dir / "manifest.json").write_text( + manifest_text, + encoding="utf-8", + ) + with pytest.raises(RuntimeError, match="manifest"): + validate_existing_manifest_runtime(run_dir, _runtime_identity()) + assert marker.read_bytes() == b"unchanged" + + +def test_runtime_guard_rejects_substantive_artifacts_without_manifest(tmp_path): + run_dir = tmp_path / "run" + run_dir.mkdir() + marker = run_dir / "checkpoint_latest.pt" + marker.write_bytes(b"unchanged") + with pytest.raises(RuntimeError, match="without manifest"): + validate_existing_manifest_runtime(run_dir, _runtime_identity()) + assert marker.read_bytes() == b"unchanged" + + +def test_protocol_fingerprint_changes_with_scientific_dependencies(monkeypatch): + cfg = _config() + metadata = _data_metadata(cfg) + monkeypatch.setattr( + config_module, + "scientific_dependency_versions", + lambda: {"torch": "2.7.0", "weightwatcher": "0.7.7"}, + ) + first = protocol_fingerprint( + cfg, + optimizer=OPTIMIZER, + seed=SEED, + data_metadata=metadata, + ) + monkeypatch.setattr( + config_module, + "scientific_dependency_versions", + lambda: {"torch": "2.8.0", "weightwatcher": "0.7.7"}, + ) + second = protocol_fingerprint( + cfg, + optimizer=OPTIMIZER, + seed=SEED, + data_metadata=metadata, + ) + assert first != second diff --git a/baseline/nanogpt_one_head/tests/test_muonclip.py b/baseline/nanogpt_one_head/tests/test_muonclip.py index 746a1a34..e331d0d7 100644 --- a/baseline/nanogpt_one_head/tests/test_muonclip.py +++ b/baseline/nanogpt_one_head/tests/test_muonclip.py @@ -5,6 +5,7 @@ import subprocess import sys +import pandas as pd import pytest import torch import yaml @@ -43,11 +44,13 @@ def hidden_matrices(model: GPT) -> list[torch.nn.Parameter]: ] -def test_historical_launcher_remains_three_optimizer_reference() -> None: - # Importing the class does not mutate the historical launcher. The extension - # is installed only by rg-onehead-muonclip. +def test_launcher_exposes_optional_adam_without_installing_muonclip() -> None: + # Importing the class does not install MuonClip. Adam is now a first-class + # optional profile while the historical YAML still requires only its + # original SGD/AdamW/Muon profiles. assert SUPPORTED_OPTIMIZERS == ( "sgd_momentum", + "adam", "adamw", "muon", ) @@ -177,6 +180,48 @@ def test_qk_clip_is_noop_below_threshold() -> None: assert optimizer.last_diagnostics["min_gamma"] == pytest.approx(1.0) +def test_qk_diagnostic_rows_are_atomic_and_idempotent(tmp_path) -> None: + parameter = torch.nn.Parameter(torch.zeros(2, 2)) + + class DummyModel: + blocks = [] + + path = tmp_path / "muonclip_qk.csv" + optimizer = MuonClip( + [parameter], + model=DummyModel(), + lr=0.0, + momentum=0.0, + nesterov=False, + weight_decay=0.0, + update_rms_scale=0.2, + qk_clip_threshold=100.0, + diagnostics_path=path, + ) + values = { + "step": 500.0, + "threshold": 100.0, + "steps_in_interval": 500.0, + "head_observations": 500.0, + "active_heads": 1.0, + "active_fraction": 0.002, + "mean_max_logit": 10.0, + "max_logit": 20.0, + "mean_gamma": 0.99, + "min_gamma": 0.5, + } + optimizer.step_index = 500 + optimizer._write_diagnostics(values) + optimizer._write_diagnostics({**values, "max_logit": 25.0}) + optimizer.step_index = 1_000 + optimizer._write_diagnostics({**values, "step": 1_000.0}) + + frame = pd.read_csv(path) + assert frame["step"].tolist() == [500.0, 1_000.0] + assert frame.loc[frame["step"].eq(500.0), "max_logit"].item() == 25.0 + assert not path.with_suffix(".csv.tmp").exists() + + def test_attention_observation_matches_native_sdpa_output() -> None: model = small_model() attention = model.blocks[0].attn diff --git a/baseline/nanogpt_one_head/tests/test_muonclip_integration.py b/baseline/nanogpt_one_head/tests/test_muonclip_integration.py index a9539e17..73d5821c 100644 --- a/baseline/nanogpt_one_head/tests/test_muonclip_integration.py +++ b/baseline/nanogpt_one_head/tests/test_muonclip_integration.py @@ -90,7 +90,9 @@ def test_tiny_muonclip_training_writes_qk_diagnostics(tmp_path) -> None: 'document_disjoint_splits': True, 'dataset_name': cfg['dataset']['name'], 'dataset_config': cfg['dataset']['config'], + 'dataset_split': cfg['dataset'].get('split', 'train'), 'dataset_revision': cfg['dataset']['revision'], + 'eot_token': 0, 'files': files, }), encoding='utf-8') diff --git a/baseline/nanogpt_one_head/tests/test_muonclip_walk.py b/baseline/nanogpt_one_head/tests/test_muonclip_walk.py index 39991b9d..b926b42d 100644 --- a/baseline/nanogpt_one_head/tests/test_muonclip_walk.py +++ b/baseline/nanogpt_one_head/tests/test_muonclip_walk.py @@ -137,7 +137,9 @@ def test_tiny_walk_run_writes_loadable_append_only_artifacts(tmp_path) -> None: 'document_disjoint_splits': True, 'dataset_name': cfg['dataset']['name'], 'dataset_config': cfg['dataset']['config'], + 'dataset_split': cfg['dataset'].get('split', 'train'), 'dataset_revision': cfg['dataset']['revision'], + 'eot_token': 0, 'files': files, }), encoding='utf-8') diff --git a/baseline/nanogpt_one_head/tests/test_one_head.py b/baseline/nanogpt_one_head/tests/test_one_head.py index bd783bc2..13e9f59a 100644 --- a/baseline/nanogpt_one_head/tests/test_one_head.py +++ b/baseline/nanogpt_one_head/tests/test_one_head.py @@ -18,6 +18,7 @@ from rg_nanogpt_one_head.analysis import mean_ci95 from rg_nanogpt_one_head.checkpoints import ( load_training_checkpoint, + load_training_checkpoint_for_resume, save_training_checkpoint, ) from rg_nanogpt_one_head.config import ( @@ -34,6 +35,7 @@ from rg_nanogpt_one_head.evaluation import fixed_probe from rg_nanogpt_one_head.model import GPT, GPTConfig, transformer_matrix_items from rg_nanogpt_one_head.optimizers import make_optimizer_handles, optimizer_step +from rg_nanogpt_one_head.run_utils import truncate_muonclip_qk_after from rg_nanogpt_one_head.spectral import summarize_spectral_frame from rg_nanogpt_one_head.training import run_one @@ -52,6 +54,16 @@ def reference_config() -> dict: def tiny_config(optimizer: str) -> dict: cfg = deepcopy(reference_config()) + if optimizer == "adam": + adam = deepcopy(cfg["optimizer_profiles"]["adamw"]) + adam.update( + { + "display_name": "Adam", + "family": "adam", + "weight_decay": 0.0, + } + ) + cfg["optimizer_profiles"]["adam"] = adam cfg["dataset"].update( { "name": "unit/fineweb", @@ -134,7 +146,9 @@ def write_tiny_data(path: Path, cfg: dict) -> None: "document_disjoint_splits": True, "dataset_name": cfg["dataset"]["name"], "dataset_config": cfg["dataset"]["config"], + "dataset_split": cfg["dataset"].get("split", "train"), "dataset_revision": cfg["dataset"]["revision"], + "eot_token": 0, "files": files, } ), @@ -220,7 +234,7 @@ def test_model_inventory_and_all_optimizer_updates_are_finite(): "W_MLP_IN", "W_MLP_OUT", } - for optimizer_name in ("sgd_momentum", "adamw", "muon"): + for optimizer_name in ("sgd_momentum", "adam", "adamw", "muon"): candidate = GPT( GPTConfig( vocab_size=64, @@ -231,7 +245,11 @@ def test_model_inventory_and_all_optimizer_updates_are_finite(): ) ) handles = make_optimizer_handles( - candidate, optimizer_profile(cfg, optimizer_name) + candidate, + optimizer_profile( + tiny_config("adam") if optimizer_name == "adam" else cfg, + optimizer_name, + ), ) x = torch.randint(0, 64, (2, 8)) _, loss = candidate(x, x) @@ -256,6 +274,7 @@ def test_document_disjoint_writer_produces_exact_verified_splits(tmp_path): dataset_metadata={ "dataset_name": "unit/fineweb", "dataset_config": "unit", + "dataset_split": "train", "dataset_revision": "unit", }, progress_every_documents=0, @@ -307,6 +326,10 @@ def test_checkpoint_roundtrip_restores_optimizer_and_generator(tmp_path): name: tensor.detach().clone() for name, tensor in model.state_dict().items() } + previous_eval_snapshot = [ + parameter.detach().float().cpu().clone() + for parameter in model.parameters() + ] path = save_training_checkpoint( tmp_path / "checkpoint.pt", model=model, @@ -320,6 +343,12 @@ def test_checkpoint_roundtrip_restores_optimizer_and_generator(tmp_path): optimizer_name="adamw", seed=13, train_generator=generator, + resume_diagnostics={ + "previous_eval_snapshot": previous_eval_snapshot, + "last_grad_pre": 1.25, + "last_grad_post": 0.75, + "last_clipped": True, + }, ) for parameter in model.parameters(): parameter.data.zero_() @@ -333,6 +362,42 @@ def test_checkpoint_roundtrip_restores_optimizer_and_generator(tmp_path): assert restored == (7, 2.5, 6, 12.0) for name, tensor in model.state_dict().items(): assert torch.equal(tensor, before[name]) + restored_with_diagnostics = load_training_checkpoint_for_resume( + path, + model=model, + handles=handles, + expected_fingerprint="abc", + train_generator=generator, + ) + assert restored_with_diagnostics[:4] == (7, 2.5, 6, 12.0) + diagnostics = restored_with_diagnostics[4] + assert diagnostics is not None + assert diagnostics["last_grad_pre"] == pytest.approx(1.25) + assert diagnostics["last_grad_post"] == pytest.approx(0.75) + assert diagnostics["last_clipped"] is True + for restored_snapshot, expected_snapshot in zip( + diagnostics["previous_eval_snapshot"], + previous_eval_snapshot, + strict=True, + ): + assert torch.equal(restored_snapshot, expected_snapshot) + + +def test_muonclip_qk_resume_truncates_only_uncheckpointed_rows(tmp_path): + run_dir = tmp_path / "muon_clip" / "seed_13" + run_dir.mkdir(parents=True) + path = run_dir / "muonclip_qk.csv" + pd.DataFrame( + { + "step": [500.0, 1_000.0, 1_500.0], + "threshold": [100.0, 100.0, 100.0], + } + ).to_csv(path, index=False) + + truncate_muonclip_qk_after(run_dir, 1_000) + + retained = pd.read_csv(path) + assert retained["step"].tolist() == [500.0, 1_000.0] def test_spectral_summary_keeps_direct_trap_and_erg_fields(): @@ -363,7 +428,10 @@ def test_student_t_interval_is_run_level(): ) -@pytest.mark.parametrize("optimizer_name", ["sgd_momentum", "adamw", "muon"]) +@pytest.mark.parametrize( + "optimizer_name", + ["sgd_momentum", "adam", "adamw", "muon"], +) def test_tiny_cpu_training_writes_restart_and_epoch_artifacts( tmp_path, monkeypatch, @@ -374,13 +442,13 @@ def test_tiny_cpu_training_writes_restart_and_epoch_artifacts( results_root = tmp_path / "results" write_tiny_data(data_root, cfg) - monkeypatch.setattr( - "rg_nanogpt_one_head.train_loop.evaluate_bleu", - lambda *args, **kwargs: {"bleu": 0.0}, - ) monkeypatch.setattr( "rg_nanogpt_one_head.run_utils.evaluate_bleu", - lambda *args, **kwargs: {"bleu": 0.0}, + lambda *args, **kwargs: { + "bleu": 0.0, + "continuation_token_accuracy": 0.0, + "continuation_exact_match": 0.0, + }, ) monkeypatch.setattr( "rg_nanogpt_one_head.train_loop.run_weightwatcher", @@ -408,6 +476,7 @@ def test_tiny_cpu_training_writes_restart_and_epoch_artifacts( epoch_metrics = pd.read_csv(run_dir / "epoch_metrics.csv") assert len(epoch_metrics) >= 1 assert epoch_metrics["test_monitoring_only"].eq(1).all() + assert epoch_metrics["test_held_out"].eq(1).all() assert float(epoch_metrics.iloc[0]["primary_lr"]) == 0.0 assert all( Path(path).is_file() for path in epoch_metrics["checkpoint_path"] @@ -480,10 +549,12 @@ def test_notebooks_are_valid_and_expose_requested_metrics(): ), "05_muonclip_esd_clip_xmax.ipynb": ( "WeightMatrixHolder", - "watcher_standard.analyze", + "watcher.analyze", "fix_fingers='clip_xmax'", "max_fingers=MAX_FINGERS", - "watcher_standard.get_ESD", + "watcher.get_ESD", + "raw_alpha", + "weightwatcher_analysis_calls", "alpha_reduction", ), "06_first_layer_esd_binning_powerlaw.ipynb": (