DeepSpeed

66fef47b - [AffineIR] Storage maps of AutoEP experts under EDP and ZeRO (#8941)

[AffineIR] Storage maps of AutoEP experts under EDP and ZeRO (#8941) ## 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)` (#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)` (#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 (#8809), fp32 master / optimizer state of experts, and expert tensor parallelism. - Replica selection policy (#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>