[jinyouzhi]

Summary

#8800 describes AutoEP expert weights as one packed [local_experts, ...] tensor per EP rank. What a process actually stores is different:

• under ZeRO-0/1/2, every expert-data-parallel (EDP) replica holds that tensor;
• under ZeRO-3, each EDP rank holds a padded fragment of it.

So a plan built from #8800's maps cannot read from training storage until the storage is first normalized to packed EP tensors.

This PR describes the storage itself. It nests #8800's placement map with #8885's ZeRO storage maps, so plan_transfer (#8623) reads directly from runtime storage:

AutoEP placement map (#8800) --nest--> ZeRO storage map per EP rank (#8885) -> storage map in global ranks -> plan_transfer (#8623)

Related: #8252, #8800, #8885.

Changes

• AutoEPMoELayer.expert_storage_affine_maps() returns one storage map each for w1, w2 and w3.
  • Ranks are global ranks; the logical shape is the parameter's own shape.
  • Each shard is the rank's flat storage, as zero_param_storage_shard(param) ([AffineIR] Affine storage maps for live ZeRO parameters and optimizer states #8885) returns it: param.data under ZeRO-0/1/2, or param.ds_tensor without its trailing padding under ZeRO-3.
  • Under ZeRO-0/1/2, all EDP replicas are listed as holders, so the planner's replica_selector can choose among them.
  • Under ZeRO-3, a rank whose partition is all padding, or whose EP rank holds no experts, has an empty shard.
  • Fails loudly instead of guessing:
    • mixed ZeRO-3 and non-ZeRO-3 expert params;
    • a storage group that does not match the EDP table;
    • a ds_tensor size that differs from even partitioning;
    • expert tensor parallelism, quantized ZeRO-3 storage and NVMe-offloaded ZeRO-3 storage (NotImplementedError).
• autoep_affine.autoep_storage_affine_map(descriptor, logical_shape, holders_by_ep_rank, partition_numel_by_ep_rank=None) is the layer-independent builder. It is also convenient for target maps.
• groups keeps every EDP group's ranks (_get_expert_data_parallel_group_table), not only the caller's.
  • Every rank already creates all groups, so the maps are built locally, without a collective, and any rank may call it alone.
  • The table is cleared together with the other group state.
• autotp_storage_affine_maps(model) ([AffineIR] Affine storage maps for live ZeRO parameters and optimizer states #8885) takes the maps of AutoEP experts from their layer. AutoTP leaves experts whole on every TP rank, so before this it described them as replicated over TP: one EP rank's packed tensor described as the whole parameter. With optimizer_states=True it raises NotImplementedError for AutoEP experts before any collective.
• Docs: autoep.rst gains a section on expert_storage_affine_maps(); affine_ir_spec.md §10 is updated (below).

Authoritative and derived copies

Storage maps list only authoritative copies: the storage the optimizer step updates and the next forward reads (param.data, ZeRO-3 ds_tensor). A write through a map therefore lasts.

The runtime also keeps derived copies that it regenerates from authoritative ones:

Derived copy	Regenerated by
gathered ZeRO-3 persistent parameter	re-gather after each step
ZeRO++ hpZ secondary partition	forward; invalidated after the step
fused inference weights (hybrid engine)	gather + fuse for generation
low-precision params on a training target	cast from the fp32 master weights at every step

A write to a derived copy alone is lost or leaves the authoritative copy stale. So maps leave derived copies out, and zero_param_storage_shard refuses a gathered parameter (call engine.empty_partition_cache() first). Whether a copy is authoritative depends on the side: the low-precision weights are authoritative for a rollout engine, which is what these maps serve, but derived on a training engine, where loading weights needs the fp32 master (not described).

Not in scope

• Executing plans ([checkpoint] Execute affine transfer plans with bounded peer staging #8809), fp32 master / optimizer state of experts, and expert tensor parallelism.
• Replica selection policy (Define affine IR for checkpointing #8252).

Follow-ups

1. Placement changes (EPLB). The maps are computed locally because the placement, the EDP group table and the ZeRO-3 even split are known to every rank. Once the runtime can change placement at run time (expert load balancing, redundant experts), this stays correct only if every rank uses the same placement: maps and plans must be rebuilt after every rank has switched, ideally checked with a placement version. The descriptor already expresses uneven and redundant placements (redundant experts become extra holders); today the runtime packs only the contiguous, uniform placement, and anything else raises NotImplementedError.
2. Expert fp32 master weights and optimizer states. Needed when a training engine is the target or for training-side resharding. Unlike the weights, this needs an all-gather over the EDP group and an exchange over EP, then nesting as here.
3. Derived copies as read-only sources. A fresh derived copy (e.g. a gathered ZeRO-3 persistent parameter after the step) holds the same values as the authoritative one and could be offered to the planner as an extra source, never as a target. Its validity is time-dependent, so it would be an explicit opt-in that does not affect cached plans.

Tests

• tests/unit/v1/moe/test_autoep_checkpoint.py::TestAutoEPExpertStorageMaps (4 ranks, EP 2), with ZeRO-2, ZeRO-3, and ZeRO-2/3 + folded AutoTP 2:
  • On a real engine, extract(full, rank) equals each rank's zero_param_storage_shard(param) bit for bit. full comes from GatheredParameters and the runtime's packing order, so the reference does not depend on the affine maps.
  • Rank 0 builds the maps alone while the others wait at a barrier: a hidden collective would deadlock. The result equals what every rank builds.
  • Planning from the storage maps into a different 4-rank placement [[3], [0], [2], [1]] and executing it into NaN-filled targets reproduces the expected experts exactly.
  • With AutoTP, autotp_storage_affine_maps(model) gives the experts the same maps as their layer (fails without the fix), and refuses expert optimizer states.
  • Quantized and NVMe-offloaded ZeRO-3 expert storage is refused.
• tests/unit/checkpoint/test_autoep_affine.py:
  • storage maps for replicated and ZeRO-3 layouts;
  • every EDP replica offered to the selector;
  • direct ZeRO-3 → new placement plans;
  • explicit rank IDs;
  • invalid holders refused.
• tests/unit/utils/test_groups.py: the EDP table survives subgroup destruction.

Validation on 8× Intel Arc Pro B70 (torch 2.14.0.dev20260804+xpu): test_autoep_affine, test_autoep_affine_universal, test_zero_affine, test_affine_transfer_plan, test_affine_shard_map, test_autoep_zero3_metadata, tests/unit/utils/test_groups, test_autoep_checkpoint, test_autoep_unit: 531 passed, 8 skipped.

After the AutoTP fix: TestAutoEPExpertStorageMaps 6 passed, test_zero_affine 27 passed.

[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-29T07:44:13.255490Z	18423bf	PR opened

ℹ️ 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: 18423bfa1f

ℹ️ 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".

[jinyouzhi]

Ready to review @delock Thanks!