Use unchanged upstream GPT-2 with MuonClip and AdamW

.github/workflows/baseline-tests.yml

@@ -71,3 +71,14 @@ jobs:
           PYTHONPATH: baseline/nanogpt_one_head/src
           MPLBACKEND: Agg
         run: pytest -q baseline/nanogpt_one_head/tests
+
+      - name: Verify upstream GPT-2, MuonClip and AdamW integration
+        env:
+          PYTHONPATH: baseline/nanogpt_one_head/src
+          MPLBACKEND: Agg
+        run: >-
+          pytest -q baseline/gpt2_small/tests/test_stock_gpt2.py
+          baseline/gpt2_small/tests/test_repeated_speedrun.py
+          baseline/gpt2_small/tests/test_muon_speedrun.py
+          baseline/gpt2_small/tests/test_speedrun30.py
+          baseline/gpt2_small/tests/test_muon_longrun.py

README.md

@@ -5,6 +5,18 @@ spectral renormalization-group program. The repository keeps the unmodified
 reference baselines separate from every RG intervention so optimizer claims can
 be tested against strong, restartable, statistically controlled experiments.
 
+## nanoGPT speedrun: repeated-seed comparison
+
+Use `python3 baseline/gpt2_small/speedrun.py plan` for the fixed-recipe
+MuonClip/AdamW comparison using the **unchanged upstream GPT-2 Small (124,439,808 parameters)**:
+12 blocks, 12 heads, width 768, context 1024. Three matched seeds per optimizer, 19,560 updates each,
+raw/clipped WeightWatcher spectra and validation token error every 250 updates.
+The [pinned original GPT-2/FineWeb baseline](baseline/gpt2_small/muon_speedrun/BENCHMARK.md) uses 700 warmup updates, cosine decay and global gradient clipping at 1.0; MuonClip is an explicit optimizer substitution.
+[Protocol, launch instructions and seed-level statistics](baseline/gpt2_small/muon_speedrun/REPEATED_SEEDS.md).
+[Architecture audit and every matrix dimension](baseline/gpt2_small/STOCK_ARCHITECTURE.md).
+The current paired runs use the standard architecture; the historical modified
+six-head speedrun and one-head experiments remain separately identified.
+
 ## Baseline status
 
 The baseline suite has completed a recipe audit, an executable audit, and a

baseline/README.md

@@ -382,3 +382,8 @@ GitHub Actions additionally compiles all Python sources, parses notebook code
 cells, runs a real pinned WeightWatcher integration, and executes a pinned
 nanochat CPU model/optimizer preflight. These bounded checks do not replace the
 full three-seed long-horizon campaigns or the required target-MPS preflight.
+
+## GPT-2 Small
+
+See [gpt2_small/README.md](gpt2_small/README.md) for the 124M-parameter,
+12-layer, 12-head FineWeb TPU validation workflow.

baseline/gpt2_small/STOCK_ARCHITECTURE.md

@@ -0,0 +1,72 @@
+# Unmodified upstream GPT-2 Small architecture
+
+Both MuonClip and AdamW instantiate `GPT` directly from the unchanged
+`speedrun30/vendor/llmc_train_gpt2.py`, pinned to llm.c commit
+`7ecd8906afe6ed7a2b2cdb731c042f26d525b820`. The adapter checks the complete file's
+SHA256 before import. No model class, layer or forward method is rewritten.
+The architecture identifier is `gpt2-small-upstream-packed-v2`.
+
+| Property | Value |
+|---|---|
+| Blocks / heads / hidden width | 12 / 12 / 768 |
+| Head width / MLP width | 64 / 3072 |
+| Context / vocabulary | 1024 / 50257 |
+| Positions | Learned absolute embeddings |
+| Normalization | Pre-LayerNorm, affine scales and biases, epsilon 1e-5 |
+| Activation | Original GPT-2 tanh GELU |
+| Linear biases | Enabled; vocabulary head has no bias |
+| Output head | Tied to token embedding |
+| Dropout | 0, as in the upstream pretraining model |
+| Unique trainable parameters | **124,439,808** |
+| Single-run initialization | Upstream initializer, seed 42 |
+
+## Stored layer weight matrices
+
+Shapes use PyTorch `[output features, input features]` storage. **QKV is packed**
+in the upstream `attn.c_attn.weight` parameter. The following four matrices occur
+in each of the twelve blocks, indexed 0–11:
+
+| Block | QKV `attn.c_attn` | O `attn.c_proj` | MLP IN `mlp.c_fc` | MLP OUT `mlp.c_proj` |
+|---|---|---|---|---|
+| L00 | 2304 × 768 | 768 × 768 | 3072 × 768 | 768 × 3072 |
+| L01 | 2304 × 768 | 768 × 768 | 3072 × 768 | 768 × 3072 |
+| L02 | 2304 × 768 | 768 × 768 | 3072 × 768 | 768 × 3072 |
+| L03 | 2304 × 768 | 768 × 768 | 3072 × 768 | 768 × 3072 |
+| L04 | 2304 × 768 | 768 × 768 | 3072 × 768 | 768 × 3072 |
+| L05 | 2304 × 768 | 768 × 768 | 3072 × 768 | 768 × 3072 |
+| L06 | 2304 × 768 | 768 × 768 | 3072 × 768 | 768 × 3072 |
+| L07 | 2304 × 768 | 768 × 768 | 3072 × 768 | 768 × 3072 |
+| L08 | 2304 × 768 | 768 × 768 | 3072 × 768 | 768 × 3072 |
+| L09 | 2304 × 768 | 768 × 768 | 3072 × 768 | 768 × 3072 |
+| L10 | 2304 × 768 | 768 × 768 | 3072 × 768 | 768 × 3072 |
+| L11 | 2304 × 768 | 768 × 768 | 3072 × 768 | 768 × 3072 |
+
+| Other weight | Shape |
+|---|---|
+| Token embedding `transformer.wte.weight` | 50257 × 768 |
+| Position embedding `transformer.wpe.weight` | 1024 × 768 |
+| Vocabulary output `lm_head.weight` | 50257 × 768; alias of token embedding |
+
+There are **48 stored block matrices**, two embedding matrices and one named tied
+output alias: **50 unique matrix parameters, 51 named entries**. The full machine
+readable list is [stock_weight_matrices.csv](stock_weight_matrices.csv).
+Each block also has LayerNorm scales/biases of length 768, packed QKV bias 2304,
+attention output bias 768, MLP input bias 3072 and output bias 768. Final LayerNorm
+has scale and bias vectors of length 768. Each block contains 7,087,872 parameters.
+The upstream causal-mask buffers are not learned parameters.
+
+Q, K and V are each a 768 × 768 **slice** of the packed 2304 × 768 matrix; a single
+head occupies 64 × 768 rows. WeightWatcher extracts these slices only from saved
+CPU weights, retaining 72 projection traces without changing trainable storage.
+
+MuonClip observes QK logits without modifying the forward output. Its parameter
+and bias rescaling occurs in the optimizer update. TPU attention selection and
+BF16 autocast live outside the upstream model. The worker requires attention
+parity plus a full-sized accumulated optimizer preflight before fresh training.
+
+[Configuration and provenance](muon_speedrun/BENCHMARK.md) ·
+[Launch instructions](muon_speedrun/README.md) ·
+[Paired-seed protocol](muon_speedrun/REPEATED_SEEDS.md)
+
+Historical six-head modified-model runs and the earlier split-QKV port have
+different architecture/protocol identifiers and are excluded from new comparisons.

baseline/gpt2_small/TPU_PORT_BUGS.md

@@ -0,0 +1,183 @@
+# TPU port bug log
+
+Finding and isolating correctness/performance defects in the PyTorch-to-TPU port is
+an explicit objective of this project. Preserve failures as evidence. A passing
+CPU smoke test does not establish TPU correctness. Do not silently lower learning
+rates, change data, or skip nonfinite checks to make a run appear successful.
+
+## 2026-10-04: Muon speedrun missing Pallas dependency and HBM exhaustion
+
+Status: two configuration failures confirmed; corrected retry requires live TPU validation.
+These observations do not establish an upstream PyTorch/XLA or hardware bug.
+
+- Run `muon-speedrun-muon-20261004-225635`, commit `1b4548d`, v5litepod-8.
+- Flash check failed with `ModuleNotFoundError: No module named 'jax'`.
+  `torch_xla[tpu]` does not include the Pallas extras; XLA 2.6 setup.py pins
+  both JAX and jaxlib to 0.4.38 for that optional dependency group.
+- Automatic math-attention fallback kept global microbatch 128 (16/chip).
+  First forward/backward compilation failed with `RESOURCE_EXHAUSTED`:
+  16.89G required versus 15.75G HBM, exceeding capacity by 1.14G.
+- No completed training update. Step-zero checkpoint and final cloud backup
+  were saved. Disk evidence:
+  `/mnt/disks/rg-data/gpt2small/muon-speedrun-muon-20261004-225635`;
+  cloud prefix: `gs://tpu-builders-504820-ww-continuous8/gpt2small/muon-speedrun-muon-20261004-225635`.
+- Correction: per-run pinned Pallas overlay; flash forward/backward check required;
+  no implicit math fallback; microbatch 64, eight accumulation passes, unchanged
+  global batch 524,288; early checkpoints after updates 1 and 5. Failure status now
+  replaces stale `training` status and includes the actual exception.
+- The smaller microbatch and kernel must still pass a live full-model run. Do not
+  report the OOM fixed solely because CPU tests or attention-only checks pass.
+
+## 2026-10-04: GPT-2 AdamW nonfinite result, followed by stalled failure reporting
+
+Status: numerical cause open; blocking diagnostic implementation replaced.
+Upstream attribution: **unconfirmed**. No upstream issue has been submitted.
+
+- Run: `ww-gpt2-validation-48h-20261004-s1337`, original commit `4631a3d`.
+- Machine: one v5litepod-8, eight chips, SPMD; project `tpu-builders-504820`,
+  zone `us-west4-a`. Reported runtime: Python 3.10.12, PyTorch 2.6.0+cpu,
+  PyTorch/XLA 2.6.0. The library's `+cpu` build string does not identify where
+  the model computations executed; the run explicitly selected XLA/TPU.
+- Model: GPT-2 Small, 124,439,808 parameters; 12 layers, 12 heads, width 768,
+  context 1024, tied embeddings. FineWeb reused from the persistent disk.
+- Before updates 1 and 2, reported losses were finite (approximately 11.01 and
+  10.24) and aggregate gradient norms were 16.177856 and 7.746079.
+- After approximately 2.5 hours, the four-update check had not completed.
+  Latest checkpoint pointer: step 0. Process 6802 had roughly 212 GiB RSS.
+- Saved Python and native stacks identify `require_finite_update`, line 141,
+  at `p.grad.detach().float().cpu()`. This branch executes only after detecting
+  a nonfinite loss or aggregate gradient norm. It does not reveal which scalar
+  failed, the first affected matrix, or the numerical root cause.
+- Native frames include `THPVariable_cpu` and tensor conversion. They confirm
+  waiting in the host transfer path; they do not prove a hardware failure,
+  compiler bug, deadlock, or that all training ran on CPU.
+- Service subsequently confirmed `MainPID=0`, `ActiveState=inactive`,
+  `SubState=dead`. TPU allocation and data remain available.
+- Evidence on disk:
+  `/mnt/disks/rg-data/gpt2small/ww-gpt2-validation-48h-20261004-s1337/diagnostics/stall-20261004-041944-714308`.
+
+### Reporting defect and correction
+
+The old failure handler copied full gradients to CPU, serially, before writing
+the failure report. That code could stall and hide the already detected failure.
+Commit `f887593` removed these copies, added an execution barrier before host
+scalar reads, saved the scalar failure immediately, and bounded validation phases.
+These changes have local test coverage; they have not established numerical
+correctness on TPU. GPT-2 now has its own `baseline/gpt2_small` package.
+
+### Next diagnostic and attribution criteria
+
+Run four fresh AdamW updates on the same allocation, with the original model,
+corpus, seed and optimizer hyperparameters. Save each completed update. Record:
+
+- Source commit, Python/torch/torch_xla/libtpu versions, relevant XLA settings.
+- Exact input-window offsets, corpus identities, initialization/rolling full states.
+- Per-tensor finite flags and extrema before clipping, after clipping, and after
+  the optimizer update; parameter names plus Adam moment names.
+- XLA compilation/execution counters, fallback counters, stage timestamps,
+  Python tracebacks and a structured failure report.
+
+Checks reduce tensors on device and transfer a small summary table, never full
+gradients. Additional synchronization is recorded as instrumentation: it can change
+fusion/compilation behavior. A pass under instrumentation does not by itself clear
+the original execution path. The diagnostic stops after four updates or 20 minutes;
+up to 10 additional minutes are reserved for verified cloud backup. Disk evidence
+remains if upload fails. No long experiment starts automatically.
+
+To attribute an upstream bug, isolate the first failing operation and compare a
+matched CPU/TPU replay with the same inputs, weights and optimizer state. Preserve
+both results and a minimal reproducer before claiming a PyTorch/XLA defect.
+
+## 2026-10-04: instrumented diagnostic aborts at its first gradient check
+
+Status: open; fatal native message still required for diagnosis.
+
+Further stack evidence localizes this abort to `port_debug.py:68`: the new
+diagnostic's `torch.stack((isfinite(value).all().float(), value.amin(), value.amax()))`.
+Native frames include `torch_xla::Stack::Stack`, `XlaNode::GetOpShape`, and
+`XLANativeFunctions::stack`. The abort occurs while assembling diagnostic summaries,
+before the first optimizer update. It therefore does not reproduce or explain the
+earlier nonfinite-result failure. CPU tests passed this operation, but TPU behavior
+has not passed validation. The preceding native assertion/status text is still
+needed; the Python abort trace alone is insufficient to identify its cause.
+
+- Run: `port-check-20261004-045633`, commit `19e2bb0`.
+- Supervisor report: child exit code `-6` (SIGABRT), with last recorded stage
+  `before_clipping_started`, update 1, Unix time `1791089838.8617651`.
+- No per-tensor results from this check were reported. This abort does not by
+  itself establish a nonfinite gradient, a particular failing operator, or an
+  upstream runtime/hardware bug. It is a separate observed failure from the
+  earlier numerical check and stalled host copy.
+- Cloud backup was explicitly verified for this run. Logs, initialization,
+  exact input-window offsets and environment metadata remain on disk and in
+  its cloud prefix. The next evidence to inspect is the fatal native message
+  immediately preceding the abort in `run.log`.
+- A later launch was blocked by an untracked repository-root `FETCH_HEAD`
+  file in Cloud Shell. This local checkout issue is separate from the TPU
+  abort; moving that file outside the checkout preserves it and clears this
+  particular cleanliness check. The real Git metadata is under `.git`.
+
+### User-authorized MuonClip bypass, 2026-10-04
+
+`run_muonclip.py` starts a separate continuous MuonClip experiment on the remaining
+48-hour allocation. Its config sets `validation_tensor_checks=false` and
+`validation_gradient_checks=false`, bypassing the crashing per-tensor stack and its
+verbose validation path. `finite_update_guard=true` still checks scalar losses and
+aggregate norm before the optimizer update; the existing numerical guard is not
+removed. Model/data/optimizer and long-run LR settings are unchanged. Additional
+synchronization, reporting and denser checkpoint/spectral measurements are explicit.
+
+The workaround has CPU integration coverage; TPU success is not claimed. It is not
+a repair or root-cause diagnosis for either SIGABRT or the earlier nonfinite result.
+The earlier source checkouts, diagnostics, initialization checkpoints and verified
+cloud archives are retained. New failures produce separate evidence in the new run.
+No automatic restart, cleanup, disk formatting or TPU reallocation is performed.
+
+### MuonClip update-2 evaluation failure, 2026-10-04
+
+Run `muonclip-night-20261004-053212`, commit
+`6f8d59214f51be00212a9956182a0c697225f086`, failed about ten minutes after launch;
+it did not train overnight. The first update and its evaluation completed. Step 1
+was saved locally and its cloud checkpoint upload was verified. The final cloud
+backup was also explicitly verified.
+
+- Before update 2, four microbatch losses were finite (approximately 10.99) and
+  the aggregate gradient norm was `14.34039306640625`.
+- The log then reported `completed_update: 2`, followed by `RuntimeError:
+  Nonfinite train NLL` in evaluation. The saved progress stage was `evaluating`,
+  completed step 2. The latest saved checkpoint was step 1.
+- The recorder evaluated before saving the checkpoint, so the failed update's
+  weights and optimizer state were not preserved. Neither a corrupt update nor
+  an evaluation/runtime fault has yet been isolated. Finite pre-update losses
+  and norm do not prove finite updated parameters or optimizer moments.
+- Evidence remains under `/mnt/disks/rg-data/gpt2small/muonclip-night-20261004-053212`
+  and `gs://tpu-builders-504820-ww-continuous8/gpt2small/muonclip-night-20261004-053212`.
+
+The new `replay_muonclip.py` diagnostic restores the saved step-1 state and exact
+next input windows without changing optimizer settings. It checks saved/restored
+state, fixed evaluation probes, pre/post-clipping gradients, and state after each
+of primary MuonClip and auxiliary AdamW. It saves the resulting diagnostic state
+before evaluation. Per-tensor finite reductions use scalar host transfers without
+the earlier stack/extrema diagnostic. No attribution to upstream PyTorch/XLA or
+TPU hardware is justified yet. These synchronization changes are recorded; a pass
+does not by itself reproduce or fix the original continuous execution path.
+
+### Saved-state TPU replay passed, 2026-10-04
+
+The user supplied the final results for `muonclip-replay-20261004-150734` at commit
+`1b039b39761b13184f1cb09585b2e1e4343f8474`: `one_update_passed`, child exit 0.
+All recorded checks passed and the cloud backup was explicitly verified.
+Post-update NLL: train `10.98291015625`, validation `10.991303443908691`,
+test `10.969128131866455`. Pre-update train/test NLL reproduced the prior saved
+checkpoint's evaluation. This localizes neither the original fault nor a fix:
+additional synchronization, reductions, process state and checkpoint restoration
+differ from the failing run.
+
+The next user-authorized continuous run retains these synchronization/reduction
+boundaries on the existing PyTorch/XLA environment, starting from initialization.
+The shared `execution_checks.py` implementation is used by both replay and training.
+Checkpointing now precedes evaluation and spectra, with a pending-measurement flag
+for explicit recovery. CPU tests cover 25 continuous updates and exact state
+recovery after injected measurement failures. No continuous TPU success is claimed
+until its output is inspected. TorchTPU migration awaits repository/package access;
+the installed environment and all previous evidence remain intact.

baseline/gpt2_small/configs/gpt2_small_cpu_smoke.yaml

@@ -0,0 +1,56 @@
+run_id: tiny_cpu_smoke
+seed: 1337
+model:
+  vocab_size: 64
+  block_size: 8
+  n_layer: 2
+  n_head: 2
+  n_embd: 32
+  dropout: 0.0
+  bias: true
+  tie_weights: true
+dataset:
+  name: HuggingFaceFW/fineweb-edu
+  config: sample-10BT
+  split: train
+  revision: 593b3a867298afb8ce42625a270ef20ddcad28f9
+  tokenizer: gpt2
+  encoding_workers: 16
+  encoding_batch_size: 256
+  train_tokens: 512
+  val_tokens: 128
+  test_tokens: 128
+runtime:
+  matmul_precision: highest
+  mps_fallback: true
+  deterministic_algorithms: false
+  empty_mps_cache_after_weightwatcher: true
+  tpu_spmd: false
+  tpu_expected_chips: 8
+training:
+  batch_size: 2
+  grad_accum_steps: 1
+  max_steps: 4
+  max_tokens: 64
+  warmup_steps: 1
+  schedule_steps: 4
+  grad_clip: 1.0
+eval_batches: 1
+metrics_interval: 2
+ww:
+  enabled: false
+  interval: 0
+  steps: []
+  logarithmic: false
+  min_evals: 20
+optimizer:
+  display_name: AdamW
+  family: adamw
+  learning_rate: 0.0006
+  min_learning_rate: 6.0e-05
+  warmup_fraction: 0.01
+  schedule: warmup_cosine
+  beta1: 0.9
+  beta2: 0.95
+  epsilon: 1.0e-08
+  weight_decay: 0.1