[AffineIR] Storage maps of AutoEP experts under EDP and ZeRO

deepspeed/checkpoint/affine_ir_spec.md

@@ -613,6 +613,8 @@ expert parameters, and `AutoEPMoELayer.plan_expert_transfer(target_placement)` c
 them with a target descriptor through `plan_transfer` into shard-to-shard copies. Plan
 ranks are EP-local ranks of the respective placements; the plan is not executed, and ZeRO
 and EDP storage must still be normalized before its offsets apply.
+`AutoEPMoELayer.expert_storage_affine_maps()` describes that storage in global ranks
+(§10).
 
 `deepspeed/checkpoint/zero_affine.py` takes the first step for live training storage. It
 describes, in global ranks, where ZeRO keeps each element of a parameter: the low-precision
@@ -689,8 +691,12 @@ cannot be merged across those rows and keeps one piece per row. Callers:
 - `autotp_storage_affine_maps(model)` nests each parameter's published AutoTP layout map
   with the storage of every TP shard. A parameter AutoTP leaves whole in a model built under
   `zero.Init` stays partitioned over the whole world, so every TP rank reports the storage of
-  one shared copy; that storage describes the parameter on its own, without nesting.
+  one shared copy; that storage describes the parameter on its own, without nesting. The
+  experts of an AutoEP layer are whole on every TP rank and differ by EP rank; their maps
+  are the layer's `expert_storage_affine_maps()`.
 - AutoEP expert storage nests the placement's layout map with each EP rank's EDP storage.
+  It needs no communication: the placement, the EDP group table and the ZeRO-3 even split
+  are known to every rank.
 
 **10.4 Replica choice** stays with the planner (§2.4): every holder appears in `locations`,
 and the planner's default or a `replica_selector` picks one.
@@ -700,5 +706,17 @@ A caller builds them once and executes the plan for every sync, until either sid
 placement, parallel sizes or ZeRO partitioning changes. A parameter whose buffer is
 reallocated keeps its plan; only the shard views must be fetched again.
 
-**10.6 Not described.** Offload tiers, MiCS / hpZ, expert tensor parallelism, quantized
-ZeRO-3 storage, non-elementwise transforms.
+**10.6 Authoritative and derived copies.** A storage map lists only *authoritative* copies:
+storage that the optimizer step updates and the next forward reads, such as `param.data` or
+a ZeRO-3 `ds_tensor`. A write through the map therefore lasts. The runtime also keeps
+*derived* copies that it regenerates from authoritative ones: the gathered copy of a
+persistent ZeRO-3 parameter, ZeRO++ hpZ secondary partitions, fused inference weights.
+Writing a derived copy alone is lost or leaves the authoritative copy stale, so maps leave
+it out and shard accessors refuse it (`zero_param_storage_shard` on a gathered parameter).
+Whether a copy is authoritative depends on the side: the low-precision parameters are
+authoritative for a rollout engine, but on a training engine they are derived from the fp32
+master weights, which the next step casts back over them.
+
+**10.7 Not described.** Offload tiers, MiCS / hpZ, expert tensor parallelism, quantized
+ZeRO-3 storage, non-elementwise transforms, derived copies, and AutoEP expert optimizer
+states.

deepspeed/checkpoint/autoep_affine.py

@@ -2,13 +2,15 @@
 # DeepSpeed Team
 """AutoEP expert placement descriptors and their affine lowering.
 
-The descriptor records placement in EP-local rank coordinates. ZeRO and EDP
-fragments are deliberately outside this map: callers must first normalize
-storage to one logical, packed expert tensor per EP rank.
+The descriptor records placement in EP-local rank coordinates, and
+``autoep_placement_to_affine_map`` describes one logical, packed expert tensor
+per EP rank. ``autoep_storage_affine_map`` adds the expert-data-parallel and
+ZeRO layer on top: it describes what each global rank actually stores.
 """
 
-from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION, AffinePiece, ParamAffineMap
+from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION, AffinePiece, ParamAffineMap, nest_affine_maps
 from deepspeed.checkpoint.affine_transfer import plan_transfer
+from deepspeed.checkpoint.zero_affine import zero_storage_affine_map
 from deepspeed.checkpoint.constants import (AFFINE_MAP, AFFINE_MAP_PARAMS, AFFINE_MAP_VERSION, AUTOEP_AFFINE_MAPS,
                                             AUTOEP_EXPERT_PLACEMENT, AUTOEP_PLACEMENT_EP_SIZE,
                                             AUTOEP_PLACEMENT_EXPERTS, AUTOEP_PLACEMENT_NUM_EXPERTS,
@@ -23,6 +25,7 @@
     'legacy_uniform_autoep_placement_descriptor',
     'autoep_placements_equal',
     'autoep_placement_to_affine_map',
+    'autoep_storage_affine_map',
     'autoep_metadata_to_affine_map',
     'extract_autoep_rank_tensor',
     'plan_autoep_placement_transfer',
@@ -169,6 +172,40 @@ def autoep_placement_to_affine_map(descriptor, logical_shape):
     return affine_map
 
 
+def autoep_storage_affine_map(descriptor, logical_shape, holders_by_ep_rank, partition_numel_by_ep_rank=None):
+    """Lower one expert parameter to the storage each global rank actually holds.
+
+    ``holders_by_ep_rank`` maps every rank ID of ``descriptor`` to the global ranks that store that
+    EP rank's packed ``[local_experts, ...]`` tensor, in expert-data-parallel (EDP) rank order.
+
+    Without ``partition_numel_by_ep_rank``, every holder stores the whole packed tensor, as EDP
+    replicas do under ZeRO-0/1/2. With it, holder ``d`` of EP rank ``r`` stores elements
+    ``[d * n, (d + 1) * n)`` of the flattened packed tensor, where ``n`` is
+    ``partition_numel_by_ep_rank[r]``, as ZeRO-3 does. ZeRO-3's trailing padding holds no
+    parameter elements and is not part of the shard, so a holder whose partition is all padding,
+    like every holder of an EP rank without experts, has an empty shard.
+
+    The returned map describes the parameter in its own shape and is keyed by global rank; each
+    shard is the holder's flat storage. It nests the placement's layout map with ZeRO's
+    storage map of each packed tensor, so every holder of an element, across EDP replicas and
+    redundant experts, is listed in the piece's ``locations``; choosing among them is left to the
+    transfer planner's replica selection.
+    """
+    layout = autoep_placement_to_affine_map(descriptor, logical_shape)
+    entries_by_rank = {entry[AUTOEP_PLACEMENT_RANK]: entry for entry in descriptor[AUTOEP_PLACEMENT_RANKS]}
+    holders_by_ep_rank = _validated_storage_holders(holders_by_ep_rank, entries_by_rank)
+
+    storage_by_ep_rank = {}
+    for ep_rank, global_ranks in holders_by_ep_rank.items():
+        packed_numel = _product(layout.shard_shapes[ep_rank])
+        partition_numel = None
+        if partition_numel_by_ep_rank is not None:
+            partition_numel = _storage_partition_numel(partition_numel_by_ep_rank, ep_rank, len(global_ranks),
+                                                       packed_numel)
+        storage_by_ep_rank[ep_rank] = zero_storage_affine_map(packed_numel, global_ranks, partition_numel)
+    return nest_affine_maps(layout, storage_by_ep_rank)
+
+
 def autoep_metadata_to_affine_map(metadata, logical_shape, param_name=None):
     """Load a persisted AutoEP affine map, falling back to legacy placement lowering."""
     serialized_map = metadata.get(AFFINE_MAP)
@@ -252,6 +289,39 @@ def _pieces_for_rank(experts, holders, expert_shape, expert_numel, source_stride
     return pieces
 
 
+def _validated_storage_holders(holders_by_ep_rank, entries_by_rank):
+    if not isinstance(holders_by_ep_rank, dict):
+        raise TypeError('holders_by_ep_rank must map each EP rank ID to its global holder ranks.')
+    if set(holders_by_ep_rank) != set(entries_by_rank):
+        raise ValueError(f'holders_by_ep_rank names EP ranks {sorted(holders_by_ep_rank)}; the placement '
+                         f'descriptor names {sorted(entries_by_rank)}.')
+    seen = set()
+    validated = {}
+    for ep_rank in sorted(holders_by_ep_rank):
+        global_ranks = list(holders_by_ep_rank[ep_rank])
+        if not global_ranks:
+            raise ValueError(f'EP rank {ep_rank} has no global holder rank.')
+        for global_rank in global_ranks:
+            if not isinstance(global_rank, int) or isinstance(global_rank, bool) or global_rank < 0:
+                raise ValueError(f'Global holder rank {global_rank!r} of EP rank {ep_rank} is not a '
+                                 'non-negative integer.')
+            if global_rank in seen:
+                raise ValueError(f'Global rank {global_rank} holds storage for more than one EP rank entry.')
+            seen.add(global_rank)
+        validated[ep_rank] = global_ranks
+    return validated
+
+
+def _storage_partition_numel(partition_numel_by_ep_rank, ep_rank, num_holders, packed_numel):
+    partition_numel = partition_numel_by_ep_rank.get(ep_rank)
+    if not isinstance(partition_numel, int) or isinstance(partition_numel, bool) or partition_numel < 0:
+        raise ValueError(f'Partition size {partition_numel!r} of EP rank {ep_rank} is not a non-negative integer.')
+    if partition_numel * num_holders < packed_numel:
+        raise ValueError(f'EP rank {ep_rank} packs {packed_numel} elements, but its {num_holders} holders of '
+                         f'{partition_numel} elements each cannot store them.')
+    return partition_numel
+
+
 def _expert_holders(entries_by_rank, num_experts):
     holders = {expert_id: set() for expert_id in range(num_experts)}
     for rank, entry in entries_by_rank.items():

deepspeed/checkpoint/zero_affine.py

@@ -268,7 +268,8 @@ def autotp_storage_affine_maps(model, optimizer_states=False, tp_group=None):
     master copy and optimizer states instead, matching ``zero_optimizer_storage_shard``, and a
     parameter without them maps to ``None``. A parameter AutoTP did not partition is replicated
     over the tensor-parallel group, unless it kept the partition over the whole world that
-    ``zero.Init`` gave it, which describes it on its own.
+    ``zero.Init`` gave it, which describes it on its own. The experts of an AutoEP layer are
+    described by its ``expert_storage_affine_maps``; their optimizer states are not supported.
 
     Collective: the tensor-parallel ranks exchange their shards' storage maps in one all-gather
     over ``tp_group`` (the AutoTP group by default), and every rank must call this with the same
@@ -286,15 +287,47 @@ def autotp_storage_affine_maps(model, optimizer_states=False, tp_group=None):
         if marker is not None:
             marker()
 
-    named_params = list(model.named_parameters())
+    # AutoTP leaves AutoEP experts whole on every tensor-parallel rank, but each EP rank holds
+    # different experts, so a replicated layout would describe one EP rank's packed tensor as the
+    # whole parameter. The AutoEP layer knows its placement and describes their storage instead.
+    from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer
+    autoep_layers = [module for module in model.modules() if isinstance(module, AutoEPMoELayer)]
+    if optimizer_states and autoep_layers:
+        raise NotImplementedError('AutoEP expert optimizer states have no affine storage map.')
+
+    # Each EP rank holds different experts, so one rank may refuse its expert storage (for example
+    # a quantized partition) while its tensor-parallel peers accept theirs. The refusal travels with
+    # the all-gather so that every peer raises instead of waiting in it.
+    local_error = None
+    expert_maps = {}
+    try:
+        for module in autoep_layers:
+            for weight_name, storage_map in module.expert_storage_affine_maps().items():
+                expert_maps[id(getattr(module.experts, weight_name))] = storage_map
+    except (NotImplementedError, ValueError, RuntimeError) as error:
+        local_error = error
+
+    named_params = []
+    for name, param in model.named_parameters():
+        if id(param) not in expert_maps:
+            named_params.append((name, param))
     params = [param for _, param in named_params]
-    if optimizer_states:
-        storage_maps = zero_optimizer_storage_affine_maps(params)
+    if local_error is not None:
+        local = {'error': f'{type(local_error).__name__}: {local_error}'}
     else:
-        storage_maps = [zero_param_storage_affine_map(param) for param in params]
-    local = [None if storage_map is None else storage_map.to_dict() for storage_map in storage_maps]
+        if optimizer_states:
+            storage_maps = zero_optimizer_storage_affine_maps(params)
+        else:
+            storage_maps = [zero_param_storage_affine_map(param) for param in params]
+        local = [None if storage_map is None else storage_map.to_dict() for storage_map in storage_maps]
     gathered = [None] * tp_world_size
     dist.all_gather_object(gathered, local, group=tp_group)
+    if local_error is not None:
+        raise local_error
+    for tp_rank, rank_maps in enumerate(gathered):
+        if isinstance(rank_maps, dict):
+            raise RuntimeError(f"Tensor-parallel rank {tp_rank} cannot describe its AutoEP expert storage: "
+                               f"{rank_maps['error']}")
 
     maps = {}
     for position, (name, param) in enumerate(named_params):
@@ -314,4 +347,10 @@ def autotp_storage_affine_maps(model, optimizer_states=False, tp_group=None):
             layout = replicated_map(layout.logical_shape, 1)
             inner = {0: inner[0]}
         maps[name] = nest_affine_maps(layout, inner)
-    return maps
+    ordered = {}
+    for name, param in model.named_parameters():
+        if id(param) in expert_maps:
+            ordered[name] = expert_maps[id(param)]
+        else:
+            ordered[name] = maps[name]
+    return ordered

deepspeed/module_inject/auto_ep_layer.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import math
 from functools import lru_cache
 from typing import Literal, NamedTuple
 
@@ -22,7 +23,8 @@
 import deepspeed.comm as dist
 from deepspeed.accelerator import get_accelerator
 from deepspeed.checkpoint.autoep_affine import (autoep_experts_for_rank, autoep_placement_to_affine_map,
-                                                autoep_placements_equal, legacy_uniform_autoep_placement_descriptor,
+                                                autoep_placements_equal, autoep_storage_affine_map,
+                                                legacy_uniform_autoep_placement_descriptor,
                                                 plan_autoep_placement_transfer)
 from deepspeed.checkpoint.affine import ParamAffineMap
 from deepspeed.checkpoint.affine_transfer import PlanBudget, TransferPlan
@@ -792,11 +794,85 @@ def expert_affine_maps(self) -> dict[str, ParamAffineMap]:
         """Return the affine map of each live expert parameter (``w1``, ``w2``, ``w3``).
 
         Map ranks are EP-local ranks; translate them with ``dist.get_global_rank(self.ep_group, rank)``.
-        Each rank's shard is the unpartitioned packed ``[local_experts, ...]`` tensor.
+        Each rank's shard is the unpartitioned packed ``[local_experts, ...]`` tensor; use
+        :meth:`expert_storage_affine_maps` for the expert-data-parallel and ZeRO-3 storage.
         """
         placement = self._live_expert_placement()
         return {name: self._live_expert_affine_map(name, placement) for name in ("w1", "w2", "w3")}
 
+    def _expert_storage_group(self, params: dict[str, nn.Parameter]):
+        partitioned = [hasattr(param, "ds_id") for param in params.values()]
+        if all(partitioned):
+            storage_groups = {id(param.ds_process_group): param.ds_process_group for param in params.values()}
+            if len(storage_groups) != 1:
+                raise ValueError("AutoEP expert parameters are ZeRO-3 partitioned over different process groups.")
+            return next(iter(storage_groups.values()))
+        if any(partitioned):
+            raise ValueError("Only some AutoEP expert parameters are ZeRO-3 partitioned.")
+        if self.folding_group_handles is not None:
+            return self.folding_group_handles.edp_group
+        from deepspeed.utils import groups
+        return groups._get_expert_data_parallel_group(self.ep_group_name)
+
+    def expert_storage_affine_maps(self) -> dict[str, ParamAffineMap]:
+        """Return the affine map of the storage each global rank holds for ``w1``, ``w2`` and ``w3``.
+
+        This is computed locally from the expert data parallel group table and the placement, so it
+        involves no communication and any rank may call it on its own.
+
+        Map ranks are global ranks and the logical shape is the parameter's own shape. Each shard is
+        a rank's flat storage: the packed ``[local_experts, ...]`` tensor under ZeRO-0/1/2, where all
+        expert-data-parallel replicas are listed as holders, or the rank's ZeRO-3 partition
+        ``ds_tensor`` without its trailing padding. Optimizer state is not described. Quantized and
+        NVMe-offloaded ZeRO-3 storage is not supported.
+        """
+        if self.folding_group_handles is not None and self.folding_group_handles.spec.etp_size > 1:
+            raise NotImplementedError("Expert storage maps do not describe expert tensor-parallel shards yet.")
+        placement = self._live_expert_placement()
+        params = {name: getattr(self.experts, name) for name in ("w1", "w2", "w3")}
+        for name in params:
+            self._live_expert_affine_map(name, placement)
+        storage_group = self._expert_storage_group(params)
+        partitioned = hasattr(params["w1"], "ds_id")
+        if partitioned:
+            from deepspeed.checkpoint.zero_affine import _check_zero3_param_storage
+            for param in params.values():
+                _check_zero3_param_storage(param)
+
+        # The EDP group of each EP rank's global rank holds that EP rank's packed tensor. Every rank
+        # built all EDP groups, so the table answers this without assuming a rank layout.
+        from deepspeed.utils import groups
+        group_of_rank = {}
+        for group_ranks in groups._get_expert_data_parallel_group_table(self.ep_group_name):
+            for global_rank in group_ranks:
+                group_of_rank[global_rank] = list(group_ranks)
+        own_storage_ranks = dist.get_all_ranks_from_group(storage_group)
+        if group_of_rank.get(dist.get_rank()) != own_storage_ranks:
+            raise RuntimeError(f"Expert storage group ranks {own_storage_ranks} do not match the expert data "
+                               f"parallel group table for {self.ep_group_name}.")
+        holders_by_ep_rank = {}
+        for ep_rank, global_rank in enumerate(dist.get_all_ranks_from_group(self.ep_group)):
+            holders_by_ep_rank[ep_rank] = group_of_rank[global_rank]
+
+        storage_maps = {}
+        for name, param in params.items():
+            logical_shape = self._expert_logical_shape(param)
+            partition_numel_by_ep_rank = None
+            if partitioned:
+                # ZeRO-3 pads each packed tensor to a multiple of the group size and splits it evenly.
+                expert_numel = math.prod(logical_shape[1:])
+                partition_numel_by_ep_rank = {}
+                for ep_rank, holders in holders_by_ep_rank.items():
+                    packed_numel = len(autoep_experts_for_rank(placement, ep_rank)) * expert_numel
+                    partition_numel_by_ep_rank[ep_rank] = -(-packed_numel // len(holders))
+                expected = partition_numel_by_ep_rank[self.ep_rank]
+                if param.ds_tensor.ds_numel != expected:
+                    raise RuntimeError(f"Expert parameter {name} has a ZeRO-3 partition of {param.ds_tensor.ds_numel} "
+                                       f"elements, but even partitioning gives {expected}.")
+            storage_maps[name] = autoep_storage_affine_map(placement, logical_shape, holders_by_ep_rank,
+                                                           partition_numel_by_ep_rank)
+        return storage_maps
+
     def plan_expert_transfer(self,
                              target_placement: dict,
                              limits: PlanBudget | None = None,