diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index f05a53867c93..118b42e03eee 100644 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -2149,12 +2149,23 @@ def get_flat_partition(self, if (not self.optimizer.state[flatten_copy]) and getattr( tensor_list[0], 'use_muon', False) and 'muon' in self.optimizer.__class__.__name__.lower(): self.optimizer.state[flatten_copy] = {} - if "momentum_buffer" not in self.optimizer.state[flatten_copy] and getattr( - tensor_list[0], 'use_muon', False) and 'muon' in self.optimizer.__class__.__name__.lower(): - # need to check the total # of elements in the parameters in this group and this partition - total_size = sum([t.numel() for t in tensor_list]) - flatten_bf_list = [torch.zeros([total_size], dtype=dtype, device=device)] - self.optimizer.state[flatten_copy]["momentum_buffer"] = self.flatten(flatten_bf_list) + if getattr(tensor_list[0], 'use_muon', False) and 'muon' in self.optimizer.__class__.__name__.lower(): + momentum_buffer = self.optimizer.state[flatten_copy].get("momentum_buffer") + if momentum_buffer is None: + # need to check the total # of elements in the parameters in this group and this partition + total_size = sum([t.numel() for t in tensor_list]) + flatten_bf_list = [torch.zeros([total_size], dtype=dtype, device=device)] + self.optimizer.state[flatten_copy]["momentum_buffer"] = self.flatten(flatten_bf_list) + elif momentum_buffer.dtype != dtype: + # A restored buffer arrives in the dtype the checkpoint holds optimizer state + # in, which is fp32, while the gradients it is combined with are in the + # gradient accumulation dtype. muon_update does momentum.lerp_(grad), which + # requires both to match, so resuming a bf16 run raised: + # RuntimeError: expected dtype torch.float32 for `end`, but got dtype + # torch.bfloat16 + # Convert rather than reallocate: the momentum a resume just restored is the + # reason the checkpoint carries it. + self.optimizer.state[flatten_copy]["momentum_buffer"] = momentum_buffer.to(dtype=dtype, device=device) buffer_idx = 0 for i, tensor in enumerate(tensor_list): diff --git a/tests/unit/ops/muon/test_muon_checkpoint.py b/tests/unit/ops/muon/test_muon_checkpoint.py new file mode 100644 index 000000000000..785cc5b9639b --- /dev/null +++ b/tests/unit/ops/muon/test_muon_checkpoint.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Muon has to survive a checkpoint round trip, in the dtype the run is using. + +ZeRO 1/2 keep Muon's momentum in a hand-flattened tensor under +``optimizer.state[flatten_copy]["momentum_buffer"]``, allocated in the gradient +accumulation dtype. A checkpoint stores optimizer state in fp32, so a restored buffer +comes back in a different dtype than the gradients it is combined with, and +``muon_update``'s ``momentum.lerp_(grad)`` requires the two to match. + +The existing Muon suite covers neither half of this: it runs in fp16, where the two +dtypes coincide, and it never saves or loads. +""" + +import pytest +import torch + +import deepspeed +from deepspeed.accelerator import get_accelerator +from unit.common import DistributedTest +from unit.simple_model import SimpleModel + + +def _config(zero_stage, dtype): + return { + "train_micro_batch_size_per_gpu": + 2, + "gradient_accumulation_steps": + 1, + "steps_per_print": + 10**9, + dtype: ({ + "enabled": True, + # A loss scale that overflows spends the run backing off instead of stepping, + # and this test is about what a checkpoint carries, not about the scaler. + "initial_scale_power": 4, + } if dtype == "fp16" else { + "enabled": True + }), + "zero_optimization": { + "stage": zero_stage, + # Muon rejects reduce-scatter, as the rest of the suite does too. + "reduce_scatter": False, + }, + "optimizer": { + "type": "Muon", + "params": { + "lr": 1e-3, + "momentum": 0.95 + } + }, + } + + +def _momentum_norm(engine): + optimizer = getattr(engine.optimizer, "optimizer", engine.optimizer) + total = 0.0 + for state in optimizer.state.values(): + buffer = state.get("momentum_buffer") if isinstance(state, dict) else None + if buffer is not None: + total += float(buffer.float().norm()**2) + return total**0.5 + + +@pytest.mark.parametrize("zero_stage", [1, 2, 3]) +@pytest.mark.parametrize("dtype", ["bf16", "fp16"]) +class TestMuonCheckpointRoundTrip(DistributedTest): + world_size = 2 + + def test_resumes_with_its_momentum(self, tmpdir, zero_stage, dtype): + """A resumed run must continue the one it resumed, not restart its momentum. + + Before the dtype was reconciled this raised on stages 1 and 2 under bf16: + + RuntimeError: expected dtype torch.float32 for `end`, + but got dtype torch.bfloat16 + """ + torch_dtype = torch.bfloat16 if dtype == "bf16" else torch.half + if torch_dtype not in get_accelerator().supported_dtypes(): + pytest.skip(f"{dtype} not supported by {get_accelerator().device_name()}") + + hidden_dim, steps = 32, 4 + ckpt_dir = str(tmpdir) + + def build(): + torch.manual_seed(1234) + model = SimpleModel(hidden_dim, nlayers=3) + engine, _, _, _ = deepspeed.initialize(model=model, + model_parameters=model.parameters(), + config=_config(zero_stage, dtype)) + return engine + + uninterrupted = build() + batch = torch.randn(2, hidden_dim, device=uninterrupted.device, dtype=torch_dtype) + labels = torch.randn(2, hidden_dim, device=uninterrupted.device, dtype=torch_dtype) + + for _ in range(steps): + uninterrupted.backward(uninterrupted(batch, labels)) + uninterrupted.step() + momentum_at_save = _momentum_norm(uninterrupted) + uninterrupted.save_checkpoint(ckpt_dir, tag="mid") + for _ in range(steps): + uninterrupted.backward(uninterrupted(batch, labels)) + uninterrupted.step() + straight_through = [float(p.float().norm()) for p in uninterrupted.module.parameters()] + + resumed = build() + resumed.load_checkpoint(ckpt_dir, tag="mid") + assert _momentum_norm(resumed) == pytest.approx(momentum_at_save, rel=1e-3), \ + "the checkpoint carries the momentum; a resume that drops it is a different run" + + for _ in range(steps): + resumed.backward(resumed(batch, labels)) + resumed.step() + after_resume = [float(p.float().norm()) for p in resumed.module.parameters()] + + for straight, restored in zip(straight_through, after_resume): + assert straight == pytest.approx(restored, abs=1e-4)