DeepSpeed

431f8ace - Give Muon's momentum the dtype of the gradient it is combined with (#8483)

Give Muon's momentum the dtype of the gradient it is combined with (#8483) ## The bug `muon_update` does `momentum.lerp_(grad)`, which requires both operands to have the same dtype. The momentum buffer is allocated in the configured gradient accumulation dtype — but the gradients only arrive in that dtype while `use_grad_accum_attribute` is on: ```python self.use_separate_grad_accum = (self.dtype != self.gradient_accumulation_dtype) if self.use_separate_grad_accum and not self.partition_gradients: self.use_grad_accum_attribute = True ``` `partition_gradients` is true at stage 2, so the attribute path is off there and `get_param_gradient_attribute` returns `param.grad` in the **parameter** dtype instead. bf16 with `grad_accum_dtype: fp32` at ZeRO stage 2 therefore cannot run Muon at all: ``` RuntimeError: expected dtype torch.float32 for `end`, but got dtype torch.bfloat16 ``` The dtype now comes from the gradient rather than from the configuration. ## Why #8433 did not already cover this There are two copies of the buffer setup. `get_flat_partition` delegates to `_get_flat_partition_unpadded` for any group with no round-robin padding: ```python if not any(self.round_robin_bit16_padding[param_group_idx]): return self._get_flat_partition_unpadded(...) ``` #8433 reconciled a restored checkpoint's buffer dtype in the first copy. **The second never got it** — and it is the one this configuration takes. I found this by adding a debug print to the code I assumed was running and watching it never fire; the traceback goes through `_get_flat_partition_unpadded`, not `get_flat_partition`. So rather than add the same reconciliation a second time — duplication is what caused the miss — both paths now call one `_muon_momentum_buffer`, which allocates in the gradient's dtype and converts a restored buffer into it. **That means this also repairs #8433's own case on the unpadded path.** Its regression tests fail on current `master`, which carries #8433 (`b726f4ed`): ``` master 4 failed, 2 passed TestMuonCheckpointRoundTrip::test_resumes_with_its_momentum branch 6 passed master's failure: RuntimeError: expected dtype torch.float32 for `end`, but got dtype torch.bfloat16 ``` I verified `master` really does contain #8433 before believing that, rather than assuming a stale checkout. ## Testing 2×H20. New case, `TestMuonSeparateGradAccumDtype`: ``` branch 1 passed master 1 failed ``` It parametrizes stage 2 only — stage 1 with the same combination selects `BF16_Optimizer`, which refuses Muon for a different reason (#8461). My first version parametrized both and failed on the branch too; that was the test's fault, not the fix's. `tests/unit/ops/muon/`, comparing failure *sets*: ``` master 70 failed, 83 passed branch 66 failed, 87 passed only failing on branch (regressions) -> none only failing on master -> the 4 TestMuonCheckpointRoundTrip cases above failing on both -> 66 ``` The 66 are pre-existing on this box — `CUDAMismatchException` between the installed toolkit and the one torch was built against, unrelated to this change. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>