From d53dd0ed0729daac35bc46d3f9799f5a1cf7fcb1 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 26 Aug 2026 22:06:17 -0700 Subject: [PATCH 1/5] Allow configurable attention-head count --- .../nanogpt_one_head/src/rg_nanogpt_one_head/model.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py index af1c3447..d40fd384 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py @@ -20,11 +20,10 @@ class GPTConfig: tie_weights: bool = True def __post_init__(self) -> None: - if self.n_layer != 1 or self.n_head != 1: - raise ValueError( - "the reference architecture is fixed to one block and one " - "attention head" - ) + if self.n_layer != 1: + raise ValueError("the reference architecture is fixed to one block") + if self.n_head < 1: + raise ValueError("n_head must be positive") if self.n_embd % self.n_head != 0: raise ValueError("n_embd must be divisible by n_head") if self.block_size < 2 or self.vocab_size < 2 or self.n_embd < 1: From 37dc062ca99e967e06b7cbbff860f8d209561c02 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 26 Aug 2026 22:06:18 -0700 Subject: [PATCH 2/5] Validate multi-head one-block configurations --- baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py | 2 -- 1 file changed, 2 deletions(-) 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 29f9643b..2499d208 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 @@ -225,8 +225,6 @@ def validate_config(cfg: dict[str, Any]) -> None: for key in ("vocab_size", "block_size", "n_layer", "n_head", "n_embd"): if int(model[key]) < 1: raise ValueError(f"model.{key} must be positive") - if int(model["n_head"]) != 1: - raise ValueError("this experiment is fixed to exactly one attention head") if int(model["n_layer"]) != 1: raise ValueError("this experiment is fixed to one transformer block") if int(model["n_embd"]) % int(model["n_head"]) != 0: From a5c1d02bdab71c1f7fa37d5074e11a1280cc9ff9 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 26 Aug 2026 22:06:19 -0700 Subject: [PATCH 3/5] Test matched four-head architecture and config --- .../nanogpt_one_head/tests/test_one_head.py | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/baseline/nanogpt_one_head/tests/test_one_head.py b/baseline/nanogpt_one_head/tests/test_one_head.py index 13e9f59a..54c24333 100644 --- a/baseline/nanogpt_one_head/tests/test_one_head.py +++ b/baseline/nanogpt_one_head/tests/test_one_head.py @@ -183,6 +183,54 @@ def test_reference_protocol_is_one_block_one_head_and_has_required_ww_flags(): ] +def test_four_head_configuration_preserves_width_and_parameter_count(): + one_head = GPT( + GPTConfig( + vocab_size=64, + block_size=8, + n_layer=1, + n_head=1, + n_embd=16, + ) + ) + four_head = GPT( + GPTConfig( + vocab_size=64, + block_size=8, + n_layer=1, + n_head=4, + n_embd=16, + ) + ) + tokens = torch.randint(0, 64, (2, 8)) + logits, loss = four_head(tokens, tokens) + + assert four_head.blocks[0].attn.n_head == 4 + assert four_head.blocks[0].attn.n_embd // four_head.blocks[0].attn.n_head == 4 + assert four_head.parameter_count() == one_head.parameter_count() + assert logits.shape == (2, 8, 64) + assert loss is not None and torch.isfinite(loss) + + +def test_four_head_baseline_matches_one_head_campaign_except_head_count(): + four_head = load_config( + EXPERIMENT_ROOT / "configs" / "four_head_4epoch_baseline.yaml" + ) + + assert four_head["model"]["n_layer"] == 1 + assert four_head["model"]["n_head"] == 4 + assert four_head["model"]["n_embd"] == 128 + assert four_head["training"]["target_epochs"] == 4.0 + assert four_head["training"]["seeds"] == [ + 1337, + 2027, + 4099, + 31415, + 271828, + ] + assert max_steps(four_head) == 39_063 + + def test_probe_identity_is_fixed_across_training_seeds(tmp_path): data = np.memmap( tmp_path / "probe.bin", From 2e46efd028fbb012f002568ef27555ba928d4f15 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 26 Aug 2026 22:06:20 -0700 Subject: [PATCH 4/5] Test four-head MuonClip QK clipping --- .../nanogpt_one_head/tests/test_muonclip.py | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/baseline/nanogpt_one_head/tests/test_muonclip.py b/baseline/nanogpt_one_head/tests/test_muonclip.py index e331d0d7..12c65502 100644 --- a/baseline/nanogpt_one_head/tests/test_muonclip.py +++ b/baseline/nanogpt_one_head/tests/test_muonclip.py @@ -152,6 +152,49 @@ def test_qk_clip_balances_query_and_key_scaling() -> None: assert optimizer.last_diagnostics["min_gamma"] == pytest.approx(0.25) +def test_qk_clip_scales_four_heads_independently() -> None: + model = GPT( + GPTConfig( + vocab_size=64, + block_size=8, + n_layer=1, + n_head=4, + n_embd=16, + dropout=0.0, + bias=False, + ) + ) + optimizer = MuonClip( + hidden_matrices(model), + model=model, + lr=0.0, + momentum=0.0, + nesterov=False, + weight_decay=0.0, + update_rms_scale=0.2, + qk_clip_threshold=100.0, + qk_clip_balance=0.5, + diagnostics_interval=1, + ) + q_before = model.blocks[0].attn.q_proj.weight.detach().clone().view(4, 4, 16) + k_before = model.blocks[0].attn.k_proj.weight.detach().clone().view(4, 4, 16) + model.blocks[0].attn._muonclip_max_logits = torch.tensor( + [25.0, 100.0, 400.0, 1_600.0] + ) + for parameter in hidden_matrices(model): + parameter.grad = torch.zeros_like(parameter) + + optimizer.step() + + expected_scales = torch.tensor([1.0, 1.0, 0.5, 0.25])[:, None, None] + q_after = model.blocks[0].attn.q_proj.weight.view(4, 4, 16) + k_after = model.blocks[0].attn.k_proj.weight.view(4, 4, 16) + assert torch.allclose(q_after, q_before * expected_scales) + assert torch.allclose(k_after, k_before * expected_scales) + assert optimizer.last_diagnostics["active_fraction"] == pytest.approx(0.5) + assert optimizer.last_diagnostics["min_gamma"] == pytest.approx(0.0625) + + def test_qk_clip_is_noop_below_threshold() -> None: model = small_model() optimizer = MuonClip( From 245586b6eccdc809b05aa850548038330541ff71 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 26 Aug 2026 22:06:22 -0700 Subject: [PATCH 5/5] =?UTF-8?q?Add=20five-seed=20four-head=20AdamW?= =?UTF-8?q?=E2=80=93MuonClip=20baseline=20config?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../configs/four_head_4epoch_baseline.yaml | 145 ++++++++++++++++++ 1 file changed, 145 insertions(+) create mode 100644 baseline/nanogpt_one_head/configs/four_head_4epoch_baseline.yaml diff --git a/baseline/nanogpt_one_head/configs/four_head_4epoch_baseline.yaml b/baseline/nanogpt_one_head/configs/four_head_4epoch_baseline.yaml new file mode 100644 index 00000000..197d7c75 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/four_head_4epoch_baseline.yaml @@ -0,0 +1,145 @@ +protocol: + name: nanogpt_four_head_2026_08_27_ww_baseline + version: 1 + description: Four-corpus-equivalent-epoch AdamW and MuonClip baselines with five paired seeds, four attention heads, and WeightWatcher raw and clip_xmax monitoring. + +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: 4 + 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