[CKPT] Plan expert transfers from the live runtime placement

[jinyouzhi]

Summary

#8623 plans shard-to-shard copies between two ParamAffineMaps. Today, though, the only AutoEP maps come from checkpoint metadata, so nothing at runtime can feed the planner. This PR closes that gap: it builds the source map from the placement the live AutoEPMoELayer actually holds, and builds the target map from a target placement descriptor. Both maps then go straight to plan_transfer.

AutoEP runtime placement -> source / target ParamAffineMap -> plan_transfer (#8623) -> copy descriptors

Related: #8252 (§8.4 AutoEP expert placement), #8230, #8544.

Changes

• AutoEPMoELayer.expert_affine_maps() returns the affine maps of the live w1/w2/w3.
  • Map ranks are EP-local ranks. The docstring explains how to translate them with dist.get_global_rank(self.ep_group, rank).
  • Each shard is the unpartitioned packed [local_experts, ...] tensor. ZeRO-3/EDP fragments stay outside the map, and the shape comes from ds_shape.
  • No collectives.
• AutoEPMoELayer.plan_expert_transfer(target_placement, limits=None, *, replica_selector=None) returns one TransferPlan per expert parameter.
  • It only plans: it moves no data, swaps no parameters and migrates no optimizer state.
• autoep_affine.plan_autoep_placement_transfer(source_placement, target_placement, logical_shape, ...) is the layer-independent form. Both placements go through the existing autoep_placement_to_affine_map, so no second placement→map conversion is introduced.
• The live placement must match how repack_expert_weights actually packed the weights; today that is contiguous and uniform. A descriptor that disagrees raises NotImplementedError, rather than producing a map that addresses shards the weights do not have.
  • A placement the descriptor can express is not necessarily one the runtime can execute. The target placement is only used for planning.

Not in scope

• Executing plans or transporting data (see the executor work in [UC] Execute a shard-to-shard transfer within a staging budget 0z5a/DeepSpeed#1), parameter swap, optimizer-state migration.
• Placement generation/versioning: the runtime does not change placement after construction yet.
• Runtime support for non-uniform placements.

Tests

tests/unit/v1/moe/test_autoep_unit.py adds three tests:

• test_live_expert_maps_address_the_packed_runtime_weights: for each rank, the map's shard shape and contents match the real packed w1/w2/w3. The reference is the full tensor from repack_expert_weights(ep_size=1), which does not depend on any affine map.
• test_plan_expert_transfer_moves_live_shards_into_target_packed_order:
  • The source is EP=2. The target has 3 ranks with reordered, uneven, empty and replicated placements.
  • Executing the plan on the live shards reproduces the target shards in target packed order, and every target element is written exactly once.
  • Planning is checked to never call rebuild/extract.
• test_live_expert_maps_reject_a_placement_the_weights_do_not_have: a mismatched descriptor and a mismatched logical shape both raise.

Results:

• test_autoep_unit.py, test_autoep_affine.py, test_autoep_affine_universal.py, test_affine_transfer_plan.py: 354 passed, 6 skipped.
• Swapping in a wrong live placement makes the new tests fail.
• Hardware: Intel Arc Pro B70; these are single-process unit tests.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-16T13:24:09.021953Z	1b8298f	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9637a10dbf

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add user documentation for the transfer-planning API

For downstream code that consumes these new public planning methods, the dedicated AutoEP documentation and affine IR specification remain unchanged, so users are not given a discoverable contract for constructing target_placement, translating the separate source/target EP-local rank namespaces, or normalizing ZeRO/EDP storage before applying segment offsets. Add this API and its placement/normalization requirements to docs/code-docs/source/autoep.rst or the affine specification rather than relying only on implementation docstrings.

AGENTS.md reference: AGENTS.md:L26-L26

Useful? React with 👍 / 👎.

[jinyouzhi]

Added in c02604d: a "Planning expert placement transfers" section in docs/code-docs/source/autoep.rst covering how to build target_placement with make_autoep_placement_descriptor, the source/target EP-local rank namespaces (and dist.get_global_rank(layer.ep_group, rank)), element offsets into the unpartitioned packed [local_experts, ...] tensor (gather under ZeRO-3; EDP replicas hold the same tensor), replica selection, PlanBudget limits, and the uniform live-placement requirement. affine_ir_spec.md §8.4 now points to the runtime API.