Muon silently discards the param groups it is given

[alanhuangyoo]

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:

# 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.