From 9cda44ec0e5475e231c0508a10820735ee7e776a Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 6 Sep 2026 21:29:18 +0800 Subject: [PATCH] Keep the param groups Muon was given instead of flattening them away Which half of Muon a parameter belongs to is a property of the parameter, so _configure_basic_optimizer has to build its own groups. It did that by flattening model_parameters into one list and rebuilding two groups from the config, which discards whatever the incoming groups set. Every other optimizer here receives model_parameters unchanged, so for them a group's own lr and weight_decay reach the optimizer. The consequence is the no-weight-decay-on-biases-and-norms grouping that most training recipes use. Passing the usual two groups, wd 0.1 and 0.0: AdamW -> lr=1.0e-03 wd=0.1 lr=1.0e-04 wd=0.0 Muon -> lr=5.0e-04 wd=0.01 lr=5.0e-04 wd=0.01 Everything falls back to the config values and nothing is reported, so parameters the user excluded from weight decay are decayed anyway. 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 the group's own. A plain parameter list is unchanged, names included. Also raise on a parameter with no use_muon attribute rather than logging an error and then failing on p.use_muon two lines later. Signed-off-by: alanhuangyoo --- deepspeed/runtime/engine.py | 99 ++++++---- .../runtime/zero/test_muon_param_groups.py | 170 ++++++++++++++++++ 2 files changed, 236 insertions(+), 33 deletions(-) create mode 100644 tests/unit/runtime/zero/test_muon_param_groups.py diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 86918bd71c5a..e7fea5eb279e 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -2187,6 +2187,71 @@ def get_optimizer_configuration(self, optimizer_parameters, adam_w_mode, allow_l return None, {} return FusedAdam, {'adam_w_mode': adam_w_mode} + # Which of the optimizer's config keys each half of a Muon param group accepts. Muon takes a + # momentum and a Newton-Schulz method; the auxiliary Adam takes betas and eps. + _MUON_HALF_KEYS = ("lr", "momentum", "weight_decay", "ns_method") + _ADAM_HALF_KEYS = ("lr", "betas", "eps", "weight_decay") + + @staticmethod + def _muon_half_defaults(optimizer_parameters, keys, lr_override): + """Config-level settings for one half, with muon_lr / adam_lr overriding the shared lr.""" + defaults = {key: optimizer_parameters[key] for key in keys if key in optimizer_parameters} + if lr_override in optimizer_parameters: + defaults["lr"] = optimizer_parameters[lr_override] + return defaults + + @staticmethod + def _muon_param_groups(model_parameters, optimizer_parameters): + """Split each incoming param group into its Muon and Adam halves. + + Muon has to build its own groups, because which half a parameter belongs to is a + property of the parameter rather than of the config. The incoming groups still have to + survive that: every other optimizer here receives `model_parameters` unchanged, so a + group's own `lr` or `weight_decay` reaches it. Flattening the groups into one list threw + those away, and the no-weight-decay-on-biases-and-norms grouping that most training + recipes use was silently ignored - the parameters the user excluded were decayed at the + config's rate instead, with nothing reported. + + Settings are resolved most-specific-last: the config's shared value, then `muon_lr` / + `adam_lr`, then whatever the group itself sets. + """ + groups, loose = [], [] + for item in model_parameters: + (groups if isinstance(item, dict) else loose).append(item) + if loose or not groups: + groups.append({"params": loose}) + + missing = [p for group in groups for p in group["params"] if not hasattr(p, "use_muon")] + if missing: + raise ValueError(f"The Muon optimizer needs every parameter tagged with use_muon, and {len(missing)} " + "are not. deepspeed.initialize tags them from the model it is given, so this means " + "model_parameters holds parameters that model does not. Set `param.use_muon = " + "True / False` on them, or pass them as part of the model.") + + muon_keys, adam_keys = DeepSpeedEngine._MUON_HALF_KEYS, DeepSpeedEngine._ADAM_HALF_KEYS + halves = ( + (True, "muon", muon_keys, DeepSpeedEngine._muon_half_defaults(optimizer_parameters, muon_keys, "muon_lr")), + (False, "adam", adam_keys, DeepSpeedEngine._muon_half_defaults(optimizer_parameters, adam_keys, + "adam_lr")), + ) + + param_groups = [] + for index, group in enumerate(groups): + overrides = {key: value for key, value in group.items() if key != "params"} + trainable = [p for p in group["params"] if p.requires_grad] + for use_muon, label, keys, defaults in halves: + half = [p for p in trainable if bool(p.use_muon) is use_muon] + if not half: + continue + settings = dict(defaults) + settings.update({key: value for key, value in overrides.items() if key in keys}) + # One incoming group is the common case and keeps the historical names; more than + # one needs distinct ones, because MoE regrouping keys its buckets by name. + prefix = overrides.get("name") or (f"group{index}" if len(groups) > 1 else None) + name = f"{prefix}-{label}-params" if prefix else f"{label}-params" + param_groups.append(dict(params=half, use_muon=use_muon, name=name, **settings)) + return param_groups + def _configure_basic_optimizer(self, model_parameters): # Copy so the pop() calls below (torch_adam, adam_w_mode, fp32_optimizer_states) do not # mutate the shared config dict returned by optimizer_params(). @@ -2267,39 +2332,7 @@ def _configure_basic_optimizer(self, model_parameters): adam_optimizer, adam_optimizer_kwargs = self.get_optimizer_configuration(optimizer_parameters, adam_w_mode, allow_legacy_fallback=True) - # 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']) - else: - all_params.append(item) - if not all([hasattr(p, 'use_muon') for p in all_params]): - msg = "Muon optimizer is used, but the use_muon attribute is NOT configured for some of the model parameters, " \ - "please set by `param.use_muon = True / False` for all params" - logger.error(msg) - muon_params = [p for p in all_params if p.use_muon and p.requires_grad] - non_muon_params = [p for p in all_params if (not p.use_muon) and p.requires_grad] - param_groups = [] - if muon_params: - accepted_parameters = dict() - for key in ["lr", "momentum", "weight_decay", "muon_lr", "ns_method"]: - if key in optimizer_parameters: - if key == "muon_lr": # muon_lr will override lr - accepted_parameters['lr'] = optimizer_parameters[key] - else: - accepted_parameters[key] = optimizer_parameters[key] - param_groups.append(dict(params=muon_params, use_muon=True, name='muon-params', **accepted_parameters)) - if non_muon_params: - accepted_parameters = dict() - for key in ["lr", "betas", "eps", "weight_decay", "adam_lr"]: - if key in optimizer_parameters: - if key == "adam_lr": # adam_lr will override lr - accepted_parameters['lr'] = optimizer_parameters[key] - else: - accepted_parameters[key] = optimizer_parameters[key] - param_groups.append( - dict(params=non_muon_params, use_muon=False, name='adam-params', **accepted_parameters)) + param_groups = self._muon_param_groups(model_parameters, optimizer_parameters) if self.has_moe_layers: from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer param_groups = split_params_into_different_moe_groups_for_optimizer(param_groups) diff --git a/tests/unit/runtime/zero/test_muon_param_groups.py b/tests/unit/runtime/zero/test_muon_param_groups.py new file mode 100644 index 000000000000..93cda0a9732f --- /dev/null +++ b/tests/unit/runtime/zero/test_muon_param_groups.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Muon has to build its own param groups, but the ones it was given must survive that. + +Which half of Muon a parameter belongs to is a property of the parameter, so `_configure_basic_optimizer` +regroups them. Every other optimizer receives `model_parameters` unchanged, so a group's own +`lr` / `weight_decay` reaches it; Muon flattened the groups away, which silently dropped the +no-weight-decay-on-biases-and-norms grouping most recipes use. +""" + +import pytest +import torch + +import deepspeed +from deepspeed.runtime.engine import DeepSpeedEngine +from unit.common import DistributedTest + + +def _param(shape, use_muon, requires_grad=True): + p = torch.nn.Parameter(torch.zeros(*shape), requires_grad=requires_grad) + p.use_muon = use_muon + return p + + +def _groups(model_parameters, **optimizer_parameters): + """`_muon_param_groups` is a staticmethod, so the grouping runs without an engine.""" + return DeepSpeedEngine._muon_param_groups(model_parameters, optimizer_parameters) + + +def _by_name(param_groups): + return {g["name"]: g for g in param_groups} + + +def test_a_plain_parameter_list_keeps_the_historical_groups(): + groups = _by_name(_groups([_param((4, 4), True), _param((4, ), False)], lr=1e-3, weight_decay=0.01)) + + assert set(groups) == {"muon-params", "adam-params"} + assert groups["muon-params"]["lr"] == 1e-3 + assert groups["adam-params"]["weight_decay"] == 0.01 + + +def test_muon_lr_and_adam_lr_still_override_the_shared_lr(): + groups = _by_name(_groups([_param((4, 4), True), _param((4, ), False)], lr=1e-4, muon_lr=2e-2, adam_lr=1e-5)) + + assert groups["muon-params"]["lr"] == 2e-2 + assert groups["adam-params"]["lr"] == 1e-5 + + +def test_a_group_keeps_its_own_weight_decay(): + """The pattern this is really about: no weight decay on biases and norms.""" + decayed, undecayed = _param((4, 4), True), _param((4, ), False) + groups = _by_name( + _groups([{ + "params": [decayed], + "weight_decay": 0.1 + }, { + "params": [undecayed], + "weight_decay": 0.0 + }], + lr=1e-3, + weight_decay=0.01)) + + assert [g["weight_decay"] for g in groups.values()] == [0.1, 0.0] + + +def test_a_group_keeps_its_own_lr_even_against_muon_lr(): + """Most specific wins: the config's shared lr, then muon_lr / adam_lr, then the group.""" + tagged = _param((4, 4), True) + groups = _groups([{"params": [tagged], "lr": 7e-3}], lr=1e-4, muon_lr=2e-2) + + assert groups[0]["lr"] == 7e-3 + + +def test_a_group_without_an_lr_still_takes_muon_lr(): + groups = _groups([{"params": [_param((4, 4), True)], "weight_decay": 0.0}], lr=1e-4, muon_lr=2e-2) + + assert groups[0]["lr"] == 2e-2 + assert groups[0]["weight_decay"] == 0.0 + + +def test_a_group_holding_both_kinds_splits_in_two_and_both_keep_its_settings(): + groups = _by_name( + _groups([{ + "params": [_param((4, 4), True), _param((4, ), False)], + "weight_decay": 0.0, + "name": "no-decay" + }], + lr=1e-3, + weight_decay=0.1)) + + assert set(groups) == {"no-decay-muon-params", "no-decay-adam-params"} + assert all(g["weight_decay"] == 0.0 for g in groups.values()) + assert all(g["lr"] == 1e-3 for g in groups.values()) + + +def test_group_names_stay_distinct_across_groups(): + """MoE regrouping keys its buckets by name, so two groups must not collide on one.""" + param_groups = _groups([{"params": [_param((4, 4), True)]}, {"params": [_param((4, 4), True)]}], lr=1e-3) + + names = [g["name"] for g in param_groups] + assert len(set(names)) == len(names) == 2 + + +def test_the_muon_half_gets_muon_keys_and_the_adam_half_gets_adam_keys(): + groups = _by_name( + _groups([_param((4, 4), True), _param((4, ), False)], + lr=1e-3, + momentum=0.9, + ns_method="standard", + betas=[0.9, 0.95], + eps=1e-8)) + + assert groups["muon-params"]["momentum"] == 0.9 + assert groups["muon-params"]["ns_method"] == "standard" + assert "betas" not in groups["muon-params"] + assert groups["adam-params"]["betas"] == [0.9, 0.95] + assert "momentum" not in groups["adam-params"] + + +def test_frozen_parameters_are_left_out(): + param_groups = _groups([_param((4, 4), True, requires_grad=False), _param((4, ), False)], lr=1e-3) + + assert [g["name"] for g in param_groups] == ["adam-params"] + + +def test_an_untagged_parameter_is_reported_rather_than_crashing_on_the_attribute(): + """Without use_muon the old path logged an error and then died on `p.use_muon` two lines later.""" + stray = torch.nn.Parameter(torch.zeros(4, 4)) + + with pytest.raises(ValueError, match="use_muon"): + _groups([stray], lr=1e-3) + + +class TestMuonParamGroupsEndToEnd(DistributedTest): + world_size = 1 + + def test_the_settings_reach_the_optimizer(self): + model = torch.nn.Sequential(torch.nn.Linear(8, 8, bias=True), torch.nn.LayerNorm(8)) + decay = [p for n, p in model.named_parameters() if p.ndim >= 2] + no_decay = [p for n, p in model.named_parameters() if p.ndim < 2] + + _, optimizer, _, _ = deepspeed.initialize(model=model, + model_parameters=[{ + "params": decay, + "weight_decay": 0.1 + }, { + "params": no_decay, + "weight_decay": 0.0 + }], + config={ + "train_micro_batch_size_per_gpu": 1, + "gradient_accumulation_steps": 1, + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": 1 + }, + "optimizer": { + "type": "Muon", + "params": { + "lr": 5e-4, + "weight_decay": 0.01 + } + }, + }) + + decays = sorted(g["weight_decay"] for g in optimizer.param_groups) + assert decays == [0.0, 0.1], f"the groups the user passed were not honoured: {decays}"