fix(optim): create Muon momentum in optim_dtype, not the gradient dtype

[StarrickLiu]

The momentum buffer was created as grad.clone(), so under mixed precision it inherited the bf16 dtype of _reduced_grad. The checkpoint load path in checkpoint.set_optimizer_state_dict instead casts to _optim_dtype, which defaults to torch.float32 in both the ddp and fsdp2 param backends. The two paths therefore disagreed:

continuous run : momentum is bfloat16 (grad.clone())
saved : upcast to float32 (lossless)
after resume : momentum is float32 (cast to _optim_dtype)

Measured on a 2-rank Qwen3.5 run, reading buf.dtype directly at the first Muon step: 167 of 172 owner-managed matrices were bfloat16 in a continuous run and float32 in a run resumed from a checkpoint (the other 5 have an fp32 _reduced_grad, so both paths already agreed).

Two consequences:

1. Resuming silently changed the precision of the momentum accumulator. Values match at load time (bf16 -> fp32 is lossless), so the first step after a resume reproduces exactly; the first momentum_buffer.mul_(mu).add_(grad) then rounds differently and the run diverges from an uninterrupted one for good. This made checkpoint/restart non-reproducible, which is what led us here: after fixing it, a run resumed from a checkpoint reproduces the continuous trajectory bit-for-bit (14/14 logged steps, relative difference 0.00e+00 on train_loss, flow_loss, grad_norm, org_vlm_grad_norm and dmuon_adamw_grad_norm at full precision).

2. Independently of resuming, every mixed-precision Muon run has been keeping the momentum accumulator at 8 mantissa bits. buf = mu*buf + g is a long-memory accumulator; bf16 drops low-order bits on every step.

_optim_dtype is the declared intent (it is the default the param backends set and what the load path restores), so the create path is the side that was wrong. All three momentum initialisation sites -- including the two on the TP path -- now resolve the dtype through one helper that mirrors the load path's resolution order exactly, so the two can no longer drift apart.

No measurable step-time cost in the run above (5.1-6.1 s/step before and after); the extra memory is one fp32 copy per owner-managed Muon matrix.