DeepSpeed

8a47f989 - [AffineIR] Plan expert transfers from the live runtime placement (#8800)

[AffineIR] Plan expert transfers from the live runtime placement (#8800) ## Summary #8623 plans shard-to-shard copies between two `ParamAffineMap`s. 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 (#8809), parameter swap, optimizer-state migration. - Storage maps of EDP replicas and ZeRO-3 fragments (follow-up PR), and replica selection policy (#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 #8623 (`plan_transfer`). - Plan execution is #8809; replica policy is #8252. - #8844 and #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. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>