Follow the optimizer's own state names in universal checkpoints

[alanhuangyoo]

Fixes #8608.

The problem

Universal checkpoint conversion read Adam's state by name, so a checkpoint written by any other optimizer could not be converted. Muon keeps momentum_buffer for the matrices it orthogonalizes, so ds_to_universal raised KeyError: 'exp_avg' before writing anything — ZeRO-1/2 at extract_zero_shards, ZeRO-3 at extract_zero_shards_stage3.

The same three names are hard-coded in two more places: the merge loops in merge_tp_slices / merge_zero3_slices, and the ZeRO-3 loader, which also reads a step that the Muon half does not have. The ZeRO-1/2 loader was already generic — it takes the names from the files it finds — which is why it needed no change.

The change

• Extraction takes the names from the param group's own state.
• The merge takes them from the fragment files the extraction wrote.
• The ZeRO-3 loader takes them per sub-group, and the step from whichever group saved one.

Per group, not per checkpoint: DeepSpeed splits a Muon run into two param groups, so one checkpoint holds momentum_buffer for the matrices and exp_avg / exp_avg_sq / step for everything else. Reading the names once for the whole checkpoint sends the ZeRO-3 loader looking for 0.bias/momentum_buffer.pt.

_SCALE_POWERS gains momentum_buffer: -1. It is only consulted for a tensor-parallel piece, where a parameter scaled by s has a gradient scaled by 1/s, and the momentum buffer is gradient-like, as exp_avg is. A state name that is in no table stops the conversion with a message rather than being merged at an assumed scale.

What this does not cover

ZeRO-1/2 with Muon. There the momentum buffer is the size of the whole group and replicated on every rank, because Newton-Schulz needs the whole matrix (_muon_momentum_buffer, stage_1_and_2.py). Every other state is shaped like the partition, and the universal format stores state in the partition's layout, so slicing the buffer by partition offsets would write fragments that are the right size and hold the wrong rows. The conversion now refuses, naming the state:

Optimizer state 'momentum_buffer' of param group 0 has 256 elements, but the group's
ZeRO partition has 128. The universal format stores optimizer state in the partition's
layout, so a state of another shape - a buffer replicated across ranks, for instance -
cannot be placed yet.

A refusal is not the end state, but it is better than a resume that looks successful. I am happy to map the replicated buffer in a follow-up, or here if you would rather have it in one piece.

autoep_universal.py carries the same three names in its own paths and is untouched.

Testing

On 2×H20, from the repo root:

• New tests/unit/checkpoint/test_universal_optimizer_states.py. Neither test passes on master: the conversion raises KeyError: 'exp_avg' on the converting rank, which is also why the round trip hangs there rather than failing - the other rank is at the barrier behind it.
  • test_muon_state_survives_a_universal_round_trip — ZeRO-3, a model with a 2-D weight and a norm so both halves are present. After converting and resuming, every state is element-wise equal to the state before the save.
  • test_a_state_that_is_not_partition_shaped_stops_the_conversion — ZeRO-1, the replicated momentum buffer, expects the ValueError.
• Adam, ZeRO-1 and ZeRO-3, converted and resumed the same way: element-wise equal, so the generalization leaves the existing path where it was.

I could not run tests/unit/checkpoint/test_universal_checkpoint.py locally — its baseline_ws2 fixture fails to resolve in my environment, on master as well as here — so that suite is down to CI.