Repository navigation
[AffineIR] Plan expert transfers from the live runtime placement - #8800
Merged
Merged
Conversation
Expose the affine maps of an AutoEPMoELayer's live expert weights and plan their movement to a target placement through the affine transfer planner (deepspeedai#8623), without reading checkpoint metadata or materializing the logical parameter. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
jinyouzhi
marked this pull request as ready for review
October 12, 2026 10:48
jinyouzhi
requested review from
hwchen2017,
loadams,
tjruwase and
tohtana
as code owners
October 12, 2026 10:48
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 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".
| affine_maps[name] = autoep_placement_to_affine_map(placement, self._expert_logical_shape(param)) | ||
| return affine_maps | ||
|
|
||
| def plan_expert_transfer(self, |
There was a problem hiding this comment.
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 👍 / 👎.
Contributor
Author
There was a problem hiding this comment.
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.
Describe expert_affine_maps() and plan_expert_transfer(): how to build a target placement, the EP-local rank namespaces of a plan, the element offsets into the unpartitioned packed tensor under ZeRO-3/EDP, replica selection, plan budgets, and the uniform live placement requirement. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
# Conflicts: # docs/code-docs/source/autoep.rst
jinyouzhi
changed the title
jinyouzhi
marked this pull request as draft
October 16, 2026 10:32
jinyouzhi
marked this pull request as ready for review
October 16, 2026 13:19
This was referenced Oct 16, 2026
…partitions Add autoep_storage_affine_map and AutoEPMoELayer.expert_storage_affine_maps so the transfer planner can read directly from runtime storage: maps are keyed by global rank over the flattened parameter, every EDP replica is offered as a holder for replica selection, ZeRO-3 partitions exclude padding, and ranks without experts or with padding-only partitions get empty shards. Compare placements by explicit rank IDs instead of list order, and verify that the published per-parameter affine map matches the live placement. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
…oep-live-affine-map
Every rank already creates all expert data parallel groups, so record their rank tables and derive each EP rank's holders locally. ZeRO-3 partition sizes follow from the placement (ceil(packed_numel / group size)) and are checked against the local ds_tensor. Clarify that plan_expert_transfer does not see EDP replicas. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
jinyouzhi
changed the title
delock
self-requested a review
October 21, 2026 00:28
The table lists every EDP subgroup of a group name, but only a subgroup's members run the release path, so dropping it there left ranks disagreeing about the topology. Clear it only on world teardown; re-creating the group name overwrites it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
A ZeRO-3 fragment that splits rows is expressible as several N-D boxes; the flat shape keeps it one interval. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Storage maps compose the expert placement with ZeRO's storage layout, which belongs with the ZeRO storage maps rather than with live placement planning. Remove expert_storage_affine_maps, autoep_storage_affine_map, the expert data parallel group table and their tests and docs; they return on top of a generic map nesting in a follow-up. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
delock
reviewed
Oct 23, 2026
Expert rebalancing is not implemented in AutoEP; describe the target placement instead. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
delock
approved these changes
Oct 23, 2026
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Oct 30, 2026
…edai#8941) ## Summary deepspeedai#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 deepspeedai#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 deepspeedai#8800's placement map with deepspeedai#8885's ZeRO storage maps, so `plan_transfer` (deepspeedai#8623) reads directly from runtime storage: ``` AutoEP placement map (deepspeedai#8800) --nest--> ZeRO storage map per EP rank (deepspeedai#8885) -> storage map in global ranks -> plan_transfer (deepspeedai#8623) ``` Related: deepspeedai#8252, deepspeedai#8800, deepspeedai#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)` (deepspeedai#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)` (deepspeedai#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 (deepspeedai#8809), fp32 master / optimizer state of experts, and expert tensor parallelism. - Replica selection policy (deepspeedai#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. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 liveAutoEPMoELayeractually holds, and builds the target map from a target placement descriptor. Both maps then go straight toplan_transfer.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 livew1/w2/w3.dist.get_global_rank(self.ep_group, rank).[local_experts, ...]tensor. ZeRO-3/EDP fragments stay outside the map, and the shape comes fromds_shape.AutoEPMoELayer.plan_expert_transfer(target_placement, limits=None, *, replica_selector=None)returns oneTransferPlanper expert parameter.autoep_affine.plan_autoep_placement_transfer(source_placement, target_placement, logical_shape, ...)is the layer-independent form. Both placements go through the existingautoep_placement_to_affine_map, so no second placement→map conversion is introduced.autoep_placements_equal), and each parameter's publishedDS_AUTOEP_UC_METAplacement and affine map must equal the live lowering.repack_expert_weightsactually packed the weights; today that is contiguous and uniform. A descriptor that disagrees raisesNotImplementedError, rather than producing a map that addresses shards the weights do not have.Not in scope
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 packedw1/w2/w3. The reference is the full tensor fromrepack_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 callsrebuild/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.
plan_transfer).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.
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.
AI assistance. This PR was developed with AI assistance (GitHub Copilot CLI). I reviewed every changed line and ran the tests above.