[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:masterdeepspeedai/DeepSpeed:masterfrom

alanhuangyoo:feat/per-head-muon-linear-attentionalanhuangyoo/DeepSpeed:feat/per-head-muon-linear-attentionCopy head branch name to clipboard