From 1f1c422fc9e5b899324d2ec1812abdba210c4776 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 6 Sep 2026 17:57:59 +0800 Subject: [PATCH 1/2] [muon] Keep the momentum out of steps the loss scaler discards Muon under ZeRO 1/2 with fp16 does not train. The first loss-scale overflow is folded into the momentum buffer before the overflow check decides to discard the step, the buffer stays non-finite for the rest of the run, and every later step overflows too until the scaler gives up: Exception: Current loss scale already at minimum - cannot decrease scale anymore. Exiting run. No parameter is ever updated. This is the configuration test_muon.py itself uses, only run for longer than five steps. The failure sustains itself because both halves of muon_update touch the gradient: momentum.lerp_(grad, 1 - beta) # inf/nan enters the momentum update = grad.lerp_(momentum, beta) if nesterov # ...and is written back to grad so the next step's gradient is already non-finite whatever the loss scale has been reduced to. The momentum now stays out of a step whose gradient is not finite, and the non-finite gradient is still returned so the overflow is seen and the step is skipped. Both are needed: on the non-nesterov path `update` is the momentum, so protecting the momentum alone would hand an overflowed step a finite update and the step would be applied instead of skipped. Evaluated on device, so this costs no synchronization. 30 steps, 2 x H20, ZeRO 1/2, SimpleModel(hidden_dim=128, nlayers=5), lr 0.05: stage scale master this commit 1 65536 0/10 moved, non-finite, dies 10/10 moved, finite 2 65536 0/10 moved, non-finite, dies 10/10 moved, finite 1 1 10/10 moved, finite 10/10 moved, finite 2 1 10/10 moved, finite 10/10 moved, finite The scale-1 rows are the control: with no overflow there was never a problem, so the failure is entirely the overflow interaction. Reported as #8432, which also records why the suite is green today: the run is too short to reach the exception, and the parameter-change assertion compares parameters captured before deepspeed.initialize -- fp32 -- against fp16 ones after training, so torch.equal is False whatever happened in between. Casting a model to fp16 and training it zero steps satisfies that assertion. Tests: 3 of the 4 new cases fail on the parent commit, including the unit-level one that pins the momentum directly. Signed-off-by: alanhuangyoo --- deepspeed/runtime/zero/muon/original_muon.py | 14 ++- tests/unit/ops/muon/test_muon_overflow.py | 104 +++++++++++++++++++ 2 files changed, 116 insertions(+), 2 deletions(-) create mode 100644 tests/unit/ops/muon/test_muon_overflow.py diff --git a/deepspeed/runtime/zero/muon/original_muon.py b/deepspeed/runtime/zero/muon/original_muon.py index 1cbc46392410..8884a8e18ed1 100644 --- a/deepspeed/runtime/zero/muon/original_muon.py +++ b/deepspeed/runtime/zero/muon/original_muon.py @@ -145,7 +145,15 @@ def zeropower_via_gram_newtonschulz(G, steps: int): @compiler.compile() def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method="gram", is_expert_group=False): orig_dtype = grad.dtype - momentum.lerp_(grad, 1 - beta) + # A step whose gradients overflowed is discarded by the loss scaler, but Muon folds the + # gradient into its momentum before that decision is made. Left alone, one overflow + # leaves the momentum non-finite for the rest of the run: with nesterov the blend is + # written back into the gradient in place, so the next step overflows too, and the + # scaler backs off until it raises "Current loss scale already at minimum". Keep the + # momentum out of it, and let the non-finite gradient through so the overflow is still + # seen and the step still skipped. Evaluated on device so this costs no synchronization. + grad_is_finite = torch.isfinite(grad).all() + momentum.copy_(torch.where(grad_is_finite, momentum.lerp(grad, 1 - beta), momentum)) update = grad.lerp_(momentum, beta) if nesterov else momentum if is_expert_group: ns_fn = zeropower_via_gram_newtonschulz if ns_method == "gram" else zeropower_via_newtonschulz5 @@ -161,7 +169,9 @@ def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method= update *= max(1, grad.size(-2) / grad.size(-1))**0.5 if update.dtype != orig_dtype: update = update.to(orig_dtype) - return update + # On the non-nesterov path `update` is the (untouched, finite) momentum, so without this + # an overflowed step would produce a finite update and be applied instead of skipped. + return torch.where(grad_is_finite, update, grad.to(orig_dtype)) class Muon(torch.optim.Optimizer): diff --git a/tests/unit/ops/muon/test_muon_overflow.py b/tests/unit/ops/muon/test_muon_overflow.py new file mode 100644 index 000000000000..3aef6da6891c --- /dev/null +++ b/tests/unit/ops/muon/test_muon_overflow.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""A step the loss scaler discards must leave Muon's momentum as it was. + +Muon folds the gradient into its momentum while the partition is filled, which happens +before the overflow check decides whether to keep the step. With nesterov the blend is +also written back into the gradient in place. One overflow would therefore leave the +momentum non-finite for the rest of the run and make every later step overflow too, +until the scaler reaches its minimum and raises. +""" + +import pytest +import torch + +import deepspeed +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime.zero.muon.original_muon import muon_update +from unit.common import DistributedTest +from unit.simple_model import SimpleModel + + +def test_an_overflowed_gradient_does_not_enter_the_momentum(): + """The unit of the behaviour, without a training loop around it.""" + device = get_accelerator().device_name() + grad = torch.randn(16, 16, device=device) + momentum = torch.randn(16, 16, device=device) + before = momentum.clone() + + overflowed = grad.clone() + overflowed[0, 0] = float("inf") + update = muon_update(overflowed, momentum) + + assert torch.equal(momentum, before), "an overflowed step must not move the momentum" + assert not torch.isfinite(update).all(), \ + "the update has to stay non-finite, or the overflow check will not skip the step" + + +def test_a_finite_gradient_still_moves_the_momentum(): + """The guard must not disable the optimizer.""" + device = get_accelerator().device_name() + grad = torch.randn(16, 16, device=device) + momentum = torch.zeros(16, 16, device=device) + + update = muon_update(grad.clone(), momentum) + + assert momentum.abs().sum() > 0 + assert torch.isfinite(update).all() + + +@pytest.mark.parametrize("zero_stage", [1, 2]) +class TestMuonSurvivesLossScaleBackoff(DistributedTest): + world_size = 2 + + def test_training_recovers_from_the_initial_overflow(self, zero_stage): + """fp16 starts at a loss scale that overflows; backing off is the normal path. + + On the parent commit this never recovers: the momentum is NaN from the first step, + the poisoned gradient keeps the overflow check firing, and DeepSpeed raises + "Current loss scale already at minimum - cannot decrease scale anymore". + """ + if torch.half not in get_accelerator().supported_dtypes(): + pytest.skip("fp16 not supported") + + hidden_dim, batch_size = 128, 8 + torch.manual_seed(0) + model = SimpleModel(hidden_dim=hidden_dim, nlayers=5) + config = { + "train_batch_size": batch_size, + "optimizer": { + "type": "Muon", + "params": { + "lr": 0.05 + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": True + }, + "zero_optimization": { + "stage": zero_stage, + "reduce_scatter": False + }, + } + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + + # Captured after initialize: before it the parameters are still fp32, and comparing + # across the fp16 cast makes any assertion about change trivially true. + before = [p.clone().cpu() for p in model.parameters()] + for _ in range(30): + x = torch.randn(batch_size, hidden_dim, device=engine.device, dtype=torch.half) + y = torch.randint(0, hidden_dim, (batch_size, ), device=engine.device) + engine.backward(engine(x, y)) + engine.step() + after = [p.clone().cpu() for p in model.parameters()] + + changed = sum(1 for b, a in zip(before, after) if not torch.equal(b, a)) + assert changed == len(before), f"only {changed}/{len(before)} parameters moved in 30 steps" + + optimizer = getattr(engine.optimizer, "optimizer", engine.optimizer) + for state in optimizer.state.values(): + buffer = state.get("momentum_buffer") if isinstance(state, dict) else None + if buffer is not None: + assert torch.isfinite(buffer.float()).all(), "the momentum did not survive the backoff" From 4ce8e623a3bfedcd2da25c26b8e792ead764364f Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Tue, 8 Sep 2026 17:59:07 +0800 Subject: [PATCH 2/2] Narrow the invariant to what the guard delivers, and pin the mixed case The module claimed 'A step the loss scaler discards must leave Muon's momentum as it was'. The guard is per tensor and the scaler's decision is global -- has_overflow reduces _has_inf_or_nan over every partitioned gradient and step() discards on that one flag -- so a tensor whose own gradient was finite still advances its momentum on a step discarded for another tensor. Measured on 1xH20, two 2-D parameters in one group, only one fed an overflowing input, fp16 + ZeRO-1, momentum read on either side of backward because that is where muon_update writes: step overflow calm momentum boom momentum params 0 False None -> 55.159618 None -> 42.780441 changed 1 True 55.159618 -> 104.801094 42.780441 -> 42.780441 discarded On master the same run gives boom 42.780441 -> inf, then nan, with every later step discarded and the loss scale halving to the minimum -- the failure this PR is for. calm's 55.159618 -> 104.801094 is identical on both sides. Reported by @ebarkhordar, who found it by reading and named which tests could not reach it: both tensor-level cases use a single tensor and the training test only asserts the momentum is finite. Docstring now says what the guard covers, and test_a_finite_tensor_still_absorbs_a_step_discarded_for_another records the gap so a later change to global scope shows up as a failing test rather than a silent improvement. Signed-off-by: alanhuangyoo --- tests/unit/ops/muon/test_muon_overflow.py | 102 +++++++++++++++++++++- 1 file changed, 100 insertions(+), 2 deletions(-) diff --git a/tests/unit/ops/muon/test_muon_overflow.py b/tests/unit/ops/muon/test_muon_overflow.py index 3aef6da6891c..131d062aae78 100644 --- a/tests/unit/ops/muon/test_muon_overflow.py +++ b/tests/unit/ops/muon/test_muon_overflow.py @@ -1,13 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team -"""A step the loss scaler discards must leave Muon's momentum as it was. +"""A tensor whose own gradient overflowed must not absorb it into its momentum. Muon folds the gradient into its momentum while the partition is filled, which happens before the overflow check decides whether to keep the step. With nesterov the blend is also written back into the gradient in place. One overflow would therefore leave the momentum non-finite for the rest of the run and make every later step overflow too, until the scaler reaches its minimum and raises. + +The guard is per tensor, and the loss scaler's decision is global: `has_overflow` reduces +`_has_inf_or_nan` over every partitioned gradient and `step()` discards the whole step on +that one flag. So a tensor whose own gradient was finite still advances its momentum on a +step discarded because some *other* tensor overflowed, and that update is thrown away. +`test_a_finite_tensor_still_absorbs_a_step_discarded_for_another` pins that, measured +rather than assumed; making it exact needs the momentum write deferred until the step is +known to survive, which costs a second buffer the size of the momentum. """ import pytest @@ -31,7 +39,7 @@ def test_an_overflowed_gradient_does_not_enter_the_momentum(): overflowed[0, 0] = float("inf") update = muon_update(overflowed, momentum) - assert torch.equal(momentum, before), "an overflowed step must not move the momentum" + assert torch.equal(momentum, before), "a tensor's own overflow must not move its momentum" assert not torch.isfinite(update).all(), \ "the update has to stay non-finite, or the overflow check will not skip the step" @@ -102,3 +110,93 @@ def test_training_recovers_from_the_initial_overflow(self, zero_stage): buffer = state.get("momentum_buffer") if isinstance(state, dict) else None if buffer is not None: assert torch.isfinite(buffer.float()).all(), "the momentum did not survive the backoff" + + +class TestMuonMixedOverflow(DistributedTest): + world_size = 1 + + def test_a_finite_tensor_still_absorbs_a_step_discarded_for_another(self): + """The scope of the guard, recorded rather than implied. + + Two 2-D parameters in one group; only `boom` is fed an input that overflows in + fp16. The step is discarded for the whole model, so neither parameter moves -- + but `calm`'s gradient was finite, so its momentum advances anyway, for an update + that is thrown away. Making that exact needs the momentum write deferred until + the step is known to survive; this pins today's behaviour so the gap is visible. + """ + if torch.half not in get_accelerator().supported_dtypes(): + pytest.skip("fp16 not supported") + + hidden_dim = 32 + numel = hidden_dim * hidden_dim + + class TwoMatrices(torch.nn.Module): + + def __init__(self): + super().__init__() + self.calm = torch.nn.Linear(hidden_dim, hidden_dim, bias=False) + self.boom = torch.nn.Linear(hidden_dim, hidden_dim, bias=False) + + def forward(self, calm_x, boom_x): + return self.calm(calm_x).sum() + self.boom(boom_x).sum() + + torch.manual_seed(0) + model = TwoMatrices() + config = { + "train_batch_size": 1, + "optimizer": { + "type": "Muon", + "params": { + "lr": 0.01, + "momentum": 0.9, + "weight_decay": 0.0 + } + }, + # low enough that a normal step does not overflow, so the only overflow is + # the one the test injects + "fp16": { + "enabled": True, + "initial_scale_power": 4 + }, + "zero_optimization": { + "stage": 1 + }, + "zero_allow_untested_optimizer": True, + } + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=list(model.parameters()), config=config) + + def momentum_halves(): + """(calm, boom) slices of the group's flat momentum buffer, in parameter order.""" + inner = getattr(engine.optimizer, "optimizer", engine.optimizer) + for state in inner.state.values(): + buffer = state.get("momentum_buffer") if isinstance(state, dict) else None + if buffer is not None and buffer.numel() >= 2 * numel: + flat = buffer.detach().float() + return flat[:numel].norm().item(), flat[numel:2 * numel].norm().item() + return None, None + + device = engine.device + calm_x = torch.randn(1, hidden_dim, device=device, dtype=torch.half) + finite_x = torch.randn(1, hidden_dim, device=device, dtype=torch.half) + overflowing_x = torch.full((1, hidden_dim), 6e4, device=device, dtype=torch.half) + + # Step 0 establishes a momentum for both; step 1 overflows only through `boom`. + for step in range(2): + # muon_update runs while the partition is filled, i.e. inside backward, so the + # momentum has to be read on either side of that rather than around step(). + calm_before, boom_before = momentum_halves() + params_before = [p.detach().float().norm().item() for p in model.parameters()] + engine.backward(engine(calm_x, overflowing_x if step == 1 else finite_x)) + calm_after, boom_after = momentum_halves() + engine.step() + params_after = [p.detach().float().norm().item() for p in model.parameters()] + + if step == 1: + assert engine.optimizer.overflow, "step 1 was meant to overflow" + assert params_before == params_after, "an overflowed step must not move parameters" + assert boom_before == boom_after, \ + "the tensor whose own gradient overflowed must not absorb it" + assert calm_before != calm_after, ( + "a tensor whose gradient was finite does advance its momentum on a step " + "discarded for another tensor -- if this starts failing, the guard has " + "become global and the module docstring should say so")