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..131d062aae78 --- /dev/null +++ b/tests/unit/ops/muon/test_muon_overflow.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""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 +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), "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" + + +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" + + +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")