Carry Muon's optimizer state through universal checkpoints

[alanhuangyoo]

Fixes #8608.

Universal checkpoints could not carry Muon's optimizer state at any ZeRO stage. With this they can at all three, including a resume on a different number of ranks.

#8633 is in now, so this is only its own four commits.

What was wrong

1. The converter and the ZeRO-3 loader read Adam's exp_avg / exp_avg_sq by name, so a Muon checkpoint failed with KeyError: 'exp_avg'.
2. ZeRO-1/2 keep Muon's momentum whole per parameter, for every parameter a partition touches, because Newton-Schulz needs the whole matrix. That buffer is not shaped like the partition, and a matrix that straddles two partitions is in both ranks' buffers.
3. ZeRO-3's universal load copied the first param group's hyperparameters onto every group. With Muon that runs the Adam half as a Muon group, so every state came back exactly and the resume still diverged (3.2e-1 on the norm's bias after two steps).

The change

• The converter and both loaders take the state names from what the optimizer saved, per param group for ZeRO-1/2 and per sub-group for ZeRO-3.
• ZeRO-1/2 record which states hold each parameter whole (whole_param_optimizer_states in the optimizer checkpoint). The converter writes those per parameter from every rank that holds it and checks the copies agree. The loader rebuilds each rank's buffer in the order of its own partition, so a different dp size works.
• ZeRO-3 gives each param group back its own hyperparameters.
• A ZeRO-1/2 Muon checkpoint saved before the layout was recorded is refused. By size alone its buffer can look partition-shaped, and reading it that way writes the wrong rows without failing.
• momentum_buffer takes scale power -1 for affine-mapped (TP) pieces, like exp_avg.

Testing

2 and 4 H20s, fp32, the same batch on every rank, clipping off. Train 3 steps, save, convert, resume, take 2 more steps, and compare with the run that never stopped. Bit-identical in every case:

• ZeRO-1, 2 and 3 on the same ranks
• dp 2→1, 1→2, 2→4 and 4→2 for ZeRO-1/2; 2→4 and 4→1 for ZeRO-3
• ZeRO-1 ↔ ZeRO-2
• bf16: ZeRO-2 2→2 and 2→4, ZeRO-1 on 4 → ZeRO-2 on 2

Also checked offload now that #8464 is in. With offload, ZeRO-1/2 keep momentum_buffer per partition like Adam's moments, so it converts the same way. All optimizer states come back bit-identical after loading, for ZeRO-1/2/3 with dp 2→4, 4→2, 2→1 and 1→2, and also with offload switched on or off between save and resume. test_universal_optimizer_states.py covers offload now too.

The new tests/unit/checkpoint/test_universal_optimizer_states.py covers stages 1 to 3 at 2→2 and 2→4, plus the refusal. Undoing each fix makes its test fail: the ZeRO-3 group copy, the ZeRO-1/2 rebuild, and the staging fix from #8633, where the conversion reports the disagreeing copies.

Existing checkpoint tests (test_universal_checkpoint.py in fp32, test_zero_optimizer.py, test_muon_checkpoint.py, test_autotp_uc_checkpoint.py): 169 passed. The one failure, TestRealCheckpointUniversalConversionTPxPP, fails the same way on master here (a Gloo process-group attribute missing in this torch build). These need pytest < 8.4 as pinned in requirements-dev.txt; DistributedFixture does not resolve under pytest 9.

Not covered

• Converting between ZeRO-1/2 and ZeRO-3. The engine looks for stage-specific model-state files, so this fails for Adam on master too.
• AutoEP's converter still hard-codes Adam's state names.
• A universal checkpoint saved without round_robin_gradients and loaded with it diverges, for Adam on master as well. Fixed separately in Link ZeRO-1/2 fragments at their round-robin offsets #8637.