[AffineIR] Plan expert transfers from the live runtime placement

deepspeed/checkpoint/affine_ir_spec.md

@@ -608,6 +608,12 @@ subset of a parameter combined with a list of ranks holding this subset" — is
 §2.1's exact `locations` implements, and the same mechanism can later describe normalized
 ZeRO/offload placement without introducing another geometry IR.
 
+At runtime, `AutoEPMoELayer.expert_affine_maps()` returns the maps of the live packed
+expert parameters, and `AutoEPMoELayer.plan_expert_transfer(target_placement)` composes
+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.
+
 ---
 
 ## 9. Staging

deepspeed/checkpoint/autoep_affine.py

@@ -8,6 +8,7 @@
 """
 
 from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION, AffinePiece, ParamAffineMap
+from deepspeed.checkpoint.affine_transfer import plan_transfer
 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,
@@ -20,9 +21,11 @@
     'validate_autoep_placement_descriptor',
     'autoep_experts_for_rank',
     'legacy_uniform_autoep_placement_descriptor',
+    'autoep_placements_equal',
     'autoep_placement_to_affine_map',
     'autoep_metadata_to_affine_map',
     'extract_autoep_rank_tensor',
+    'plan_autoep_placement_transfer',
 ]
 
 
@@ -114,6 +117,24 @@ def legacy_uniform_autoep_placement_descriptor(num_experts, num_local_experts, e
     return make_autoep_placement_descriptor(num_experts, experts_by_rank)
 
 
+def autoep_placements_equal(descriptor, other):
+    """Return whether two descriptors place the same experts, in the same order, on each rank ID.
+
+    Rank entries are identified by their explicit rank IDs, so listing them in another order does
+    not change the placement.
+    """
+    validate_autoep_placement_descriptor(descriptor)
+    validate_autoep_placement_descriptor(other)
+    if descriptor[AUTOEP_PLACEMENT_NUM_EXPERTS] != other[AUTOEP_PLACEMENT_NUM_EXPERTS]:
+        return False
+    if descriptor[AUTOEP_PLACEMENT_EP_SIZE] != other[AUTOEP_PLACEMENT_EP_SIZE]:
+        return False
+    for rank in range(descriptor[AUTOEP_PLACEMENT_EP_SIZE]):
+        if autoep_experts_for_rank(descriptor, rank) != autoep_experts_for_rank(other, rank):
+            return False
+    return True
+
+
 def autoep_placement_to_affine_map(descriptor, logical_shape):
     """Lower one ``[num_experts, ...]`` parameter to a :class:`ParamAffineMap`."""
     validate_autoep_placement_descriptor(descriptor)
@@ -188,6 +209,24 @@ def extract_autoep_rank_tensor(full_param, target_map, ep_rank):
     return target_map.extract(full_param, ep_rank)
 
 
+def plan_autoep_placement_transfer(source_placement,
+                                   target_placement,
+                                   logical_shape,
+                                   limits=None,
+                                   *,
+                                   replica_selector=None):
+    """Plan copies that move one expert parameter from ``source_placement`` to ``target_placement``.
+
+    Both placements are lowered with :func:`autoep_placement_to_affine_map`, so plan ranks are
+    EP-local ranks of the respective placements, and offsets follow each rank's local packed order.
+    The plan only describes the copies; it does not require that the AutoEP runtime can execute
+    with ``target_placement``.
+    """
+    source_map = autoep_placement_to_affine_map(source_placement, logical_shape)
+    target_map = autoep_placement_to_affine_map(target_placement, logical_shape)
+    return plan_transfer(target_map, source_map, limits, replica_selector=replica_selector)
+
+
 def _pieces_for_rank(experts, holders, expert_shape, expert_numel, source_strides, dest_strides):
     pieces = []
     run_start = 0

deepspeed/module_inject/auto_ep_layer.py

@@ -22,7 +22,10 @@
 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,
-                                                legacy_uniform_autoep_placement_descriptor)
+                                                autoep_placements_equal, legacy_uniform_autoep_placement_descriptor,
+                                                plan_autoep_placement_transfer)
+from deepspeed.checkpoint.affine import ParamAffineMap
+from deepspeed.checkpoint.affine_transfer import PlanBudget, TransferPlan
 from deepspeed.checkpoint.constants import (AFFINE_MAP, AUTOEP_EXPERT_PLACEMENT, AUTOEP_PARAM_EP_RANK,
                                             AUTOEP_PARAM_LOCAL_EXPERTS, AUTOEP_PARAM_LOGICAL_SHAPE, DS_AUTOEP_UC_META)
 from deepspeed.module_inject.auto_ep_config import AutoEPConfig, MoELayerSpec, resolve_autoep_config_defaults
@@ -648,8 +651,7 @@ def __init__(
         self.experts.w3.requires_grad_(w3_requires_grad)
         local_experts = autoep_experts_for_rank(self.expert_placement_descriptor, self.ep_rank)
         for param in (self.experts.w1, self.experts.w2, self.experts.w3):
-            physical_shape = getattr(param, 'ds_shape', param.shape)
-            logical_shape = [self.num_experts, *physical_shape[1:]]
+            logical_shape = self._expert_logical_shape(param)
             affine_map = autoep_placement_to_affine_map(self.expert_placement_descriptor, logical_shape)
             setattr(
                 param,
@@ -757,6 +759,66 @@ def _wait_async_split_plan(self, pending: _PendingSplitPlan) -> SplitPlan:
         self._async_split_plan_pending = None
         return plan
 
+    def _expert_logical_shape(self, param: nn.Parameter) -> list[int]:
+        # ZeRO-3 partitions the packed local tensor in place; ds_shape keeps its unpartitioned shape.
+        physical_shape = getattr(param, 'ds_shape', param.shape)
+        return [self.num_experts, *physical_shape[1:]]
+
+    def _live_expert_placement(self) -> dict:
+        # repack_expert_weights packs each rank's experts as one contiguous, uniform range. A
+        # descriptor that says otherwise would address shards the weights do not have.
+        packed_placement = legacy_uniform_autoep_placement_descriptor(self.num_experts, self.num_local_experts,
+                                                                      self.ep_size)
+        if not autoep_placements_equal(self.expert_placement_descriptor, packed_placement):
+            raise NotImplementedError("AutoEP expert weights are packed with the contiguous uniform placement; "
+                                      "the layer's expert placement descriptor does not describe them.")
+        return self.expert_placement_descriptor
+
+    def _live_expert_affine_map(self, name: str, placement: dict) -> ParamAffineMap:
+        param = getattr(self.experts, name)
+        affine_map = autoep_placement_to_affine_map(placement, self._expert_logical_shape(param))
+        # Checkpoints publish this map; planning from a different one would move other elements.
+        published = getattr(param, DS_AUTOEP_UC_META, None)
+        if published is not None:
+            if not autoep_placements_equal(published[AUTOEP_EXPERT_PLACEMENT], placement):
+                raise ValueError(f"AutoEP expert parameter '{name}' publishes a placement that differs from the "
+                                 "layer's expert placement descriptor.")
+            if published[AFFINE_MAP] != affine_map.to_dict():
+                raise ValueError(f"AutoEP expert parameter '{name}' publishes an affine map that differs from its "
+                                 "expert placement descriptor.")
+        return affine_map
+
+    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.
+        """
+        placement = self._live_expert_placement()
+        return {name: self._live_expert_affine_map(name, placement) for name in ("w1", "w2", "w3")}
+
+    def plan_expert_transfer(self,
+                             target_placement: dict,
+                             limits: PlanBudget | None = None,
+                             *,
+                             replica_selector=None) -> dict[str, TransferPlan]:
+        """Plan the copies that move the live expert parameters to ``target_placement``.
+
+        ``target_placement`` is an AutoEP placement descriptor in the target's EP-local ranks. This
+        only plans: it moves no data and does not check that this layer can run with the target.
+        """
+        placement = self._live_expert_placement()
+        plans = {}
+        for name in ("w1", "w2", "w3"):
+            param = getattr(self.experts, name)
+            self._live_expert_affine_map(name, placement)
+            plans[name] = plan_autoep_placement_transfer(placement,
+                                                         target_placement,
+                                                         self._expert_logical_shape(param),
+                                                         limits,
+                                                         replica_selector=replica_selector)
+        return plans
+
     def set_deepspeed_parallelism(
         self,
         use_data_before_expert_parallel_: bool = False,

docs/code-docs/source/autoep.rst

@@ -644,3 +644,55 @@ or ``None`` for automatic sizing, even when the selected backend does not use
 row blocking. Invalid settings are rejected before a backend runs. A model
 with a non-writable stock ``loss_function`` raises an actionable error rather
 than being silently left unmodified.
+
+Planning expert placement transfers
+-----------------------------------
+
+An ``AutoEPMoELayer`` can describe where its live expert weights are, and plan
+the copies that would move them to another expert placement, for example one
+with a different EP size. Planning moves no data, does not
+plan optimizer states, and does not change the placement the layer runs with.
+
+.. code-block:: python
+
+    from deepspeed.checkpoint.autoep_affine import make_autoep_placement_descriptor
+
+    # layer: an AutoEPMoELayer after deepspeed.initialize()
+    affine_maps = layer.expert_affine_maps()  # {"w1": ParamAffineMap, "w2": ..., "w3": ...}
+
+    # Move 8 experts onto 2 EP ranks: rank 0 holds experts 0-2, rank 1 holds experts 3-7.
+    target = make_autoep_placement_descriptor(8, [[0, 1, 2], [3, 4, 5, 6, 7]])
+    plans = layer.plan_expert_transfer(target)  # {"w1": TransferPlan, "w2": ..., "w3": ...}
+    for segment in plans["w1"].segments:
+        ...  # copy segment.shape elements from the source shard into the target shard
+
+- ``make_autoep_placement_descriptor(num_experts, experts_by_rank)`` builds the
+  target placement. ``experts_by_rank[r]`` lists the global expert IDs that
+  target rank ``r`` holds, in the order they are packed in its local tensor.
+  Ranks may hold different numbers of experts, non-contiguous experts, no
+  experts, or experts that other ranks also hold, but every expert must be
+  held by at least one rank.
+- Ranks in a plan are EP-local. ``segment.source_rank`` is a rank of the
+  layer's EP group; translate it with
+  ``dist.get_global_rank(layer.ep_group, segment.source_rank)``.
+  ``segment.target_rank`` indexes ``experts_by_rank`` of the target placement;
+  mapping it to a process is up to the caller, since the target EP group may
+  differ from the current one.
+- ``source_offset``/``target_offset`` and the stride vectors count elements of
+  a rank's packed local expert tensor ``[local_experts, ...]`` in row-major
+  order. That tensor is the unpartitioned parameter: under ZeRO-3, gather the
+  parameter (for example with ``deepspeed.zero.GatheredParameters``) before
+  applying the offsets. Ranks are EP-local, so expert-data-parallel replicas
+  are not candidates here.
+- When several source ranks hold an element, the plan reads it from the
+  lowest-ranked holder. Pass ``replica_selector``, a callable that receives a
+  ``ReplicaSelectionRequest`` and returns the ``source_rank`` of one of its
+  ``candidates``, to choose a different holder.
+- ``limits`` takes a ``deepspeed.checkpoint.affine_transfer.PlanBudget``. A plan
+  that exceeds it raises ``AffineTransferError`` instead of being truncated.
+- The runtime packs experts with the contiguous, uniform placement. If the
+  layer's placement descriptor describes anything else, both methods raise
+  ``NotImplementedError``.
+
+The maps and plans use the affine representation that Universal Checkpoint
+uses for expert parameters; see ``deepspeed/checkpoint/affine_ir_spec.md``.

tests/unit/checkpoint/test_autoep_affine.py

@@ -9,10 +9,11 @@
 from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION, ParamAffineMap
 from deepspeed.checkpoint.autoep_affine import (AUTOEP_PLACEMENT_VERSION, autoep_experts_for_rank,
                                                 autoep_metadata_to_affine_map, autoep_placement_to_affine_map,
-                                                extract_autoep_rank_tensor, legacy_uniform_autoep_placement_descriptor,
+                                                autoep_placements_equal, extract_autoep_rank_tensor,
+                                                legacy_uniform_autoep_placement_descriptor,
                                                 make_autoep_placement_descriptor, validate_autoep_placement_descriptor)
 from deepspeed.checkpoint.constants import (AFFINE_MAP, AFFINE_MAP_PARAMS, AFFINE_MAP_VERSION, AUTOEP_AFFINE_MAPS,
-                                            AUTOEP_EXPERT_PLACEMENT)
+                                            AUTOEP_EXPERT_PLACEMENT, AUTOEP_PLACEMENT_RANKS)
 
 
 def _oracle_shards(full_tensor, experts_by_rank):
@@ -236,3 +237,12 @@ def test_extract_target_rank_uses_target_packed_order():
     actual = extract_autoep_rank_tensor(universal_tensor, target_map, 1)
 
     assert torch.equal(actual, expected)
+
+
+def test_placements_equal_identifies_ranks_by_id():
+    descriptor = make_autoep_placement_descriptor(4, [[3, 0], [], [2, 1, 0]])
+    reordered = dict(descriptor)
+    reordered[AUTOEP_PLACEMENT_RANKS] = list(reversed(descriptor[AUTOEP_PLACEMENT_RANKS]))
+
+    assert autoep_placements_equal(descriptor, reordered)
+    assert not autoep_placements_equal(descriptor, make_autoep_placement_descriptor(4, [[0, 3], [], [2, 1, 0]]))