[muon] Per-head Muon for linear-attention layers: read head geometry from the owning module by alanhuangyoo · Pull Request #8436 · deepspeedai/DeepSpeed

alanhuangyoo · 2026-09-06T11:30:46Z

Closes #8420. Stacked on #8384 — that branch is the base, so the diff to review is the single commit 622c28d; the rest is #8384.

@delock answered the two model questions in #8420: Muon Split applies to KDA, and not to GLM's indexer. This implements exactly that.

The problem

#8384 derives candidate geometries from the config: head counts through AutoTPMeta, per-head widths from head_dim or the MLA fields. Kimi-K3's linear-attention layers are built from linear_attn_config instead:

# modeling_kimi_k3_linear.py, KimiDeltaAttention
self.head_dim  = config.linear_attn_config["head_dim"]    # 32
self.num_heads = config.linear_attn_config["num_heads"]   # 8
self.q_proj = nn.Linear(self.hidden_size, self.head_k_dim * self.num_k_heads, bias=False)

q_proj is (256, 1024) = 8 x 32. Nothing at the top level says 32: head_dim is 74 and qk_nope + qk_rope is 96, so every candidate failed the width check. The head count was right and the width was wrong, so the projections were declined rather than recognized.

The change

Map the parameter to the module that owns it and ask that module for the geometry it was built with. It is one more candidate, not an override — #8384's shape confirmation still decides, and a module that disagrees with the config on the head count is still ambiguous and still skipped.

Only modules that say they are attention are asked. That is what keeps the indexer out: GlmMoeDsaIndexer.wq_b is (512, 512) and index_n_heads 8 * index_head_dim 64 is exactly 512, so reading n_heads/head_dim off any module carrying them would tag it. It is declined by the module, not by the shape, and there is a test asserting the geometry does confirm so that the reason stays honest.

The cost of that choice is stated in a test: Qwen3-Next's Qwen3NextGatedDeltaNet does not say attention, so it keeps the full-matrix path. That is #8384's behaviour rather than a regression, and widening it is one marker. The alternative — ask every module, name the ones to skip — tags the indexer by default, which is the opposite of what #8420 concluded.

Both models, instantiated

inference-optimization/Kimi-K3-0.40B, 318 Muon parameters:

owner	parameter	shape	before	after
`KimiMLAAttention` (2 layers)	`q_b_proj`	(768, 256)	8 heads (`mla-q`)	unchanged
`KimiMLAAttention` (2 layers)	`kv_b_proj`	(1024, 128)	8 heads (`mla-kv`)	unchanged
`KimiDeltaAttention` (6 layers)	`q_proj`	(256, 1024)	full matrix	8 heads (`owner-q`)
`KimiDeltaAttention` (6 layers)	`k_proj`	(256, 1024)	full matrix	8 heads (`owner-k`)
`KimiDeltaAttention` (6 layers)	`v_proj`	(256, 1024)	full matrix	8 heads (`owner-v`)
`KimiDeltaAttention`	`g_proj`	(256, 1024)	full matrix	unchanged
`KimiDeltaAttention`	`f_b_proj`, `b_proj`, `o_proj`	—	full matrix	unchanged

4 tagged -> 22. g_proj is the case worth noting: same 256 x 1024 shape as q_proj, in the same module, and it stays whole because its leaf name is not a projection name. There is a test for it.

inference-optimization/GLM-5.2-0.8B-A0.8B, 69 Muon parameters — unchanged at 12:

TAGGED:   q_b_proj (4096, 512) x6 -> 16 (mla-q)
          kv_b_proj (5120, 128) x6 -> 16 (mla-kv)
DECLINED: wq_b (512, 512) x3   width-mismatch
          wk (64, 2048) x3     width-mismatch

End to end, real checkpoints

bf16, ZeRO-1, per_head_muon: true, 6 steps, loaded weights rather than from_config (the KDA layer has an uninitialized dt_bias, so a randomly initialized hybrid gives NaN on step 0 regardless of the optimizer):

Kimi-K3-0.40B   tagged=22 (18 KDA + 4 MLA)   loss 20.00 -> 18.25   all params finite
GLM-5.2-0.8B    tagged=12  indexer_tagged=0  loss 12.29 ->  9.98   all params finite

Tests

12 cases added to tests/unit/runtime/zero/test_per_head_muon_tagging.py, 54 passing. Reverting the code change fails exactly the four that assert the new behaviour:

FAILED test_linear_attention_geometry_comes_from_the_owning_module[q_proj] - assert None == 8
FAILED test_linear_attention_geometry_comes_from_the_owning_module[k_proj] - assert None == 8
FAILED test_linear_attention_geometry_comes_from_the_owning_module[v_proj] - assert None == 8
FAILED test_the_value_projection_uses_the_value_width                      - assert None == 8

The other eight pass both before and after by design — they pin that nothing widened: the indexer, a module that does not say attention, geometry that does not match the shape, the non-projection matrices of a KDA layer, and a standard Llama where the config and the module agree.

The two existing tests that assert Kimi's KDA is declined on a config-only model are unchanged and still pass. A SimpleNamespace config has no modules, so they now document the narrower thing they were always testing: the config alone cannot describe this layout.

Whether per-head helps on linear-attention heads is a separate question from whether the split is well-defined on them, and #8384 has the measurements on that.

Full-matrix orthogonalization treats every attention head as one coupled block, so heads with larger momentum dominate the shared update direction while smaller-scale heads get insufficiently normalized updates. Kimi K3 (arXiv:2607.24653 5 2.5) and GLM-5 Muon Split (arXiv:2602.15763) both orthogonalize per head instead. With num_heads set, the update for a [num_heads * head_dim, in_features] projection is viewed as [num_heads, head_dim, in_features] and Newton-Schulz runs on that batch, with the existing max(1, m/n)**0.5 scaling applied per head block. Both NS kernels are already batch-safe, so this reuses the path the expert-group branch takes. Kernel only; the metadata plumbing and config surface for deepspeedai#8367 follow separately. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

…rough Adds the metadata and config half of deepspeedai#8367 on top of the kernel. set_optimizer_flags now also tags muon_num_heads next to use_muon, gated on an opt-in optimizer.params.per_head_muon. Head structure comes from the model config rather than AutoTP, so it does not require AutoTP to be enabled: q/o projections are blocked by num_attention_heads, k/v by num_key_value_heads, which differ under GQA. Deliberately conservative about what it claims to recognize. A fused QKV matrix is left on the full-matrix path - its three sections split separately, and under GQA they do not even share a head count - and any projection whose output dim does not divide by the head count is skipped with a warning rather than reshaped on a guess. All six muon_update call sites pass the tag through. Each one already operates on a whole parameter rather than a flat shard: the ZeRO-1/2 path views the momentum back to tensor.size() and asserts ndim > 1, and the ZeRO-3 path takes param.grad directly. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

Two mistakes in the previous commit's tagging, both of which produced a wrong update rather than an error. o_proj / out_proj were tagged with the query head count, but their head structure is on the input dimension ([hidden, num_heads * head_dim]) while the split is on dim 0. With the usual hidden == num_heads * head_dim they still divide evenly, so the matrix was silently cut across the wrong axis. Q/K/V only now. 'dense' was matched anywhere in the parameter path, which also names MLP matrices - intermediate.dense, output.dense, dense_h_to_4h, dense_4h_to_h - so a matrix with no head structure at all was split by the head count. Matching is now on the leaf module name against explicit Q/K/V names, and the shape has to confirm the layout: dim 0 divisible by the head count, and equal to num_heads * head_dim wherever the config states head_dim. Regression tests cover both; against the previous logic all five of the names they pin come back tagged. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

The synthetic module the other cases use has the leaf names I chose, which is circular for a change whose whole job is recognizing real ones. These build actual HF configs instead. llama / qwen2 / mistral (split QKV, GQA): q_proj tagged with the query head count, k_proj and v_proj with the kv count, o_proj and the MLP projections left alone. gpt_neox / falcon (fused QKV): nothing tagged. These name their output projection 'dense' and their MLP matrices 'dense_h_to_4h' / 'dense_4h_to_h', which is exactly what the previous substring matching got wrong - all three came back tagged. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

Addresses the review on deepspeedai#8384. MLA blocks its two up-projections by head instead of using q/k/v: q_b_proj is [num_heads * (qk_nope + qk_rope), rank] and kv_b_proj is [num_heads * (qk_nope + v_head_dim), rank]. That is the split GLM-5's Muon Split applies, and neither width is head_dim, so the previous shape check rejected both. The down-projections (q_a_proj, kv_a_proj_with_mqa) mix latent and rope components and stay on the full-matrix path. Verified against the real tensor shapes of the two models named in deepspeedai#8367: inference-optimization/GLM-5.2-0.8B-A0.8B tags q_b_proj (4096, 512) and kv_b_proj (5120, 128) with 16 heads, matching 16*(192+64) and 16*(192+128), and leaves the down-projections, o_proj and the MLP alone. inference-optimization/Kimi-K3-0.40B is kimi_linear rather than MLA - q_proj is [256, 1024] against 8 heads of 74 - so the shape check keeps it off the per-head path. Head counts now come from AutoTPMeta.from_model_config, the repo's stated single source of truth, which descends into text_config and probes the several spellings models use (num_heads, n_head, attention_heads) instead of one hardcoded name. Adds end-to-end training over ZeRO 1/2/3 at world size 2: tags survive deepspeed.initialize into the ZeRO call sites, per_head_muon stays off by default, and both paths train. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

An MLA attention block only builds the q_a/q_b pair when q_lora_rank is set. Without it the query up-projection is a plain q_proj, and its per-head width is still qk_nope_head_dim + qk_rope_head_dim rather than head_dim: self.q_proj = nn.Linear(hidden_size, num_heads * self.qk_head_dim) if self.q_lora_rank is None else None DeepSeek-V2-Lite is in that shape. Measured by instantiating DeepseekV2Attention on the released deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct config under transformers 5.16.1: heads=16 head_dim=64 qk_nope=128 qk_rope=64 q_lora_rank=None q_proj.weight (3072, 2048) 16 * 192 kv_b_proj.weight (4096, 512) 16 * 256 kv_a_proj_with_mqa.weight (576, 2048) o_proj.weight (2048, 2048) head_dim is present and equal to 64, so the shape check compared 3072 against 16 * 64 = 1024 and put the model back on the full-matrix path. The fallback is the behaviour before this feature, so nothing was wrong, but a whole class of MLA checkpoints never reached the per-head split the feature exists for. q_proj now takes the MLA width when the config carries the MLA head dimensions, and falls through to head_dim when it does not, so ordinary attention models are unchanged. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

Adding the per-head helper above muon_update left @compiler.compile() attached to the new function instead of to muon_update, so muon_update lost the torch.compile it has on master. That was not intended and it applies to every Muon user, not only the per-head path. The decorated set now matches master again: zeropower_via_newtonschulz5, zeropower_via_gram_newtonschulz and muon_update. _per_head_orthogonalize is called from inside muon_update and does not need its own. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

@delock

Addresses @delock's review. The dispatch was an ordered cascade, so which branch claimed a leaf depended on which config fields happened to exist rather than on the model. q_proj sat in both the MLA and the standard tables; Kimi-K3's KDA q_proj was claimed by the MLA branch on the strength of top-level qk_nope/qk_rope that belong to its two MLA layers, and survived only because 8 x 96 != 256. Safety by width mismatch, not by recognition. Now every geometry the config makes plausible for a leaf is collected and the shape confirms one: classify(leaf) -> kind geometry_candidates(kind, meta, cfg) -> [(heads, width, source)] confirm(param, candidates) -> (num_heads | None, reason) Candidates that confirm and agree on the head count are not a conflict, since the head count is the whole output. Candidates that confirm and disagree are, and the parameter is skipped with both named in the warning. This also removes a weaker rule that was hiding in the old code: with no head_dim the tagger fell back to divisibility alone, and rows % heads == 0 holds for matrices with no head structure at all -- the way o_proj slipped through in the first version. Every candidate now carries a width, derived as hidden_size // num_attention_heads where a config omits head_dim. An explicit opt-in that silently does nothing is its own failure. deepspeed.initialize now raises when per_head_muon is set and no projection could be tagged, listing the likely causes. The systemic one is tensor parallelism: the config describes the whole model while each rank holds a shard, so every width check fails and per-head is off model-wide while the user believes it is on. Non-systemic misses are aggregated per leaf as a warning so a hybrid model keeps per-head on its recognized layers, and an info line reports what was tagged. Also from the review: - AutoTPMeta.from_model_config is built once per model rather than per parameter. The model-taking _attention_head_count is kept for single-parameter callers. - ns_compute_dtype moves into original_muon.py, so the kernels and the test tolerances read the NS precision from one place instead of restating it. - per_head_muon is documented in config-json.md: what is tagged, what deliberately is not, the shape-confirms-the-name contract, and the reporting behaviour. - The three test files carried a Microsoft copyright line; they use the repo's two-line header now. Tests added: a Kimi-K3 hybrid fixture pinning that the KDA projections are declined by candidate confirmation rather than by branch order, a synthetic ambiguity case, agreement-is-not-ambiguity, the divisibility-alone rejection, head_dim derivation, and the initialize-time error. 52 unit, 9 end-to-end on 2 x H20. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

…lelism set_optimizer_flags runs before the engine partitions the model, so the count it records is the model's, not the rank's. AutoTP replaces the parameter's .data in place, so the tag rides onto a column-parallel shard and nothing catches it: with tp=2 a tag of 8 heads lands on a shard holding 4 heads' worth of rows, out_features % num_heads still divides, and Newton-Schulz runs on half of each head. Measured on a Llama with 8 heads of 32 and autotp_size 2: every q/k/v tagged 8 heads of 16. The per-head width is what a column-parallel split leaves alone, so record it and re-derive the count from the shard. A shard whose rows are not a multiple of the width does not hold whole heads and is dropped. This is not only a repair. A column-parallel shard holds whole heads, so per-head Newton-Schulz on the shard is bit-identical to the corresponding blocks of per-head Newton-Schulz on the whole matrix - the split is along the same axis the batch is taken over. There is a test asserting equality, not closeness, for both kernels at tp=2 and tp=4, against the whole-matrix path, which differs by more than 10% on the same input. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

zero.Init replaces a partitioned parameter's data with a flat placeholder and records the layer's shape as ds_shape, so param.shape is 1-D for every parameter in the model. The width check then confirms nothing and the flag raises on a model whose layout it could read perfectly well. Same root cause as deepspeedai#8438, which fixes the use_muon test on master; this applies it to the tagger's shape reads as well. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

…as none Kimi-K3's linear-attention layers keep their head count and width in linear_attn_config and build q/k/v at num_k_heads * head_k_dim, which no top-level config field describes, so the tagger declined them on a width check rather than recognising them. Ask the module that built the projection; it holds both numbers as attributes. Only modules that say they are attention are asked. GLM-5.2's DSA indexer carries an n_heads/head_dim pair of its own and its wq_b is exactly index_n_heads * index_head_dim, so an unconditional read would tag it. Muon Split covers attention, not the indexer that picks the keys attention will see (deepspeedai#8420). The module is another candidate, not an override: the shape still has to confirm it, and a candidate that disagrees with the config on the head count is still ambiguous and still skipped. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

alanhuangyoo · 2026-09-06T12:14:55Z

Rebased onto the updated #8384 branch (0a1fbe1). Two fixes landed in the base since this was opened, both from the same review pass and both affecting what this PR sits on:

Tensor parallelism. set_optimizer_flags runs before AutoTP partitions, and AutoTP replaces the parameter's .data in place, so the head count made against the whole model rode onto a shard and Newton-Schulz ran on half-heads. The count is now re-derived from the shard, which also makes per-head exact under TP — torch.equal, not close.
zero.Init. Under ZeRO-3's zero.Init every parameter reports as 1-D, so nothing confirmed and the flag raised. Shapes are read from ds_shape now.

Neither changes what this PR does: Kimi-K3 still goes 4 tagged -> 22, GLM-5.2 stays at 12 with the indexer untouched, and the whole per-head suite is 70 passed.

The commit to review here is still the single one on top.

@pengdurice

@pengdurice asked for fewer files. Five, one per concern, become two split by what they need to run: tests/unit/runtime/zero/test_per_head_muon.py CPU the arithmetic (was test_per_head_muon.py) which parameters are tagged, and with how many heads (was test_per_head_muon_tagging.py) re-resolving the count against a tensor-parallel shard (was test_per_head_muon_tensor_parallel.py) tests/unit/ops/muon/test_per_head_muon_accelerator.py GPU a shard of the per-head result is the per-head result of the shard (was test_per_head_muon_under_sharding.py) end-to-end training (was test_per_head_muon_e2e.py) Nothing is dropped or rewritten: 77 tests before, 77 after, same names. The tagging module's _Attn and the tensor-parallel module's _Attn were different classes with the same name, so the latter is now _ShardedAttn. The accelerator file is not called test_per_head_muon.py because these directories have no __init__.py, and two same-named modules break collection when both are selected in one run. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

@pengdurice

…ted files deepspeedai#8384's five test files became two (@pengdurice asked for fewer). The linear attention and DSA indexer cases this branch added to test_per_head_muon_tagging.py move into section 3 of the consolidated tests/unit/runtime/zero/test_per_head_muon.py; nothing is rewritten. 89 passed across the two files on 2xH20 (77 from the base, 12 from here).

alanhuangyoo added 8 commits September 1, 2026 19:33

alanhuangyoo requested review from loadams, tjruwase and tohtana as code owners September 6, 2026 11:30

alanhuangyoo mentioned this pull request Sep 6, 2026

[Per-Head Muon] Head geometry that lives outside num_attention_heads/head_dim: Kimi-K3 KDA layers and GLM-5.2's DSA indexer #8420

Open

alanhuangyoo force-pushed the feat/per-head-muon-linear-attention branch from 622c28d to 823e1e5 Compare September 6, 2026 11:58

alanhuangyoo added 2 commits September 6, 2026 20:13

alanhuangyoo force-pushed the feat/per-head-muon-linear-attention branch from 823e1e5 to 0a1fbe1 Compare September 6, 2026 12:14

alanhuangyoo mentioned this pull request Sep 7, 2026

[muon] Per-head Newton-Schulz for attention projections #8384

Open

Merge remote-tracking branch 'upstream/master' into feat/per-head-muon

61b4253

delock self-requested a review September 7, 2026 13:56

alanhuangyoo added 2 commits September 9, 2026 03:21

Navigation Menu

Uh oh!

[muon] Per-head Muon for linear-attention layers: read head geometry from the owning module - #8436

[muon] Per-head Muon for linear-attention layers: read head geometry from the owning module#8436
alanhuangyoo wants to merge 14 commits into
deepspeedai:masterfrom
alanhuangyoo:feat/per-head-muon-linear-attention

alanhuangyoo commented Sep 6, 2026

Uh oh!

alanhuangyoo commented Sep 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

alanhuangyoo commented Sep 6, 2026

The problem

The change

Both models, instantiated

End to end, real checkpoints

Tests

Uh oh!

alanhuangyoo commented Sep 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant