[faresobeid]

Stacked on #4114.

Extends optim.heads_per_block to MLA attention (GLM-5.x, glm_moe_dsa). The query up-projection q_b_proj [H·qk_head_dim, q_lora] and the key/value up-projection kv_b_proj [H·(qk_nope+v), kv_lora] are orthogonalized in blocks of heads_per_block heads. This is GLM-5's "Muon Split" for W_UQ/W_UK/W_UV (GLM-5 report, section 2). GlmMoeDsaForCausalLM.muon_head_partitions returns these blocks; with #4114 alone, glm_moe_dsa raises. q_a_proj, kv_a_proj_with_mqa and o_proj stay whole matrices. Default None is unchanged.

Measurements

GLM-5, 4 layers, random init, 1 node of 8xH200, offload = "none", optimizer step only, steady state, max over ranks:

heads_per_block	step	NS calls per step
None	0.197 s	(16384,2048) x4, (28672,512) x4
4	0.208 s	per-block
1	0.240 s	(256,2048) x256, (448,512) x256

• dion runs one Newton-Schulz call per partition in a Python loop (the TODO in dion muon_orthogonalize_update). That costs about +11 ms per layer at heads_per_block = 1. Batching equal partitions in dion would remove this.
• The update norm is unchanged, as rms_norm predicts: the squared parameter delta is 0.0153 after 4 steps in all three runs.

Checks

• ruff check, ruff format --check. tests/unit/test_configs.py -m "not gpu".

---

Note

Medium Risk
Changes optimizer parameter partitioning for a specific model architecture; incorrect block boundaries would affect training dynamics but scope is limited to Muon + heads_per_block on glm_moe_dsa.

Overview
Extends Muon heads_per_block to GLM MoE DSA sparse MLA (glm_moe_dsa), implementing GLM-5’s “Muon Split” on the query and key/value up-projections (q_b_proj, kv_b_proj) while leaving down-projections and o_proj as single matrices.

GlmMoeDsaAttention.muon_head_partitions defines per-head row blocks via head_blocks; GlmMoeDsaForCausalLM aggregates them with attention_head_partitions so the trainer can orthogonalize Muon updates on those weights (previously this path would fail for this model). The heads_per_block config docstring now documents MLA behavior; default None is unchanged.

^{Reviewed by Cursor Bugbot for commit 41f92bf. Bugbot is set up for automated code reviews on this repo. Configure here.}