Muon silently discards the param groups it is given (#8440)
Muon's param groups are built by DeepSpeed rather than taken from the
caller, because which half of Muon a parameter belongs to is a property
of the parameter. Doing that flattened away whatever the incoming groups
set.
## What happens
`_configure_basic_optimizer` collapsed `model_parameters` into one list
and rebuilt two groups from the config:
```python
# Flatten param group dicts (created by MoE/EP) into a raw parameter list
all_params = []
for item in model_parameters:
if isinstance(item, dict):
all_params.extend(item['params'])
...
```
The comment says MoE/EP, and for those groups it is harmless — the MoE
identity rides on `param.group_name` and
`split_params_into_different_moe_groups_for_optimizer` rebuilds them
afterwards. But the same line eats a caller's own groups, and every
other optimizer here receives `model_parameters` unchanged, so for them
a group's `lr` and `weight_decay` reach the optimizer.
The pattern this breaks is the one in most training recipes: no weight
decay on biases and norms. Passing the usual two groups, `wd 0.1` and
`0.0`, `lr 1e-3` and `1e-4`, against a config of `lr 5e-4, weight_decay
0.01`:
```
AdamW group 1: lr=1.0e-03 wd=0.1 group 2: lr=1.0e-04 wd=0.0
Muon muon-params: lr=5.0e-04 wd=0.01 adam-params: lr=5.0e-04 wd=0.01
```
Everything falls back to the config values. The parameters the caller
excluded from weight decay are decayed at 0.01 anyway, and nothing says
so.
## The change
Split each incoming group into its Muon and Adam halves and carry that
group's settings onto both. Settings resolve most-specific-last: the
config's shared value, then `muon_lr` / `adam_lr`, then whatever the
group itself sets.
After:
```
Muon group0-muon-params: lr=1.0e-03 wd=0.1 group1-adam-params: lr=1.0e-04 wd=0.0
```
Also raise on a parameter with no `use_muon` attribute. Today that logs
an error and then dies two lines later on `p.use_muon` with an
`AttributeError`, so the helpful message is followed by an unhelpful
traceback.
## Backward compatibility
A plain parameter list — the common case, and every existing test — is
unchanged, group names included. Verified byte-for-byte against master:
```
master at init muon-params: lr=2.000e-02 adam-params: lr=1.000e-05
this PR at init muon-params: lr=2.000e-02 adam-params: lr=1.000e-05
```
Names only gain a prefix when there is more than one incoming group,
which is required anyway: MoE regrouping keys its buckets by `name`, so
two groups must not collide on one. Nothing in the repo reads
`'muon-params'` / `'adam-params'`.
`muon_lr` / `adam_lr` still override the shared `lr`, and now lose to a
group that sets its own `lr`, which is the same precedence every other
optimizer has.
## Tests
`tests/unit/runtime/zero/test_muon_param_groups.py`, 11 cases. Ten call
the grouping directly (it is a staticmethod, so they need no engine and
no GPU): the historical names and values for a plain list,
`muon_lr`/`adam_lr` overriding, a group keeping its own weight decay and
its own lr, a group without an lr still taking `muon_lr`, a mixed group
splitting in two with both halves keeping its settings, distinct names
across groups, each half getting only the keys its optimizer accepts,
frozen parameters left out, and the untagged-parameter error.
The eleventh is the regression test: a real `deepspeed.initialize` with
the two-group no-decay pattern. On master it fails with
```
AssertionError: the groups the user passed were not honoured: [0.01, 0.01]
```
and passes here with `[0.0, 0.1]`. The other ten exercise a helper that
does not exist on master, so they fail there for the trivial reason; the
end-to-end one is the one that fails for the right reason.
yapf and flake8 clean.
Related: #7657 asked for different learning rates for the Muon and Adam
halves, which `muon_lr`/`adam_lr` covers at the config level; this makes
the general per-group form work too. #7713 is a different way the
configured Muon lr fails to reach the optimizer (an LR scheduler
broadcasting one scalar over both groups) and is not addressed here.
Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
