From 0a7bb9c809587291807d27e19fecdc4d6c8f752e Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 6 Sep 2026 22:17:55 +0800 Subject: [PATCH 1/3] Run Newton-Schulz when no ZeRO optimizer has MuonWithAuxAdam.step applied an update it assumed had been orthogonalized already. That holds under ZeRO, where the parameters it sees are flat partitions and get_flat_partition or the ZeRO-3 sub-group loop did the work. With no ZeRO optimizer nothing did, and p.add_(p.grad, alpha=-lr) on a raw gradient is SGD. zero_optimization.stage defaults to 0, so a config that just names Muon got that. Counting the Newton-Schulz calls on one step, before: no zero_optimization, fp32 MuonWithAuxAdam 0 max|w-SGD| 1.5e-08 stage 0, bf16 FP16_UnfusedOptimizer 0 max|w-SGD| 4.9e-04 stage 0, fp16 FP16_UnfusedOptimizer 0 stage 1, fp32 DeepSpeedZeroOptimizer 2 max|w-SGD| 9.8e-02 Training ran and the loss fell either way. The two cases are distinguishable by shape: ZeRO hands step() a flat 1-D partition, while an unwrapped optimizer and FP16_UnfusedOptimizer - which keeps per-parameter fp32 clones rather than a flat buffer - hand it the 2-D weight. So orthogonalize when the parameter is a matrix and keep applying the update as-is when it is a partition. After, stage 0 fp32 gives the same max|w-SGD| as stage 1, 9.772e-02. The existing Muon tests parametrize stages 1, 2 and 3 and assert that training progresses, which SGD also does. The new tests count the orthogonalizations, and count them around the training step only, since FP16_UnfusedOptimizer also steps once at construction to allocate state. Signed-off-by: alanhuangyoo --- deepspeed/runtime/zero/muon/muon_optimizer.py | 23 +++- .../zero/test_muon_without_zero_optimizer.py | 129 ++++++++++++++++++ 2 files changed, 149 insertions(+), 3 deletions(-) create mode 100644 tests/unit/runtime/zero/test_muon_without_zero_optimizer.py diff --git a/deepspeed/runtime/zero/muon/muon_optimizer.py b/deepspeed/runtime/zero/muon/muon_optimizer.py index 3f199aeb97ea..ce129fd1c1a4 100644 --- a/deepspeed/runtime/zero/muon/muon_optimizer.py +++ b/deepspeed/runtime/zero/muon/muon_optimizer.py @@ -8,6 +8,7 @@ try: from deepspeed.runtime.zero.muon.original_muon import MuonWithAuxAdam as BaseMuonWithAuxAdam from deepspeed.runtime.zero.muon.original_muon import adam_update + from deepspeed.runtime.zero.muon.original_muon import muon_update except ImportError: pass @@ -60,11 +61,27 @@ def step(self, closure=None, step_id=None): loss = closure() for group in self.param_groups: if group["use_muon"]: - # we move the muon update part to the deepspeed's optimizer since the parameter here is a flat version - # thus not suitable for muon update for p in group["params"]: + if p.grad is None: + p.grad = torch.zeros_like(p) # force synchronization + if p.dim() < 2: + # A flat ZeRO partition. ZeRO 1/2 orthogonalizes in get_flat_partition and + # ZeRO-3 in its sub-group loop, so the gradient already holds the update + # and only the weight decay and step size are left to apply. + update = p.grad + else: + # The weight itself, so nothing has orthogonalized it: no ZeRO optimizer is + # in play. Muon has to run here or the step degenerates to SGD. + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p) + update = muon_update(p.grad, + state["momentum_buffer"], + beta=group["momentum"], + ns_method=group.get("ns_method", "gram"), + is_expert_group=getattr(p, "is_expert_group", False)) p.mul_(1 - group["lr"] * group["weight_decay"]) - p.add_(p.grad.reshape(p.shape), alpha=-group["lr"]) + p.add_(update.reshape(p.shape), alpha=-group["lr"]) aux_param_groups = [group for group in self.param_groups if not group["use_muon"]] if self.aux_optimizer is not None: diff --git a/tests/unit/runtime/zero/test_muon_without_zero_optimizer.py b/tests/unit/runtime/zero/test_muon_without_zero_optimizer.py new file mode 100644 index 000000000000..2109aa0caeeb --- /dev/null +++ b/tests/unit/runtime/zero/test_muon_without_zero_optimizer.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Newton-Schulz has to run whether or not a ZeRO optimizer is there to do it. + +`MuonWithAuxAdam.step` applied an update it assumed had been orthogonalized already, which holds +under ZeRO - the parameters it sees there are flat partitions and `get_flat_partition` or the +ZeRO-3 sub-group loop did the work. At stage 0, the default, no ZeRO optimizer exists to have +done it, and applying a raw gradient is SGD. Training runs and the loss falls, which is why +counting the Newton-Schulz calls is the assertion that means something here. +""" + +import contextlib + +import pytest +import torch + +import deepspeed +import deepspeed.runtime.zero.muon.original_muon as original_muon +from unit.common import DistributedTest + +NS_KERNELS = ("zeropower_via_gram_newtonschulz", "zeropower_via_newtonschulz5") + + +@contextlib.contextmanager +def counting_newton_schulz(): + """Counts every Newton-Schulz call, whichever kernel the config selects. + + Patched inside the test body rather than in a fixture: `DistributedTest` runs the body in a + worker process that a fixture in the parent would not reach. + """ + calls = [] + originals = {name: getattr(original_muon, name) for name in NS_KERNELS} + + def counted(kernel): + + def wrapper(*args, **kwargs): + calls.append(1) + return kernel(*args, **kwargs) + + return wrapper + + for name, kernel in originals.items(): + setattr(original_muon, name, counted(kernel)) + try: + yield calls + finally: + for name, kernel in originals.items(): + setattr(original_muon, name, kernel) + + +def _model(): + return torch.nn.Sequential(torch.nn.Linear(32, 32, bias=False), torch.nn.Linear(32, 32, bias=False)) + + +def _config(stage, dtype="fp32"): + config = { + "train_micro_batch_size_per_gpu": 2, + "gradient_accumulation_steps": 1, + "gradient_clipping": 0.0, + "optimizer": { + "type": "Muon", + "params": { + "lr": 0.02 + } + }, + } + if stage is not None: + config["zero_optimization"] = {"stage": stage, "reduce_scatter": stage != 3} + if dtype != "fp32": + config[dtype] = {"enabled": True} + if dtype == "fp16": + config[dtype]["initial_scale_power"] = 4 + return config + + +class TestMuonRunsWithoutAZeroOptimizer(DistributedTest): + world_size = 1 + + @pytest.mark.parametrize("dtype", ["fp32", "bf16", "fp16"]) + def test_newton_schulz_runs_at_stage_zero(self, dtype): + """Every stage-0 wrapper: unwrapped for fp32, FP16_UnfusedOptimizer for bf16 and fp16. + + Each hands `step` the weight itself rather than a flat partition, so nothing upstream has + orthogonalized it. On master all three do zero orthogonalizations and train as SGD. + """ + model = _model() + engine, _, _, _ = deepspeed.initialize(model=model, + model_parameters=model.parameters(), + config=_config(0, dtype)) + + # Counting starts after initialize: FP16_UnfusedOptimizer steps once at construction to + # allocate state, and that call must not be what the assertion below is satisfied by. + with counting_newton_schulz() as calls: + x = torch.ones(2, 32, device=engine.device, dtype=next(engine.module.parameters()).dtype) + engine.backward(engine(x).square().sum()) + engine.step() + + assert len(calls) == 2, f"Newton-Schulz ran {len(calls)} times for two Muon matrices; expected one each" + + def test_the_default_config_runs_muon(self): + """`zero_optimization.stage` defaults to 0, so this is the plainest Muon config there is.""" + model = _model() + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=_config(None)) + + with counting_newton_schulz() as calls: + x = torch.ones(2, 32, device=engine.device) + engine.backward(engine(x).square().sum()) + engine.step() + + assert len(calls) == 2, f"Newton-Schulz ran {len(calls)} times; the default config trained as SGD" + + @pytest.mark.parametrize("stage", [1, 2, 3]) + def test_newton_schulz_runs_on_the_supported_stages(self, stage): + """The positive control, and the assertion the existing tests were missing. + + They check that training progresses, which SGD does too. Counting the orthogonalizations + is what distinguishes Muon from the update it degenerates to. + """ + model = _model() + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=_config(stage)) + + with counting_newton_schulz() as calls: + x = torch.ones(2, 32, device=engine.device, dtype=next(engine.module.parameters()).dtype) + engine.backward(engine(x).square().sum()) + engine.step() + + assert len(calls) == 2, \ + f"Newton-Schulz ran {len(calls)} times for two Muon matrices; the step was not Muon" From 9adc53719bcb1925603f646625e9b893a68417ae Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Mon, 7 Sep 2026 12:47:31 +0800 Subject: [PATCH 2/3] Refuse Muon with BF16_Optimizer, which the shape test cannot catch MuonWithAuxAdam.step tells the two cases apart by shape: a matrix is the weight itself and is orthogonalized there, a 1-D tensor is a ZeRO partition whose update the ZeRO optimizer already applied. BF16_Optimizer breaks that reading - it replaces the param groups with flat fp32 partitions and knows nothing about use_muon, so the update is never applied and the shape test reads its partitions as already done. Enumerating every wrapper a Muon config can select, before this branch: s0-fp32 MuonWithAuxAdam ndims=[2] NS=0 s0-bf16 FP16_UnfusedOptimizer ndims=[2] NS=0 s0-fp16 FP16_UnfusedOptimizer ndims=[2] NS=0 s1-bf16 DeepSpeedZeroOptimizer ndims=[1] NS=2 s1-bf16-ga32 BF16_Optimizer ndims=[1] NS=0 s1-fp16 DeepSpeedZeroOptimizer ndims=[1] NS=2 s2-bf16 DeepSpeedZeroOptimizer ndims=[1] NS=2 s3-bf16 DeepSpeedZeroOptimizer_S3 ndims=[] NS=2 s1-bf16-ga32 is bf16 with grad_accum_dtype fp32 at stage 1, and it was broken before this branch in the same silent way: max|w - SGD| of 4.9e-04, bf16 rounding away from plain SGD, against 6.8e-02 for the same config one flag apart. The original shapes are not recoverable from a flat partition, so this refuses the combination at initialize rather than fixing it; implementing Muon inside BF16_Optimizer is a separate change. Signed-off-by: alanhuangyoo --- deepspeed/runtime/engine.py | 22 ++++++++++++ .../zero/test_muon_without_zero_optimizer.py | 34 +++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 86918bd71c5a..2932335b09bb 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -2123,6 +2123,7 @@ def _configure_optimizer(self, client_optimizer, model_parameters): log_dist(f"DeepSpeed Basic Optimizer = {basic_optimizer.__class__.__name__}", ranks=[0]) optimizer_wrapper = self._do_optimizer_sanity_check(basic_optimizer) + self._check_muon_can_reach_its_parameters(basic_optimizer, optimizer_wrapper) if optimizer_wrapper == ZERO_OPTIMIZATION: self.optimizer = self._configure_zero_optimizer(basic_optimizer) @@ -2147,6 +2148,27 @@ def _configure_optimizer(self, client_optimizer, model_parameters): self.compression_scheduler = self._configure_compression_scheduler() self.quantizer = self._configure_quantization() + def _check_muon_can_reach_its_parameters(self, basic_optimizer, optimizer_wrapper): + """Refuse the one wrapper that hands Muon flat partitions and does not orthogonalize them. + + `MuonWithAuxAdam.step` tells the two cases apart by shape: a matrix is the weight itself + and is orthogonalized there, a 1-D tensor is a ZeRO partition whose update the ZeRO + optimizer already applied. `BF16_Optimizer` breaks that reading - it replaces the param + groups with flat fp32 partitions (`param_group['params'] = [self.fp32_groups_flat_partition[i]]`) + and knows nothing about `use_muon`, so the update is never applied and the step is SGD. + + The original shapes are not recoverable from `step`, so this is a refusal rather than a + fix; implementing Muon inside BF16_Optimizer is its own change. Reached by bf16 with + `grad_accum_dtype: fp32` at ZeRO stage 1. + """ + if not isinstance(basic_optimizer, MuonWithAuxAdam) or optimizer_wrapper != BFLOAT16: + return + raise ZeRORuntimeException( + "Muon cannot be used with the BF16_Optimizer, which this configuration selects: bf16 " + "with grad_accum_dtype fp32 at ZeRO stage 1. That optimizer hands Muon flat fp32 " + "partitions and never applies the Newton-Schulz update, so training would silently " + "proceed as SGD. Drop grad_accum_dtype, or use ZeRO stage 2 or 3.") + def _configure_autoep_folding_optimizer_gradient_reduction(self): configure = getattr(self.optimizer, "configure_autoep_folding_tp_gradient_reduction", None) if configure is None: diff --git a/tests/unit/runtime/zero/test_muon_without_zero_optimizer.py b/tests/unit/runtime/zero/test_muon_without_zero_optimizer.py index 2109aa0caeeb..42c152408930 100644 --- a/tests/unit/runtime/zero/test_muon_without_zero_optimizer.py +++ b/tests/unit/runtime/zero/test_muon_without_zero_optimizer.py @@ -17,6 +17,7 @@ import deepspeed import deepspeed.runtime.zero.muon.original_muon as original_muon +from deepspeed.runtime.zero.utils import ZeRORuntimeException from unit.common import DistributedTest NS_KERNELS = ("zeropower_via_gram_newtonschulz", "zeropower_via_newtonschulz5") @@ -127,3 +128,36 @@ def test_newton_schulz_runs_on_the_supported_stages(self, stage): assert len(calls) == 2, \ f"Newton-Schulz ran {len(calls)} times for two Muon matrices; the step was not Muon" + + +class TestMuonRefusesBF16Optimizer(DistributedTest): + """The one wrapper that hands Muon flat partitions without orthogonalizing them. + + `BF16_Optimizer` replaces the param groups with flat fp32 partitions and knows nothing about + `use_muon`, so the shape test in `step` reads them as "ZeRO already did the update" and the + step is SGD. The original shapes are not recoverable there, so this is refused rather than + fixed. Selected by bf16 with `grad_accum_dtype: fp32` at ZeRO stage 1. + """ + world_size = 1 + + def test_bf16_optimizer_with_muon_is_refused(self): + model = _model() + config = _config(1, "bf16") + config["data_types"] = {"grad_accum_dtype": "fp32"} + + with pytest.raises(ZeRORuntimeException, match="BF16_Optimizer"): + deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + + def test_the_same_config_without_grad_accum_dtype_still_runs_muon(self): + """The neighbouring config, so the refusal is shown to be narrow.""" + model = _model() + engine, _, _, _ = deepspeed.initialize(model=model, + model_parameters=model.parameters(), + config=_config(1, "bf16")) + + with counting_newton_schulz() as calls: + x = torch.ones(2, 32, device=engine.device, dtype=next(engine.module.parameters()).dtype) + engine.backward(engine(x).square().sum()) + engine.step() + + assert len(calls) == 2 From c8f34171f99aadb42ad4f3218829845e67ef9655 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Tue, 8 Sep 2026 20:22:18 +0800 Subject: [PATCH 3/3] Skip the mixed-precision cases where the accelerator refuses the dtype cpu-torch-latest fails test_newton_schulz_runs_at_stage_zero[fp16] with ValueError: Type fp16 is not supported on your device. which _do_sanity_check raises on not get_accelerator().is_fp16_supported(). That is a different predicate from supported_dtypes(): the runner reports fp16 in the latter and False from the former, so guarding on the usual one would still fail there. _skip_if_unsupported mirrors the engine's own check, and the two bf16 cases in TestMuonRefusesBF16Optimizer get the same guard. Verified on 1xH20 (9 passed) and on the CPU accelerator, where both predicates return True so the skip does not fire -- I could not reproduce the runner's False branch locally, so the guard is matched to the raising condition rather than to an observed skip. Signed-off-by: alanhuangyoo --- .../zero/test_muon_without_zero_optimizer.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/unit/runtime/zero/test_muon_without_zero_optimizer.py b/tests/unit/runtime/zero/test_muon_without_zero_optimizer.py index 42c152408930..5b6eede9a79f 100644 --- a/tests/unit/runtime/zero/test_muon_without_zero_optimizer.py +++ b/tests/unit/runtime/zero/test_muon_without_zero_optimizer.py @@ -16,6 +16,7 @@ import torch import deepspeed +from deepspeed.accelerator import get_accelerator import deepspeed.runtime.zero.muon.original_muon as original_muon from deepspeed.runtime.zero.utils import ZeRORuntimeException from unit.common import DistributedTest @@ -75,6 +76,22 @@ def _config(stage, dtype="fp32"): return config +def _skip_if_unsupported(dtype): + """Mirror the check the engine itself makes. + + `_do_sanity_check` raises `Type fp16 is not supported on your device.` on + `not get_accelerator().is_fp16_supported()`, which is a different predicate from + `supported_dtypes()` -- the cpu-torch-latest runner reports fp16 in the latter and + False from the former, so guarding on the wrong one still fails there. + """ + supported = { + "fp16": get_accelerator().is_fp16_supported, + "bf16": get_accelerator().is_bf16_supported, + }.get(dtype) + if supported is not None and not supported(): + pytest.skip(f"{dtype} not supported on this accelerator") + + class TestMuonRunsWithoutAZeroOptimizer(DistributedTest): world_size = 1 @@ -85,6 +102,7 @@ def test_newton_schulz_runs_at_stage_zero(self, dtype): Each hands `step` the weight itself rather than a flat partition, so nothing upstream has orthogonalized it. On master all three do zero orthogonalizations and train as SGD. """ + _skip_if_unsupported(dtype) model = _model() engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), @@ -141,6 +159,7 @@ class TestMuonRefusesBF16Optimizer(DistributedTest): world_size = 1 def test_bf16_optimizer_with_muon_is_refused(self): + _skip_if_unsupported("bf16") model = _model() config = _config(1, "bf16") config["data_types"] = {"grad_accum_dtype": "fp32"} @@ -150,6 +169,7 @@ def test_bf16_optimizer_with_muon_is_refused(self): def test_the_same_config_without_grad_accum_dtype_still_runs_muon(self): """The neighbouring config, so the refusal is shown to be narrow.""" + _skip_if_unsupported("bf16") model = _model() engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(),