Restore a shard from the map the layer published

[Achyuthan-S]

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

[Copilot]

Copilot review overview

🟡 Changes recommended

Optimizer moments currently use the default scaling power, which can rescale mapped states incorrectly.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1

Open (1)

• Use correct affine scale powers when restoring optimizer states · New

What changed in this PR

Updates AutoTP universal-checkpoint restore to use the target layer’s affine shard map, with legacy metadata fallback.

Changes:

• Restores mapped parameters through affine shard extraction.
• Adds tests for map precedence and fallback behavior.

File	Summary
tests/​unit/​runtime/​tensor_parallel/​test_autotp_universal_checkpoint.py	Tests map precedence and legacy fallback.
deepspeed/​checkpoint/​universal_checkpoint.py	Implements affine-map-based shard restoration.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Achyuthan-S]

Fixed by threading the state key through: restore resolves a partition once per state file,
so each one now takes its own power of a piece's scale rather than the parameter's.

The power table moves to 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 instead of two
that could drift — the same reasoning as shared_qk_value_head_ids in #8575.

Worth noting what threading it also buys: _reject_scaled_optimizer_state only triggers on
a non-parameter power, so the refusal built for scaled optimizer states needs the real power
to reach it. With the state threaded, a scaled moment now raises rather than being rescaled.

No layer emits a scaled map today — the only scaled case, Yuan's o_proj bias, is still
unsupported — so this is ahead of the path rather than behind it. It becomes live when
rollout needs scale, which is the mid-term goal in #8230.

Test asserts the parameter restores through its own power and both moments raise; reverting
the threading fails it.

[Achyuthan-S]

@delock #8622 is rebased onto master and conflict-free — it had picked up a conflict with #8544's import block, which is resolved. This is the last step-4 piece: restore reads the map, so convert and restore are the same piece list run in both directions.
112 tests pass on the rebased tree (the affine suite, the resume matrix, the coverage test, tensor_parallel, and #8544's new AutoEP suite). Could you take a look when you have a moment?

[delock]

Hi @Achyuthan-S the formatting error needs to be fixed, thanks!

[Achyuthan-S]

@delock Fixed and rebased. Adding SCALE_POWER_BY_STATE took the export list to 131 characters, and yapf 0.43.0 rewraps it where the version I had locally left it alone — which is why it passed here and failed in CI. I'm on the pinned version now.