Add nanoGPT Muon-HyperBall baseline

[charlesmartin14]

Purpose

Add a separate one-head nanoGPT experiment family for a matched comparison of ordinary long-horizon Muon and Muon with a relative Frobenius HyperBall projection. Existing baseline/nanogpt_one_head source and results are untouched.

HyperBall rule

For each hidden transformer matrix, reference Muon first proposes its complete displacement, including matrix weight decay. The new arm applies

[
\Delta W_{HB}=\Delta W\min\left(1,\frac{\rho\lVert W\rVert_F}{\lVert\Delta W\rVert_F+\epsilon}\right),
]

with reference radius rho = 0.01. The projection is radial, so it preserves Muon's proposed direction and changes only its magnitude. The tied embedding/head, normalization parameters, and other non-hidden parameters retain auxiliary AdamW.

Long-horizon schedule correction

The failed ten-epoch pilot stretched the original five-percent warmup over the full ten-epoch horizon. Muon's LR therefore kept increasing until epoch 0.5 and the run became non-finite near step 4250. This experiment separates training and LR-schedule horizons:

• ten corpus-equivalent training epochs;
• original 488-step warmup;
• original cosine decay completed at epoch 1;
• LR floors held through epochs 1–10.

Both muon and muon_hyperball use this exact schedule, data, seed, architecture, evaluation probes, and WeightWatcher cadence.

Diagnostics

The new runtime records HyperBall activation fraction, mean/min projection scale, proposed/applied update-to-weight ratios, proposed/applied update norms, and radius statistics. It also aborts on non-finite train/validation metrics before invoking WeightWatcher, preventing optimizer divergence from surfacing later as an SVD-convergence error.

Contents

• separate baseline/nanogpt_muon_hyperball package and /tmp/rg-nanogpt-muon-hyperball result namespace
• matched muon and muon_hyperball profiles
• restart-safe one-head FineWeb-Edu training runtime
• run and comparison notebooks
• dedicated CI workflow
• tests for the Frobenius cap, radial direction preservation, infinite-radius equivalence to Muon, six-matrix partition, schedule horizon, configuration validation, and persisted diagnostics

Validation

• local focused test suite: 7 passed
• Python compilation of the new source/tests: passed
• notebook JSON/code-cell parsing: passed

The full ten-epoch scientific run has not been executed in CI.