[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: 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 (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.
• Docs: autoep.rst gains a section on expert_storage_affine_maps(); affine_ir_spec.md gets a pointer.

Not in scope

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

Tests

• tests/unit/v1/moe/test_autoep_checkpoint.py::TestAutoEPExpertStorageMaps (4 ranks, EP 2), with ZeRO-2, ZeRO-3, and ZeRO-3 + folded AutoTP 2:
  • On a real engine, extract(full, rank) equals each rank's live param.data / ds_tensor 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.
• 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.

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

[chatgpt-codex-connector]

Reject encoded or swapped-out ZeRO-3 expert storage

When frozen AutoEP experts use zero_quantized_nontrainable_weights, or ZeRO-Infinity has offloaded their partitions to NVMe, this branch still derives an ordinary elementwise map from param.ds_tensor.ds_numel. The quantized tensor is encoded and has a separate ds_quant_scale, while the NVMe tensor may only be a released placeholder, so the advertised offsets cannot be applied to the runtime storage and executing a plan can copy invalid weights. zero_param_storage_affine_map already rejects both cases in _check_zero3_param_storage; perform equivalent validation before returning these maps.

Useful? React with 👍 / 👎.

[jinyouzhi]

Good catch, fixed in b7a95bc. expert_storage_affine_maps() now runs the same _check_zero3_param_storage as zero_param_storage_affine_map on each ZeRO-3 expert parameter, so quantized and NVMe-offloaded storage raise NotImplementedError. Added test_zero3_storage_maps_refuse_storage_without_plain_elements covering both cases; they fail without the fix.