deepspeedai / deepspeedai/DeepSpeed
[REQUEST] Muon Optimizer - Different LR for Different Groups
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 43.1k
- Forks
- 5k
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 112
Description
Is your feature request related to a problem? Please describe.
When using Muon optimizer, now the method is to setup that in config. Params will be grouped into muon & adam automatically per #7555. However, there's a single LR that passed into both groups - this is not the standard pattern. We should have different LR for the muon group and the adam group (usually, muon should have larger LR). I wonder if there's a way to enable this behavior easier from my side. If not, then we probably need to add a support for different LR.
Describe the solution you'd like
Easiest way to achieve this is probably to have "muon_lr" and "adam_lr" instead of a single "lr" key. Then in different groups, fetch accordingly.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
No file or test is named. Start by tracing the Muon optimizer configuration and the automatic Muon/Adam parameter grouping described here and in #7555. Determine where the shared lr is applied, then verify that separate muon_lr and adam_lr values are accepted and used by their respective groups.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100