From ee385e15289ca783fcdfc60424fa34ec0d021ae1 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 6 Sep 2026 20:11:42 +0800 Subject: [PATCH] Decide use_muon on the layer's shape, not on the ZeRO-3 partition's deepspeed.zero.Init replaces a partitioned parameter's data with a flat placeholder - torch.Size([0]) on ranks that do not hold it - and records the shape it has as a layer in ds_shape. set_optimizer_flags tests p.ndim >= 2 because Muon is defined on matrices, so under zero.Init every parameter in the model looks 1-D, none are tagged, ZeRO-3 finds no sub-group using Muon, and training continues with every parameter on the AdamW branch. Nothing reports it. Measured on a 2-layer Llama, ZeRO-3, Muon: without zero.Init 14 parameters tagged and muon_update called 84 times over 6 steps; with zero.Init 0 and 0. After this change the two agree. zero.Init is how models that do not fit on one device are built, so this is the configuration ZeRO-3 exists for. Signed-off-by: alanhuangyoo --- deepspeed/__init__.py | 17 ++- .../test_muon_use_muon_under_zero_init.py | 128 ++++++++++++++++++ 2 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 tests/unit/runtime/zero/test_muon_use_muon_under_zero_init.py diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py index 26bfc7c77df6..27b52f57a21b 100755 --- a/deepspeed/__init__.py +++ b/deepspeed/__init__.py @@ -81,10 +81,25 @@ def _parse_version(version_str): dist = None +def _layer_shape(param: torch.Tensor): + """The parameter's shape as a layer, rather than as a ZeRO-3 partition. + + Under ``deepspeed.zero.Init`` a partitioned parameter's data is a flat placeholder - + ``torch.Size([0])`` on the ranks that do not hold it - and the shape it has as a layer is + recorded as ``ds_shape``. Reading ``param.shape`` there sees a 1-D tensor for every + parameter in the model. + """ + ds_shape = getattr(param, "ds_shape", None) + return tuple(param.shape) if ds_shape is None else tuple(ds_shape) + + def set_optimizer_flags(config_class: DeepSpeedConfig, model: torch.nn.Module) -> None: if config_class.optimizer_name == MUON_OPTIMIZER: for name, p in model.named_parameters(): - if p.ndim >= 2 and not any(keyword in name.lower() for keyword in ("embed", "lm_head")): + # Muon is defined on matrices, so the test is on the layer's shape. `zero.Init` + # makes every parameter report as 1-D, which would switch Muon off for the whole + # model without anything saying so. + if len(_layer_shape(p)) >= 2 and not any(keyword in name.lower() for keyword in ("embed", "lm_head")): setattr(p, "use_muon", True) else: setattr(p, "use_muon", False) diff --git a/tests/unit/runtime/zero/test_muon_use_muon_under_zero_init.py b/tests/unit/runtime/zero/test_muon_use_muon_under_zero_init.py new file mode 100644 index 000000000000..6127e6455b76 --- /dev/null +++ b/tests/unit/runtime/zero/test_muon_use_muon_under_zero_init.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""`use_muon` has to be decided on the layer's shape, not on the ZeRO-3 partition's. + +`deepspeed.zero.Init` replaces a partitioned parameter's data with a flat placeholder and +records the shape it has as a layer in `ds_shape`. Muon is applied to matrices, so a check +against `param.shape` sees a 1-D tensor for every parameter in the model, tags none of them, +and ZeRO-3 then finds no sub-group using Muon. Training continues with every parameter on the +AdamW branch and nothing says so. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import deepspeed +from deepspeed.runtime.config import MUON_OPTIMIZER +from unit.common import DistributedTest + + +class _Model(torch.nn.Module): + + def __init__(self, hidden=16): + super().__init__() + self.attn = torch.nn.Linear(hidden, hidden, bias=False) + self.mlp = torch.nn.Linear(hidden, hidden, bias=False) + self.norm = torch.nn.Parameter(torch.ones(hidden)) + self.embed_tokens = torch.nn.Embedding(8, hidden) + + def forward(self, x): + return (self.mlp(self.attn(x)) * self.norm).square().sum() + + +def _muon_flags(model): + config = SimpleNamespace(optimizer_name=MUON_OPTIMIZER, optimizer_params={}) + deepspeed.set_optimizer_flags(config, model) + return {name: p.use_muon for name, p in model.named_parameters()} + + +def _partition_like_zero3(model): + """What `zero.Init` leaves behind: a flat placeholder plus the real shape on `ds_shape`.""" + for p in model.parameters(): + p.ds_shape = torch.Size(p.shape) + p.data = torch.zeros(0, dtype=p.dtype) + + +def test_matrices_are_tagged_when_the_model_is_not_partitioned(): + flags = _muon_flags(_Model()) + + assert flags["attn.weight"] is True + assert flags["mlp.weight"] is True + assert flags["norm"] is False + assert flags["embed_tokens.weight"] is False + + +def test_matrices_are_still_tagged_once_zero3_has_partitioned_them(): + model = _Model() + _partition_like_zero3(model) + + flags = _muon_flags(model) + + assert model.attn.weight.ndim == 1, "the partitioned parameter really is 1-D here" + assert flags["attn.weight"] is True + assert flags["mlp.weight"] is True + + +def test_the_partitioned_shape_does_not_promote_a_vector(): + """A 1-D parameter stays off Muon; `ds_shape` is read for its rank, not assumed to be >= 2.""" + model = _Model() + _partition_like_zero3(model) + + flags = _muon_flags(model) + + assert flags["norm"] is False + assert flags["embed_tokens.weight"] is False, "the name exclusions still apply" + + +class TestMuonRunsUnderZeroInit(DistributedTest): + """The end-to-end consequence: `muon_update` is reached at all.""" + world_size = 1 + + @pytest.mark.parametrize("zero_init", [False, True]) + def test_muon_update_is_called(self, zero_init): + import deepspeed.runtime.zero.stage3 as stage3 + + config = { + "train_micro_batch_size_per_gpu": 1, + "gradient_accumulation_steps": 1, + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": 3, + "reduce_scatter": False + }, + "optimizer": { + "type": "Muon", + "params": { + "lr": 1e-3 + } + }, + } + + if zero_init: + with deepspeed.zero.Init(config_dict_or_path=config): + model = _Model() + else: + model = _Model() + + calls = [] + original = stage3.muon_update + + def counting(*args, **kwargs): + calls.append(1) + return original(*args, **kwargs) + + stage3.muon_update = counting + try: + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + x = torch.ones(1, 16, dtype=torch.bfloat16, device=engine.device) + engine.backward(engine(x)) + engine.step() + finally: + stage3.muon_update = original + + assert calls, "Muon never ran; every parameter was left on the AdamW branch"