Add matched four-head nanoGPT baseline

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

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:

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:

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(

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",