[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

It also describes what each process actually stores under expert data parallelism (EDP) and ZeRO, so a plan can read directly from runtime storage instead of first normalizing it to one packed tensor per EP rank.

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.
• AutoEPMoELayer.expert_storage_affine_maps() returns maps of the storage each global rank holds, built by autoep_affine.autoep_storage_affine_map(descriptor, logical_shape, holders_by_ep_rank, partition_numel_by_ep_rank=None).
  • Map ranks are global ranks. The logical shape is the flattened parameter (numel,), so each shard is a rank's flat storage, and a ZeRO-3 fragment that splits an expert row stays one interval.
  • ZeRO-0/1/2: every EDP replica of an EP rank, including folded AutoTP peers, is listed in locations. The map makes no choice among them; the planner's default or replica_selector does (replica policy stays with Define affine IR for checkpointing #8252).
  • ZeRO-3: holder d stores [d*n, (d+1)*n) of the packed tensor with trailing padding excluded, so a padding-only partition, or any holder of an EP rank without experts, has an empty shard and is never a source or target.
  • Replicas are cut at every holder's fragment boundaries, so they expose identical pieces as the planner requires.
  • Collective over the default group: one all_gather_object gives each rank the holders of every EP rank, without assuming a rank layout. Expert tensor parallelism (etp_size > 1) raises NotImplementedError; optimizer state is not described.
• Placements are compared by explicit rank IDs rather than list order (autoep_placements_equal), and each parameter's published DS_AUTOEP_UC_META placement and affine map must equal the live lowering.
• 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 ([checkpoint] Execute affine transfer plans with bounded peer staging #8809), parameter swap, optimizer-state migration.
• Replica selection policy (Define affine IR for checkpointing #8252); this PR only exposes the candidates.
• 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:

• 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: from EP=2 into 3 ranks with reordered, uneven, empty and replicated placements; executing the plan reproduces the target shards, every element is written once, and planning never calls rebuild/extract.
• test_live_expert_maps_reject_a_placement_the_weights_do_not_have, test_live_expert_maps_identify_placement_ranks_by_id, test_live_expert_maps_reject_a_published_map_the_placement_does_not_produce.

tests/unit/checkpoint/test_autoep_affine.py (storage maps): each global rank's shard equals storage built by indexing the full tensor; every EDP replica is offered as a candidate and a selector can pick another; a ZeRO-3 source with padding-only and empty ranks plans directly into a new placement; explicit rank IDs; invalid holders/partitions are rejected.

tests/unit/v1/moe/test_autoep_checkpoint.py::TestAutoEPExpertStorageMaps (4 processes, EP=2/EDP=2, ZeRO-2, ZeRO-3, ZeRO-3 + folded AutoTP=2): on real engines, extract of the gathered packed tensors equals each rank's param.data or ds_tensor, and a plan from the live storage into a 4-rank one-expert-per-rank placement reproduces the target exactly.

Results (Intel Arc Pro B70):

• test_autoep_unit.py, test_autoep_affine.py, test_autoep_affine_universal.py, test_affine_transfer_plan.py, test_affine_shard_map.py, test_autoep_zero3_metadata.py: 443 passed, 6 skipped.
• TestAutoEPExpertStorageMaps: 3 passed on 4× B70.
• Mutations (single-holder locations, cutting at only the first replica's boundaries, keeping padding in shards, dropping ZeRO-3 partition sizes in the layer) each make at least one new test fail.

Contribution checklist

Related work.

• This PR builds on [UC] Plan a scale-1 affine shard-to-shard transfer without rebuilding the parameter #8623 (plan_transfer).
• Plan execution is [checkpoint] Execute affine transfer plans with bounded peer staging #8809; replica policy is Define affine IR for checkpointing #8252.
• Universal Checkpointing: read AutoEP ZeRO-3 optimizer files one at a time #8844 and AutoEP + AutoTP folding: Universal Checkpointing under ZeRO-3 #8848 change Universal Checkpointing for AutoEP. They do not produce runtime maps.

A search for AutoEP placement and transfer and for affine maps found no overlapping PR.

Commands and results. Hardware and software: Intel Arc Pro B70 (XPU, XCCL), Level Zero driver 1.15.39122 (compute runtime 26.27.39122.11), torch 2.14.0.dev20260804+xpu.

pytest -p no:randomly tests/unit/v1/moe/test_autoep_unit.py tests/unit/checkpoint/test_autoep_affine.py \
  tests/unit/checkpoint/test_autoep_affine_universal.py tests/unit/checkpoint/test_affine_transfer_plan.py \
  tests/unit/checkpoint/test_affine_shard_map.py tests/unit/checkpoint/test_autoep_zero3_metadata.py
pytest -p no:randomly tests/unit/v1/moe/test_autoep_checkpoint.py::TestAutoEPExpertStorageMaps   # 4 ranks

Results are listed under Tests above.

Numerics, performance, memory. Training is not affected. The new methods only run when called. expert_storage_affine_maps() issues one all_gather_object; the other methods use no collectives.

Compatibility matrix.

	Validated	Not validated
ZeRO	stage 2, stage 3, stage 3 + folded AutoTP=2 on real engines (EP=2, EDP=2); stage 0/1 replicas in unit tests	Real engines with stage 0/1; offload
Parallelism	EP=2 with EDP=2; target placements on 3–4 ranks	etp_size > 1 (raises NotImplementedError)
Precision	Dtype-independent: maps describe shapes and offsets
Accelerator	XPU	CUDA, HPU

AI assistance. This PR was developed with AI assistance (GitHub Copilot CLI). I reviewed every changed line and ran the tests above.

[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.