DeepSpeed
[muon] Per-head Muon for linear-attention layers: read head geometry from the owning module
#8436
Merged
Overview
Commits
25
Changes
View On GitHub

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

delock merged 25 commits into deepspeedai:master from alanhuangyoo:feat/per-head-muon-linear-attention
alanhuangyoo
alanhuangyoo [muon] Add per-head Newton-Schulz to muon_update
15e3bf4f
alanhuangyoo [muon] Tag attention projections with their head count and wire it th…
4644933f
alanhuangyoo [muon] Only tag projections whose output dim is head-blocked
11876c64
alanhuangyoo [muon] Cover the tagging against real model architectures
e9584d3f
alanhuangyoo [muon] Support MLA up-projections and read head counts from AutoTPMeta
870ee004
alanhuangyoo [muon] Cover MLA models that have no q_lora_rank
c6a7793b
alanhuangyoo [muon] Keep muon_update compiled
c3e93341
alanhuangyoo [muon] Resolve head geometry by candidate, and report what the flag did
37740f6f
alanhuangyoo alanhuangyoo requested a review from tjruwase tjruwase 46 days ago
alanhuangyoo alanhuangyoo requested a review from tohtana tohtana 46 days ago
alanhuangyoo alanhuangyoo requested a review from loadams loadams 46 days ago
alanhuangyoo [muon] Re-resolve the head count against the shard under tensor paral…
943fd03b
alanhuangyoo alanhuangyoo force pushed from 622c28de to 823e1e53 46 days ago
alanhuangyoo [muon] Read the layer's shape, not the ZeRO-3 partition's
5d164954
alanhuangyoo [muon] Read head geometry from the attention module when the config h…
0a1fbe1d
alanhuangyoo alanhuangyoo force pushed from 823e1e53 to 0a1fbe1d 46 days ago
alanhuangyoo
alanhuangyoo Merge remote-tracking branch 'upstream/master' into feat/per-head-muon
61b42532
delock delock requested a review from delock delock 45 days ago
alanhuangyoo Consolidate the per-head Muon tests into two files
789a1e7c
alanhuangyoo Merge feat/per-head-muon: the tests it added moved into the consolida…
6681fb1d
alanhuangyoo Merge master into feat/per-head-muon
1f143619
alanhuangyoo Merge feat/per-head-muon into feat/per-head-muon-linear-attention
fbf60f66
alanhuangyoo
alanhuangyoo Drop the _layer_shape copy the master merge duplicated
0cd2bc68
alanhuangyoo Merge feat/per-head-muon into feat/per-head-muon-linear-attention
21d11e50
pengdurice Merge branch 'master' into feat/per-head-muon
ff07664e
alanhuangyoo Allow for BLAS kernel differences in the per-head NS agreement test
effad68a
alanhuangyoo Merge master into feat/per-head-muon
f090ecf9
alanhuangyoo Merge feat/per-head-muon into feat/per-head-muon-linear-attention
69f04350
alanhuangyoo Merge master into feat/per-head-muon-linear-attention
6269b13e
alanhuangyoo
delock
delock
delock commented on 2026-09-18
delock
alanhuangyoo Read linear-attention geometry only from listed modules, and let it win
64bfdcbb
alanhuangyoo
delock
delock approved these changes on 2026-09-20
delock delock enabled auto-merge 33 days ago
delock Merge branch 'master' into feat/per-head-muon-linear-attention
6f3600ca
delock delock merged 141ee9f8 into master 33 days ago

Login to write a write a comment.

Login via GitHub

Reviewers
Assignees
No one assigned
Labels
None yet
Milestone
No milestone
Terms Of Use Privacy Policy Refund Policy Impressum