Add MLA support for Muon optimizer
Open
community-request
enhancement
waiting-on-customer
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Is your feature request related to a problem? Please describe.**
Currently, the Muon optimizer implementation in `megatron/core/optimizer/muon.py` lacks support for MLA (Multi-Latent Attention) parameters, as indicated by the TODO comment (`TODO(deyuf): support MLA`).
MLA-related parameters are not tagged with `is_qkv` and thus undergo direct orthogonalization. Maybe it lead to some unpredictable result.
**Describe the solution you'd like**
Add MLA support to the Muon optimizer.
Contributor guide
Assessment
This issue has not been assessed yet.