[muon] Reconcile the momentum dtype when a checkpoint is restored

[alanhuangyoo]

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:

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.