[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

Maps of what each process actually stores under expert data parallelism (EDP) and ZeRO are left to a follow-up, which builds them by nesting these layout maps with ZeRO's storage maps.

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.
• 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.
• Storage maps of EDP replicas and ZeRO-3 fragments (follow-up PR), and replica selection policy (Define affine IR for checkpointing #8252).
• 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: test_placements_equal_identifies_ranks_by_id.

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: 432 passed, 6 skipped.
• test_autoep_checkpoint.py (8× B70): 32 passed, 2 skipped.

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

Results are listed under Tests above.

Numerics, performance, memory. Training is not affected. The new methods only run when called and use no collectives.

Compatibility matrix.

	Validated	Not validated
ZeRO	Independent: maps describe the unpartitioned packed tensors	Storage under EDP/ZeRO (follow-up)
Parallelism	EP=2; target placements on 3 ranks
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.