DeepSpeed

b726f4ed - [muon] Reconcile the momentum dtype when a checkpoint is restored (#8433)

[muon] Reconcile the momentum dtype when a checkpoint is restored (#8433) Closes #7746. ## Problem Resuming a Muon run under ZeRO 1/2 raises as soon as the first gradient arrives: ``` RuntimeError: expected dtype torch.float32 for `end`, but got dtype torch.bfloat16 ``` 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 fp32 while the gradients it is combined with are not, and `muon_update`'s `momentum.lerp_(grad)` requires the two to match. `get_flat_partition` only asked whether a buffer existed: ```python if "momentum_buffer" not in self.optimizer.state[flatten_copy] and ...: self.optimizer.state[flatten_copy]["momentum_buffer"] = ... ``` so it left the restored one alone. Traced across a resume: ``` fresh run momentum torch.bfloat16 after initialize no buffer yet after load momentum torch.float32 <- crashes on the next step ``` ## Fix The buffer is converted to the dtype it is about to be combined with. Discarding and reallocating it would also stop the crash, but the momentum a resume just restored is the reason the checkpoint carries it — starting from zero is a different run, and Muon's update is built on that momentum. ## Verification 2 × H20, ZeRO 1/2/3, an uninterrupted run against one interrupted at the halfway point and resumed in a fresh engine: | stage | master | this branch | | --- | --- | --- | | 1 | `RuntimeError` | momentum restored (0.003303 → 0.003303), max param norm drift **0.0** | | 2 | `RuntimeError` | momentum restored (0.003303 → 0.003303), max param norm drift **0.0** | | 3 | passes | unchanged | Drift 0.0 is the part worth noting: the resumed run does not merely survive, it produces the same parameters as the run it resumed. Stage 3 keeps its momentum through a different path (`_create_momentum_buffer`) and was never affected. ## Tests `tests/unit/ops/muon/test_muon_checkpoint.py`, ZeRO 1/2/3 × bf16/fp16. Four of the six fail on the parent commit: ``` FAILED test_resumes_with_its_momentum[bf16-1] FAILED test_resumes_with_its_momentum[bf16-2] FAILED test_resumes_with_its_momentum[fp16-1] FAILED test_resumes_with_its_momentum[fp16-2] 4 failed, 2 passed ``` and all six pass here. It asserts both halves: that the restored momentum matches what was saved, and that continuing from the checkpoint lands on the same parameters as never having stopped. It is a new file rather than a case in `test_muon.py` because that module skips itself entirely where fp16 is unsupported, which would take the bf16 coverage with it. ## Why the suite could not have caught this `tests/unit/ops/muon/test_muon.py` never saves or loads a checkpoint. It also runs fp16 at the default loss scale, where — separately from this PR, filed as #8432 — every step overflows, so no momentum buffer is ever allocated in a dtype that could disagree with anything. The fp16 cases here set `initial_scale_power: 4` for that reason: this test is about what a checkpoint carries, not about the scaler. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>