Add per-head Muon QKV orthogonalization and fix TP>num_query_groups with Muon. Generalize to MLA and GQA - #6326

Add per-head Muon QKV orthogonalization and fix TP>num_query_groups with Muon. Generalize to MLA and GQA#6326

mkhona-nvidia wants to merge 8 commits into

NVIDIA:mainNVIDIA/Megatron-LM:mainfrom

mkhona-nvidia:agent/muon-split-qkv-per-headmkhona-nvidia/Megatron-LM:agent/muon-split-qkv-per-headCopy head branch name to clipboard