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}"