feat(optim): heads_per_block for MLA up-projections (GLM-5 Muon Split)

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -440,7 +440,7 @@ class MuonConfig(BaseOptimizerConfig):
     """β2 for the AdamW/Lion sub-optimizer used on non-Muon params."""
 
     heads_per_block: int | None = Field(None, ge=1)
-    """Orthogonalize the attention query, key and value projections in blocks of this many heads, each block on its own (per-head Muon, Kimi K3 report, section 2.5; ``1`` = per head). GQA key/value projections are split by their own heads; ``o_proj`` stays one matrix. If None, each projection is one matrix."""
+    """Orthogonalize the attention query, key and value projections in blocks of this many heads, each block on its own (per-head Muon, Kimi K3 report, section 2.5; ``1`` = per head). GQA key/value projections are split by their own heads; ``o_proj`` stays one matrix. For MLA, the query and key/value up-projections are split by head (GLM-5 "Muon Split"). If None, each projection is one matrix."""
 
 
 class SignSGDConfig(BaseOptimizerConfig):

src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py

@@ -15,6 +15,7 @@
 from prime_rl.trainer.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig, _index_cache_skip_topk
 from prime_rl.trainer.models.glm_moe_dsa.converting_glm_moe_dsa import conversion_chain
 from prime_rl.trainer.models.glm_moe_dsa.sparse_mla_attention import GlmMoeDsaAttention, SparseMlaAttentionArgs
+from prime_rl.trainer.models.layers.attn import attention_head_partitions
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput
 from prime_rl.trainer.models.layers.mlp import FeedForward
 from prime_rl.trainer.models.layers.moe import MoE, MoEArgs
@@ -345,5 +346,8 @@ def init_buffers_post_meta(self):
             )
             rotary_emb.inv_freq.copy_(inv_freq)
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        return attention_head_partitions((layer.self_attn for layer in self.model.layers), heads_per_block)
+
 
 __all__ = ["GlmMoeDsaConfig", "GlmMoeDsaPreTrainedModel", "GlmMoeDsaModel", "GlmMoeDsaForCausalLM"]

src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py

@@ -6,6 +6,7 @@
 
 from prime_rl.configs.trainer import DSABackend
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
+from prime_rl.trainer.models.layers.attn import head_blocks
 from prime_rl.trainer.models.layers.norms import LayerNorm, RMSNorm, RMSNormConfig
 from prime_rl.trainer.models.layers.rotary_emb import rotate_half
 from prime_rl.utils.cp import CPContext, gather_for_cp
@@ -149,6 +150,14 @@ def __init__(self, args: SparseMlaAttentionArgs):
 
         self.cp_context = CPContext()
 
+    def muon_head_partitions(self, heads_per_block: int) -> dict[nn.Parameter, tuple[int, ...]]:
+        # Per-head blocks of the query and key/value up-projections (GLM-5 "Muon Split")
+        kv_head_rows = self.qk_nope_head_dim + self.v_head_dim
+        return {
+            self.q_b_proj.weight: head_blocks(self.q_b_proj.out_features, self.qk_head_dim, heads_per_block),
+            self.kv_b_proj.weight: head_blocks(self.kv_b_proj.out_features, kv_head_rows, heads_per_block),
+        }
+
     def mla_latents(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         q_latent = self.q_a_layernorm(self.q_a_proj(hidden_states))
         compressed_kv = self.kv_a_proj_with_mqa(hidden_states)