[snimu]

What

New option optim.heads_per_block for Muon (MuonConfig, shared by SFT, RL and pretraining). Default None = today's behavior.

• heads_per_block = 1: each attention head's query, key and value rows are orthogonalized on their own (per-head Muon). 2 uses blocks of 2 heads, and so on.
• GQA-aware: key/value projections are split by their own (KV) heads.
• o_proj and every other matrix stay as they are.
• Works with the fused qkv runtime fusion and with the unfused q/k/v projections. Packed weights keep the q/k/v boundaries.
• Implemented as dion row partitions (matrix_partitions, already used for fused weights). So FSDP/EP sharding works as before (the all-to-all rebuilds the whole matrix, then each block is orthogonalized), there is no extra communication, and the rms_norm LR adjustment uses each block's shape.
• Raises a clear error if heads_per_block does not divide a projection's heads.

Which models: each model decides through a hook, PreTrainedModelPrimeRL.muon_head_partitions(heads_per_block). It returns the head blocks of all the model's attention projections, or raises. The default raises and names the model class, so a model is never covered only in part.

• Covered: llama, qwen3, qwen3_moe, glm4_moe, minimax_m2, nemotron_h (its attention layers), Laguna without output gate, and gpt-oss. They share qkv_head_partitions in layers/attn.py (packed qkv_proj or separate q/k/v).
• Gated queries: Qwen3.5 and Qwen3.8 store each head's query and gate together in q_proj ([q_h; gate_h]), so a block is 2 * head_dim rows per head. AFMoE: q/k/v plus its per-head gate_proj.
• Raise: DeepSeek V4, GLM-5 (glm_moe_dsa), Laguna with its output gate (g_proj), the Qwen3.5 VLM (vision encoder). Linear-attention layers (GatedDeltaNet, Mamba) are not softmax attention and are left as they are.
• feat(optim): heads_per_block for MLA up-projections (GLM-5 Muon Split) #4119 (GLM-5 MLA, stacked on this PR) should become glm_moe_dsa's implementation of the hook.

Code: muon_matrix_partitions() in trainer/optim/__init__.py builds the partitions (fused weights + model.muon_head_partitions(...) when heads_per_block is set). With the default it returns exactly the fused-weight partitions computed before.

Why

Kimi K3 (https://arxiv.org/abs/2607.24653, section 2.5, "Per-Head Muon") orthogonalizes Q/K/V per head. They report more balanced heads and better stability at scale. GLM-5's "Muon Split" (https://arxiv.org/abs/2602.15763) does the same for MLA up-projections. Group Muon (https://arxiv.org/abs/2605.08933) finds head-wise faster early but overtaken late on GPT-2 small, with groups of heads best, hence heads_per_block > 1.

Evidence (pretraining, MoE d384 / L6 / 6 heads, 2 KV, E64 K6, Dolma 3, 537M tokens)

Pretrain study: https://github.com/snimu/pretrain-experiments/blob/main/experiments/init/2026-10-15-muon-per-head/report.md (private).

• Per head vs whole projections: a tie.
  • +0.0014 ± 0.0012 nats over 8 seed-paired runs (4 LRs x 2 seeds); -0.0024 at the best LR. The best LR is the same for both.
  • Per head is ~0.1 nats ahead at the end of warmup and caught up by ~25% of training.
• Pairs of heads: -0.003 vs per head on one seed; a second seed is running.
• Step time: +2% on 1 GPU (dion loops over the blocks; equal-size blocks could be batched in dion later).

Default stays off.

Verification

• With the default, muon_matrix_partitions() returns exactly the fused-weight partitions computed before, so Muon receives the same inputs.
• Hook (e4a6e0935), CPU check on meta-device models: for every covered model, with heads_per_block 1 and 2 and with the qkv fusion where supported, the hook returns every q/k/v (and gate) weight of every attention layer, the blocks sum to each weight's rows, and each block has heads_per_block heads. The models listed under "Raise" raise, and a non-dividing heads_per_block raises.
• SFT smoke after the hook (Qwen3-0.6B, fake data, 5 steps, Muon lr 2e-3, 1 H200, job 4588): main 0956a8858 vs e4a6e0935. With the default every logged metric is bitwise identical to main at all 5 steps (loss 8.464153, 14.507845, 6.840884, 12.007687, 7.399467). heads_per_block = 1 / 2 reproduce the losses of the previous implementation (c48d13dd1): step 1 8.464153, step 5 7.101643 / 7.144592.
• SFT smoke (Qwen3-0.6B, fake data, 5 steps, 1 H200): main 0956a8858 vs this branch at c48d13dd1 (Muon, lr 2e-3, 1 H200, job 4394). With the default, every logged metric is bitwise identical to main at all 5 steps (loss 8.464153, 14.507845, 6.840884, 12.007687, 7.399467). With heads_per_block = 1 / 2, step 1 is identical (8.464153, no update yet) and later steps differ (step 5 loss 7.1016 / 7.1446 vs 7.3995).