Restore a shard from the map the layer published (#8622)
Restore now takes this rank's shard from the map the layer published,
which is the same piece list conversion uses run in the opposite
direction. That completes the symmetry §7 of `affine_ir_spec.md`
describes and that conversion has had since #8385: until now the two
sides described the same layout in two different ways, and could
disagree.
Follows #8519 (layers emit the map) and #8575 (GPTBigCode and Yuan
describe themselves).
### This one is not additive
The map lives on the parameter, built by this job's layer. It is `M_t` —
it describes the topology being restored into, regardless of how the
source checkpoint was converted. So the new path is taken for **every**
AutoTP restore, not only for checkpoints that were converted through a
map. A parameter whose layer published no map falls back to the existing
category keys unchanged.
That is the correct reading of the spec rather than a shortcut: §6.1
keeps `M_t` out of the file precisely because the restoring job derives
it from its own layers. But it does mean the blast radius here is wider
than the three PRs before it, so the equivalence is worth checking
rather than assuming.
I ran `tests/unit/checkpoint/` and `tests/unit/runtime/tensor_parallel/`
(385 tests) against this branch and against upstream master in a
separate worktree. The failing set is **identical by name**, not merely
the same count — 75 failed, 187 passed either way, all of them FusedAdam
JIT, fp16-on-CPU, or the pre-existing `st_st_size` typo in
`test_convert_checkpoint.py`.
### The tests are built to discriminate
`test_restore_prefers_the_map_over_the_category_keys` gives one
parameter deliberately contradictory geometry: the category keys
describe a plain dim-0 split, so rank 0 would take the first half, while
the map hands rank 0 the second half. Which half rank 0 receives says
which side of the metadata restore actually consulted.
This matters because the old path covers the same layouts correctly.
Without a discriminating case the map branch could be deleted and every
existing test would stay green. I checked: removing the branch fails
that test, while `test_restore_falls_back_when_no_map_was_published`
keeps passing.
### Validation
112 passed on CPU/gloo (`DS_ACCELERATOR=cpu LOCAL_SIZE=4`): the affine
suite, the resume matrix, the coverage test,
`tests/unit/runtime/tensor_parallel/`, and #8544's AutoEP affine suite,
which this now sits on top of after the rebase.
### Scale powers
Restore resolves a partition once per state file, so each state takes
its own power of a piece's scale. The power table lives in affine.py as
SCALE_POWER_BY_STATE — ds_to_universal.py had its own copy since #8519,
so conversion and restore now read one definition rather than two that
could drift.
This also reaches the refusal built for scaled optimizer states, which
only triggers on a non-parameter power. No layer emits a scaled map
today, so it is ahead of the path rather than behind it.
Related: #8252, #8230.
cc @delock
---------
Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
