Add affine AutoEP checkpoint placement

deepspeed/checkpoint/affine.py

@@ -234,6 +234,8 @@ def rebuild(self, shards, scale_power=1):
 
         written = {}
         for rank, pieces in self.pieces_by_rank.items():
+            if not pieces:
+                continue
             flat_shard = _flat_buffer(shards[rank])
             for piece in pieces:
                 target = piece.source_view(full_param)

deepspeed/checkpoint/affine_ir_spec.md

@@ -568,11 +568,22 @@ and fp16 alike, because dividing by `2^k` only shifts the exponent — so a bias
 bit-exactly at every TP degree in normal use. It is lossy for non-power-of-two `N` (3, 6,
 12), where a converted-and-restored bias may differ in the last bits from the original.
 
-**8.4 ZeRO and offload placement.** delock's extension in #8230 — "a subset of a parameter
-combined with a list of ranks holding this subset" — is what §2.1's `locations` implements.
-ZeRO-1/3 partitions and offload replicas should fall out as pieces whose `locations`
-describe the DP group rather than the TP group, but this spec does not yet work through
-AutoEP's expert placement, where locations are per-expert rather than per-parameter.
+**8.4 AutoEP, ZeRO, and offload placement.** Phase 1 additively introduces a versioned
+AutoEP placement descriptor in EP-local rank coordinates and lowers one logical
+`[num_experts, ...]` tensor through the existing `AffinePiece` / `ParamAffineMap` IR. The
+descriptor records each rank's ordered global expert IDs, including uneven,
+non-contiguous, replicated, and empty placements. It describes placement provenance, not
+a scheduling policy. ZeRO and EDP fragments remain outside the map: callers first
+normalize storage to one logical packed expert tensor per EP rank.
+
+This phase does not change the current runtime's uniform contiguous scheduling, choose an
+arbitrary future expert schedule, or implement direct phase-2 shard-to-shard transfer.
+Those remain follow-on work. Phase 2 may derive a target descriptor from runtime
+scheduling and transfer directly between source and target maps; until then extraction
+from the universal full tensor uses the target map. delock's extension in #8230 — "a
+subset of a parameter combined with a list of ranks holding this subset" — is still what
+§2.1's exact `locations` implements, and the same mechanism can later describe normalized
+ZeRO/offload placement without introducing another geometry IR.
 
 ---
 

deepspeed/checkpoint/autoep_affine.py

@@ -0,0 +1,198 @@
+# SPDX-License-Identifier: Apache-2.0
+# 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.
+"""
+
+from deepspeed.checkpoint.affine import AffinePiece, ParamAffineMap
+from deepspeed.checkpoint.constants import (AUTOEP_PLACEMENT_EP_SIZE, AUTOEP_PLACEMENT_EXPERTS,
+                                            AUTOEP_PLACEMENT_NUM_EXPERTS, AUTOEP_PLACEMENT_RANK,
+                                            AUTOEP_PLACEMENT_RANKS, AUTOEP_PLACEMENT_VERSION,
+                                            AUTOEP_PLACEMENT_VERSION_KEY)
+
+__all__ = [
+    'AUTOEP_PLACEMENT_VERSION',
+    'make_autoep_placement_descriptor',
+    'validate_autoep_placement_descriptor',
+    'legacy_uniform_autoep_placement_descriptor',
+    'autoep_placement_to_affine_map',
+    'extract_autoep_rank_tensor',
+]
+
+
+def make_autoep_placement_descriptor(num_experts, experts_by_rank):
+    """Build a versioned descriptor from expert IDs in rank-local packed order."""
+    ranks = [{
+        AUTOEP_PLACEMENT_RANK: rank,
+        AUTOEP_PLACEMENT_EXPERTS: list(experts),
+    } for rank, experts in enumerate(experts_by_rank)]
+    descriptor = {
+        AUTOEP_PLACEMENT_VERSION_KEY: AUTOEP_PLACEMENT_VERSION,
+        AUTOEP_PLACEMENT_NUM_EXPERTS: num_experts,
+        AUTOEP_PLACEMENT_EP_SIZE: len(ranks),
+        AUTOEP_PLACEMENT_RANKS: ranks,
+    }
+    validate_autoep_placement_descriptor(descriptor)
+    return descriptor
+
+
+def validate_autoep_placement_descriptor(descriptor):
+    """Validate a plain scalar/list/dict AutoEP placement descriptor."""
+    if not isinstance(descriptor, dict):
+        raise TypeError('AutoEP placement descriptor must be a dict.')
+
+    version = descriptor.get(AUTOEP_PLACEMENT_VERSION_KEY)
+    if version != AUTOEP_PLACEMENT_VERSION:
+        raise ValueError(f'Unsupported AutoEP placement descriptor version {version!r}; '
+                         f'expected {AUTOEP_PLACEMENT_VERSION}.')
+
+    num_experts = _positive_int(descriptor.get(AUTOEP_PLACEMENT_NUM_EXPERTS), AUTOEP_PLACEMENT_NUM_EXPERTS)
+    ep_size = _positive_int(descriptor.get(AUTOEP_PLACEMENT_EP_SIZE), AUTOEP_PLACEMENT_EP_SIZE)
+    ranks = descriptor.get(AUTOEP_PLACEMENT_RANKS)
+    if not isinstance(ranks, list):
+        raise TypeError(f'{AUTOEP_PLACEMENT_RANKS} must be a list.')
+    if len(ranks) != ep_size:
+        raise ValueError(f'AutoEP placement descriptor has {len(ranks)} rank entries; expected exactly {ep_size}.')
+
+    seen_ranks = set()
+    covered_experts = set()
+    for entry in ranks:
+        if not isinstance(entry, dict):
+            raise TypeError('Each AutoEP placement rank entry must be a dict.')
+        rank = entry.get(AUTOEP_PLACEMENT_RANK)
+        if not isinstance(rank, int) or isinstance(rank, bool) or not 0 <= rank < ep_size:
+            raise ValueError(f'AutoEP placement rank {rank!r} is outside [0, {ep_size}).')
+        if rank in seen_ranks:
+            raise ValueError(f'AutoEP placement rank {rank} appears more than once.')
+        seen_ranks.add(rank)
+
+        experts = entry.get(AUTOEP_PLACEMENT_EXPERTS)
+        if not isinstance(experts, list):
+            raise TypeError(f'Experts for AutoEP placement rank {rank} must be a list defining local packed order.')
+        local_experts = set()
+        for expert_id in experts:
+            if not isinstance(expert_id, int) or isinstance(expert_id, bool) or not 0 <= expert_id < num_experts:
+                raise ValueError(f'AutoEP placement expert ID {expert_id!r} on rank {rank} is outside '
+                                 f'[0, {num_experts}).')
+            if expert_id in local_experts:
+                raise ValueError(f'AutoEP placement rank {rank} contains duplicate expert ID {expert_id}.')
+            local_experts.add(expert_id)
+            covered_experts.add(expert_id)
+
+    missing = sorted(set(range(num_experts)) - covered_experts)
+    if missing:
+        raise ValueError(f'AutoEP placement descriptor does not cover global expert IDs {missing}.')
+
+
+def legacy_uniform_autoep_placement_descriptor(num_experts, num_local_experts, ep_size):
+    """Synthesize the legacy contiguous, uniform AutoEP placement."""
+    num_experts = _positive_int(num_experts, 'num_experts')
+    num_local_experts = _positive_int(num_local_experts, 'num_local_experts')
+    ep_size = _positive_int(ep_size, 'ep_size')
+    if num_local_experts * ep_size != num_experts:
+        raise ValueError(f'Inconsistent legacy AutoEP metadata: num_local_experts ({num_local_experts}) * '
+                         f'ep_size ({ep_size}) != num_experts ({num_experts}).')
+    experts_by_rank = []
+    for rank in range(ep_size):
+        start = rank * num_local_experts
+        experts_by_rank.append(list(range(start, start + num_local_experts)))
+    return make_autoep_placement_descriptor(num_experts, experts_by_rank)
+
+
+def autoep_placement_to_affine_map(descriptor, logical_shape):
+    """Lower one ``[num_experts, ...]`` parameter to a :class:`ParamAffineMap`."""
+    validate_autoep_placement_descriptor(descriptor)
+    logical_shape = tuple(int(dim) for dim in logical_shape)
+    num_experts = descriptor[AUTOEP_PLACEMENT_NUM_EXPERTS]
+    if not logical_shape or logical_shape[0] != num_experts:
+        raise ValueError(f'Expert parameter logical shape must start with num_experts ({num_experts}); '
+                         f'got {logical_shape}.')
+    if any(dim < 0 for dim in logical_shape):
+        raise ValueError(f'Expert parameter logical shape cannot contain negative dimensions: {logical_shape}.')
+
+    entries_by_rank = {entry[AUTOEP_PLACEMENT_RANK]: entry for entry in descriptor[AUTOEP_PLACEMENT_RANKS]}
+    holders = _expert_holders(entries_by_rank, num_experts)
+    source_strides = _row_major_strides(logical_shape)
+    expert_shape = logical_shape[1:]
+    expert_numel = _product(expert_shape)
+    pieces_by_rank = {}
+    shard_shapes = {}
+
+    for rank in range(descriptor[AUTOEP_PLACEMENT_EP_SIZE]):
+        experts = entries_by_rank[rank][AUTOEP_PLACEMENT_EXPERTS]
+        shard_shape = (len(experts), ) + expert_shape
+        shard_shapes[rank] = shard_shape
+        dest_strides = _row_major_strides(shard_shape)
+        pieces_by_rank[rank] = _pieces_for_rank(experts, holders, expert_shape, expert_numel, source_strides,
+                                                dest_strides)
+
+    affine_map = ParamAffineMap(logical_shape=logical_shape, shard_shapes=shard_shapes, pieces_by_rank=pieces_by_rank)
+    # Descriptor validation already proves expert-level coverage. Expanding that
+    # check to every tensor element is prohibitively expensive for expert weights.
+    affine_map.validate()
+    return affine_map
+
+
+def extract_autoep_rank_tensor(full_param, target_map, ep_rank):
+    """Extract one EP rank's packed local tensor from a universal full tensor."""
+    if not isinstance(target_map, ParamAffineMap):
+        raise TypeError('target_map must be a ParamAffineMap.')
+    if ep_rank not in target_map.shard_shapes:
+        raise ValueError(f'EP rank {ep_rank} is not present in the target affine map.')
+    return target_map.extract(full_param, ep_rank)
+
+
+def _pieces_for_rank(experts, holders, expert_shape, expert_numel, source_strides, dest_strides):
+    pieces = []
+    run_start = 0
+    for local_index in range(1, len(experts) + 1):
+        at_end = local_index == len(experts)
+        if not at_end:
+            previous_expert = experts[local_index - 1]
+            current_expert = experts[local_index]
+            homogeneous = (current_expert == previous_expert + 1
+                           and holders[current_expert] == holders[previous_expert]
+                           and len(holders[current_expert]) == 1)
+        if at_end or not homogeneous:
+            first_expert = experts[run_start]
+            run_length = local_index - run_start
+            pieces.append(
+                AffinePiece(shape=(run_length, ) + expert_shape,
+                            source_offset=first_expert * expert_numel,
+                            source_strides=source_strides,
+                            dest_offset=run_start * expert_numel,
+                            dest_strides=dest_strides,
+                            locations=holders[first_expert]))
+            run_start = local_index
+    return pieces
+
+
+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():
+        for expert_id in entry[AUTOEP_PLACEMENT_EXPERTS]:
+            holders[expert_id].add(rank)
+    return holders
+
+
+def _positive_int(value, name):
+    if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
+        raise ValueError(f'{name} must be a positive integer; got {value!r}.')
+    return value
+
+
+def _product(shape):
+    count = 1
+    for dim in shape:
+        count *= dim
+    return count
+
+
+def _row_major_strides(shape):
+    strides = [1] * len(shape)
+    for axis in range(len(shape) - 2, -1, -1):
+        strides[axis] = strides[axis + 1] * shape[axis + 1]
+    return tuple(strides)

deepspeed/checkpoint/autoep_universal.py

@@ -16,9 +16,15 @@
 
 from .constants import (
     AUTOEP_EP_SIZE,
+    AUTOEP_EXPERT_PLACEMENT,
     AUTOEP_EXPERT_KEY_PREFIX,
     AUTOEP_NUM_EXPERTS,
     AUTOEP_NUM_LOCAL_EXPERTS,
+    AUTOEP_PLACEMENT_EP_SIZE,
+    AUTOEP_PLACEMENT_EXPERTS,
+    AUTOEP_PLACEMENT_NUM_EXPERTS,
+    AUTOEP_PLACEMENT_RANK,
+    AUTOEP_PLACEMENT_RANKS,
     AUTOEP_ZERO12_REQUIRED_FIELDS,
     PARAM,
     CAT_DIM,
@@ -40,6 +46,7 @@
     FOLDING_FAMILY,
     FOLDING_PARAM_FAMILIES,
 )
+from .autoep_affine import autoep_placement_to_affine_map, validate_autoep_placement_descriptor
 
 
 def make_folding_metadata(*,
@@ -229,22 +236,49 @@ def get_autoep_zero12_expert_param_info(autoep_layers_metadata):
         if not isinstance(prefix, str) or not prefix:
             raise RuntimeError("AutoEP expert_key_prefix must be a non-empty string.")
 
-        for field in (AUTOEP_NUM_EXPERTS, AUTOEP_NUM_LOCAL_EXPERTS, AUTOEP_EP_SIZE):
+        for field in (AUTOEP_NUM_EXPERTS, AUTOEP_EP_SIZE):
             value = layer_info[field]
             if isinstance(value, bool) or not isinstance(value, int) or value < 1:
                 raise RuntimeError(f"AutoEP {field} must be a positive integer, got {value!r}.")
 
         num_experts = layer_info[AUTOEP_NUM_EXPERTS]
         num_local_experts = layer_info[AUTOEP_NUM_LOCAL_EXPERTS]
         ep_size = layer_info[AUTOEP_EP_SIZE]
-        if num_experts != num_local_experts * ep_size:
+        placement = layer_info.get(AUTOEP_EXPERT_PLACEMENT)
+        if isinstance(num_local_experts, bool) or not isinstance(num_local_experts, int):
+            raise RuntimeError(f"AutoEP {AUTOEP_NUM_LOCAL_EXPERTS} must be an integer, got "
+                               f"{num_local_experts!r}.")
+        minimum_local_experts = 0 if placement is not None else 1
+        if num_local_experts < minimum_local_experts:
+            qualifier = "non-negative" if placement is not None else "positive"
+            raise RuntimeError(f"AutoEP {AUTOEP_NUM_LOCAL_EXPERTS} must be a {qualifier} integer, got "
+                               f"{num_local_experts!r}.")
+        if placement is not None:
+            try:
+                validate_autoep_placement_descriptor(placement)
+            except (TypeError, ValueError) as exc:
+                raise RuntimeError(f"Invalid AutoEP expert placement for {prefix}: {exc}") from exc
+            if (placement[AUTOEP_PLACEMENT_NUM_EXPERTS] != num_experts
+                    or placement[AUTOEP_PLACEMENT_EP_SIZE] != ep_size):
+                raise RuntimeError(f"AutoEP expert placement disagrees with layer metadata for {prefix}.")
+            ep_rank = layer_info.get('ep_rank')
+            if ep_rank is not None:
+                rank_entries = {entry[AUTOEP_PLACEMENT_RANK]: entry for entry in placement[AUTOEP_PLACEMENT_RANKS]}
+                if ep_rank not in rank_entries:
+                    raise RuntimeError(f"AutoEP expert placement does not contain metadata ep_rank {ep_rank}.")
+                expected_local_experts = len(rank_entries[ep_rank][AUTOEP_PLACEMENT_EXPERTS])
+                if num_local_experts != expected_local_experts:
+                    raise RuntimeError("AutoEP num_local_experts disagrees with the placement entry for "
+                                       f"EP rank {ep_rank}: {num_local_experts} != {expected_local_experts}.")
+        elif num_experts != num_local_experts * ep_size:
             raise RuntimeError(f"AutoEP expert count mismatch for {prefix}: num_experts={num_experts}, "
                                f"num_local_experts={num_local_experts}, ep_size={ep_size}.")
 
         normalized = {
             'num_experts': num_experts,
             'num_local_experts': num_local_experts,
             'ep_size': ep_size,
+            'expert_placement': placement,
         }
         for weight_name in ('w1', 'w2', 'w3'):
             param_name = f"{prefix}.{weight_name}"
@@ -288,19 +322,27 @@ def consolidate_autoep_zero12_expert_states(temp_dir, output_dir, expert_param_i
 
         ep_size = metadata['ep_size']
         num_experts = metadata['num_experts']
-        num_local_experts = metadata['num_local_experts']
+        placement = metadata.get('expert_placement')
 
         local_shape = tuple(slice_shapes[param_name])
-        if not local_shape or local_shape[0] != num_local_experts:
+        if not local_shape:
+            raise RuntimeError(f"AutoEP local shape is empty for {param_name}.")
+
+        affine_map = None
+        if placement is not None:
+            logical_shape = (num_experts, ) + local_shape[1:]
+            affine_map = autoep_placement_to_affine_map(placement, logical_shape)
+        elif local_shape[0] != metadata['num_local_experts']:
             raise RuntimeError(f"AutoEP local shape mismatch for {param_name}: shape={local_shape}, "
-                               f"num_local_experts={num_local_experts}.")
+                               f"num_local_experts={metadata['num_local_experts']}.")
 
         param_dir = os.path.join(output_dir, "zero", param_name)
         os.makedirs(param_dir, exist_ok=True)
 
         for state_name in ('fp32', 'exp_avg', 'exp_avg_sq'):
-            ep_tensors = []
+            ep_tensors = {}
             for ep_rank in range(ep_size):
+                expected_shape = affine_map.shard_shapes[ep_rank] if affine_map is not None else local_shape
                 fragments = []
                 dp_ranks = _autoep_zero12_dp_ranks(ep_rank, dp_degree, ep_size, use_data_before_expert_parallel)
                 for dp_rank in dp_ranks:
@@ -315,17 +357,24 @@ def consolidate_autoep_zero12_expert_states(temp_dir, output_dir, expert_param_i
                                            f"in {fragment_path}.")
                     fragments.append(fragment.flatten())
 
+                expected_numel = torch.Size(expected_shape).numel()
+                if not fragments and expected_numel == 0:
+                    local_tensor = torch.empty(expected_shape, dtype=torch.float32)
+                    ep_tensors[ep_rank] = local_tensor
+                    continue
                 if not fragments:
                     raise RuntimeError(f"Missing AutoEP {state_name} fragments for {param_name}, EP rank {ep_rank}.")
 
                 local_tensor = torch.cat(fragments, dim=0)
-                expected_numel = torch.Size(local_shape).numel()
                 if local_tensor.numel() != expected_numel:
                     raise RuntimeError(f"AutoEP {state_name} fragment size mismatch for {param_name}, "
                                        f"EP rank {ep_rank}: got {local_tensor.numel()}, expected {expected_numel}.")
-                ep_tensors.append(local_tensor.reshape(local_shape))
+                ep_tensors[ep_rank] = local_tensor.reshape(expected_shape)
 
-            full_tensor = torch.cat(ep_tensors, dim=0)
+            if affine_map is not None:
+                full_tensor = affine_map.rebuild(ep_tensors)
+            else:
+                full_tensor = torch.cat([ep_tensors[rank] for rank in range(ep_size)], dim=0)
             if full_tensor.shape[0] != num_experts:
                 raise RuntimeError(f"AutoEP consolidated expert count mismatch for {param_name}: "
                                    f"got {full_tensor.shape[0]}, expected {num_experts}.")
@@ -376,6 +425,11 @@ def consolidate_autoep_expert_files(checkpoint_dir, output_dir, autoep_layers_me
         moe_layer_id = layer_info['moe_layer_id']
         num_experts = layer_info['num_experts']
         prefix = layer_info['expert_key_prefix']
+        placement = layer_info.get(AUTOEP_EXPERT_PLACEMENT)
+        if placement is not None:
+            validate_autoep_placement_descriptor(placement)
+            if placement[AUTOEP_NUM_EXPERTS] != num_experts:
+                raise RuntimeError(f"AutoEP expert placement disagrees with num_experts for {prefix}.")
 
         for wname in ('w1', 'w2', 'w3'):
             expert_tensors = []