Add an opt-in MuonClip nanoGPT baseline

[charlesmartin14]

Purpose

Add a dedicated one-head nanoGPT baseline for Moonshot AI's MuonClip optimizer while leaving the historical rg-onehead-train launcher, SGD/AdamW/Muon configs, default analysis tables, and existing result fingerprints unchanged.

MuonClip is implemented as:

Muon momentum
+ Newton--Schulz orthogonalization
+ decoupled weight decay
+ 0.2 * sqrt(max(fan_in, fan_out)) RMS-matched updates
+ per-head QK-Clip

For regular multi-head attention, the implementation measures the maximum finite causal pre-softmax QK logit over every gradient-accumulation micro-batch, computes

gamma_h = min(1, threshold / max_logit_h)

and applies the balanced Kimi-K2 scaling

W_Q[h] *= gamma_h ** 0.5
W_K[h] *= gamma_h ** 0.5

The reference threshold is 100.

Isolation from existing baselines

MuonClip is deliberately opt-in through:

rg-onehead-muonclip

Importing the optimizer class does not mutate the ordinary three-arm launcher. The dedicated command installs the extension in-process and delegates to the existing data, checkpoint, evaluation, WeightWatcher, MPS/CUDA/TPU, and long-horizon implementation.

Profiles

configs/muonclip_reference.yaml provides a matched one-epoch reference:

base LR:              2e-4
minimum LR:           2e-5
warmup:               500 / 9,766 steps
weight decay:         0.1
momentum:             0.95
Newton--Schulz:       5 steps
RMS scale:            0.2 * sqrt(max(n, m))
QK threshold:         100
Q/K balance:          0.5 / 0.5

configs/muonclip_10epochs.yaml adds a one-seed long spectral-relaxation run with 97,657 training steps, the matched 9,766-step LR horizon, 500-step warmup, and the 2e-5 floor held after epoch 1.

The compact baseline intentionally keeps the repository's matched warmup-cosine protocol. It does not claim to reproduce Kimi K2's 15.5T-token WSD schedule or distributed bfloat16 training.

QK diagnostics

Each MuonClip run writes muonclip_qk.csv at the configured interval with:

step
threshold
steps_in_interval
head_observations
active_heads
active_fraction
mean_max_logit
max_logit
mean_gamma
min_gamma

The diagnostics are monitoring-only and are never used for checkpoint selection or threshold tuning. Device values are aggregated without a host synchronization on every step; the host transfer occurs only at the diagnostic interval.

Accelerator behavior

• CPU/MPS/CUDA retain the existing native scaled-dot-product-attention output path; a detached QK matmul is used only to observe the clipping statistic.
• TPU/XLA uses the existing mathematical attention path, so the already-computed logits are observed directly.
• Checkpoints include MuonClip momentum and partial diagnostic interval state and remain portable through the existing CPU checkpoint serialization.

Validation

The new tests verify:

• the ordinary launcher still exposes exactly the historical three optimizers;
• the reference parameters and 500-step warmup;
• exact 0.2 * sqrt(max(n,m)) update scaling;
• balanced Q/K clipping and the below-threshold no-op;
• equality of the tracked CPU attention output and the native SDPA output;
• isolated installation and config validation through the dedicated extension.