From e88ffcf9df80263d720b086dd5708bd289cdd87f Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Tue, 15 Sep 2026 11:06:46 +0000 Subject: [PATCH 1/6] Add affine AutoEP checkpoint placement Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- deepspeed/checkpoint/affine.py | 2 + deepspeed/checkpoint/affine_ir_spec.md | 21 +- deepspeed/checkpoint/autoep_affine.py | 198 +++++++ deepspeed/checkpoint/autoep_universal.py | 72 ++- deepspeed/checkpoint/autoep_zero3_metadata.py | 90 +++- deepspeed/checkpoint/constants.py | 12 + deepspeed/checkpoint/ds_to_universal.py | 87 ++- deepspeed/checkpoint/universal_checkpoint.py | 81 ++- deepspeed/module_inject/auto_ep_layer.py | 21 + deepspeed/runtime/engine.py | 51 +- deepspeed/runtime/zero/stage3.py | 8 +- tests/unit/checkpoint/test_autoep_affine.py | 185 +++++++ .../test_autoep_affine_universal.py | 494 ++++++++++++++++++ .../checkpoint/test_autoep_zero3_metadata.py | 124 +++++ tests/unit/v1/moe/test_autoep_unit.py | 86 +++ 15 files changed, 1476 insertions(+), 56 deletions(-) create mode 100644 deepspeed/checkpoint/autoep_affine.py create mode 100644 tests/unit/checkpoint/test_autoep_affine.py create mode 100644 tests/unit/checkpoint/test_autoep_affine_universal.py create mode 100644 tests/unit/checkpoint/test_autoep_zero3_metadata.py diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index 8ca076ae6667..ac2a829dc611 100644 --- a/deepspeed/checkpoint/affine.py +++ b/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) diff --git a/deepspeed/checkpoint/affine_ir_spec.md b/deepspeed/checkpoint/affine_ir_spec.md index 4fd1ddf26bfc..993163d6cf15 100644 --- a/deepspeed/checkpoint/affine_ir_spec.md +++ b/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. --- diff --git a/deepspeed/checkpoint/autoep_affine.py b/deepspeed/checkpoint/autoep_affine.py new file mode 100644 index 000000000000..cd615c8c43de --- /dev/null +++ b/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) diff --git a/deepspeed/checkpoint/autoep_universal.py b/deepspeed/checkpoint/autoep_universal.py index 9ecfd11d35a1..c4d82b428f3a 100644 --- a/deepspeed/checkpoint/autoep_universal.py +++ b/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,7 +236,7 @@ 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}.") @@ -237,7 +244,33 @@ def get_autoep_zero12_expert_param_info(autoep_layers_metadata): 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}.") @@ -245,6 +278,7 @@ def get_autoep_zero12_expert_param_info(autoep_layers_metadata): '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 = [] diff --git a/deepspeed/checkpoint/autoep_zero3_metadata.py b/deepspeed/checkpoint/autoep_zero3_metadata.py index f854d17c2738..6b7766dd428f 100644 --- a/deepspeed/checkpoint/autoep_zero3_metadata.py +++ b/deepspeed/checkpoint/autoep_zero3_metadata.py @@ -4,7 +4,15 @@ # DeepSpeed Team """Shared validation for AutoEP ZeRO-3 checkpoint metadata.""" +from deepspeed.checkpoint.autoep_affine import (legacy_uniform_autoep_placement_descriptor, + validate_autoep_placement_descriptor) from deepspeed.checkpoint.constants import ( + AUTOEP_EXPERT_PLACEMENT, + AUTOEP_PLACEMENT_EP_SIZE, + AUTOEP_PLACEMENT_EXPERTS, + AUTOEP_PLACEMENT_NUM_EXPERTS, + AUTOEP_PLACEMENT_RANK, + AUTOEP_PLACEMENT_RANKS, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY, @@ -39,6 +47,7 @@ def is_autoep_zero3_partitioned_entry(entry): def validate_autoep_zero3_partitioned_metadata(autoep_metadata, require_partitioned=True, expected_expert_prefixes=None, + expected_runtime_layers=None, version_context="This DeepSpeed build"): if not isinstance(autoep_metadata, list): raise RuntimeError(f"ds_autoep_layers metadata is malformed: expected list, got " @@ -47,6 +56,7 @@ def validate_autoep_zero3_partitioned_metadata(autoep_metadata, seen_layer_ids = set() seen_prefixes = set() partitioned_count = 0 + placements = [] for entry in autoep_metadata: if not isinstance(entry, dict): @@ -66,6 +76,10 @@ def validate_autoep_zero3_partitioned_metadata(autoep_metadata, raise RuntimeError(f"ds_autoep_layers metadata has duplicate expert_key_prefix: {prefix}") seen_prefixes.add(prefix) + placement = _validated_placement(entry) + placements.append(placement) + _validate_runtime_layer(entry, placement, expected_runtime_layers) + if not is_autoep_zero3_partitioned_entry(entry): continue @@ -78,20 +92,21 @@ def validate_autoep_zero3_partitioned_metadata(autoep_metadata, f"{version}. {version_context} supports version " f"{AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION}.") - num_experts = entry['num_experts'] - num_local_experts = entry['num_local_experts'] - ep_size = entry['ep_size'] - if num_local_experts * ep_size != num_experts: - raise RuntimeError("AutoEP ZeRO-3 checkpoint metadata is inconsistent: " - f"num_local_experts={num_local_experts}, ep_size={ep_size}, " - f"num_experts={num_experts}") - - expected_start = entry['ep_rank'] * num_local_experts - expected_end = expected_start + num_local_experts - if entry['global_expert_start'] != expected_start or entry['global_expert_end'] != expected_end: - raise RuntimeError("AutoEP ZeRO-3 checkpoint metadata has inconsistent global expert range: " - f"got [{entry['global_expert_start']}, {entry['global_expert_end']}), " - f"expected [{expected_start}, {expected_end})") + if AUTOEP_EXPERT_PLACEMENT not in entry: + num_experts = entry['num_experts'] + num_local_experts = entry['num_local_experts'] + ep_size = entry['ep_size'] + if num_local_experts * ep_size != num_experts: + raise RuntimeError("AutoEP ZeRO-3 checkpoint metadata is inconsistent: " + f"num_local_experts={num_local_experts}, ep_size={ep_size}, " + f"num_experts={num_experts}") + + expected_start = entry['ep_rank'] * num_local_experts + expected_end = expected_start + num_local_experts + if entry['global_expert_start'] != expected_start or entry['global_expert_end'] != expected_end: + raise RuntimeError("AutoEP ZeRO-3 checkpoint metadata has inconsistent global expert range: " + f"got [{entry['global_expert_start']}, {entry['global_expert_end']}), " + f"expected [{expected_start}, {expected_end})") if expected_expert_prefixes is not None: module_path = entry['module_path'] @@ -107,3 +122,50 @@ def validate_autoep_zero3_partitioned_metadata(autoep_metadata, if require_partitioned and partitioned_count == 0: raise RuntimeError("AutoEP ZeRO-3 partition-native checkpoint metadata was expected but no " "partitioned AutoEP layer entries were found") + return placements + + +def _validated_placement(entry): + try: + if AUTOEP_EXPERT_PLACEMENT in entry: + placement = entry[AUTOEP_EXPERT_PLACEMENT] + validate_autoep_placement_descriptor(placement) + else: + placement = legacy_uniform_autoep_placement_descriptor(entry['num_experts'], entry['num_local_experts'], + entry['ep_size']) + except (TypeError, ValueError) as exc: + raise RuntimeError(f"AutoEP placement metadata is invalid: {exc}") from exc + + if placement[AUTOEP_PLACEMENT_NUM_EXPERTS] != entry['num_experts']: + raise RuntimeError("AutoEP placement metadata num_experts does not match its enclosing layer entry: " + f"{placement[AUTOEP_PLACEMENT_NUM_EXPERTS]} != {entry['num_experts']}") + if placement[AUTOEP_PLACEMENT_EP_SIZE] != entry['ep_size']: + raise RuntimeError("AutoEP placement metadata ep_size does not match its enclosing layer entry: " + f"{placement[AUTOEP_PLACEMENT_EP_SIZE]} != {entry['ep_size']}") + return placement + + +def _validate_runtime_layer(entry, placement, expected_runtime_layers): + ep_rank = entry.get('ep_rank') + if ep_rank is None: + return + ranks = {rank_entry[AUTOEP_PLACEMENT_RANK]: rank_entry for rank_entry in placement[AUTOEP_PLACEMENT_RANKS]} + if ep_rank not in ranks: + raise RuntimeError(f"AutoEP placement metadata does not contain current ep_rank {ep_rank}") + listed_experts = ranks[ep_rank][AUTOEP_PLACEMENT_EXPERTS] + if len(listed_experts) != entry['num_local_experts']: + raise RuntimeError("AutoEP placement metadata current-rank expert count does not match num_local_experts: " + f"{len(listed_experts)} != {entry['num_local_experts']}") + + if expected_runtime_layers is None: + return + module_path = entry['module_path'] + if module_path not in expected_runtime_layers: + raise RuntimeError(f"AutoEP checkpoint metadata references missing module: {module_path}") + runtime = expected_runtime_layers[module_path] + if ep_rank != runtime['ep_rank']: + raise RuntimeError(f"AutoEP checkpoint metadata ep_rank {ep_rank} does not match runtime ep_rank " + f"{runtime['ep_rank']} for module {module_path}") + if listed_experts != runtime['local_experts']: + raise RuntimeError("AutoEP placement metadata current-rank expert order does not match runtime local packing: " + f"{listed_experts} != {runtime['local_experts']} for module {module_path}") diff --git a/deepspeed/checkpoint/constants.py b/deepspeed/checkpoint/constants.py index f84ef8718064..230f22fabc1e 100644 --- a/deepspeed/checkpoint/constants.py +++ b/deepspeed/checkpoint/constants.py @@ -113,6 +113,11 @@ AUTOEP_NUM_EXPERTS = 'num_experts' AUTOEP_NUM_LOCAL_EXPERTS = 'num_local_experts' AUTOEP_EP_SIZE = 'ep_size' +AUTOEP_EXPERT_PLACEMENT = 'expert_placement' +DS_AUTOEP_UC_META = 'ds_autoep_uc_meta' +AUTOEP_PARAM_LOGICAL_SHAPE = 'logical_shape' +AUTOEP_PARAM_EP_RANK = 'ep_rank' +AUTOEP_PARAM_LOCAL_EXPERTS = 'local_experts' AUTOEP_ZERO12_REQUIRED_FIELDS = ( AUTOEP_EXPERT_KEY_PREFIX, AUTOEP_NUM_EXPERTS, @@ -123,6 +128,13 @@ AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT = 'zero3_partitioned' AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY = 'checkpoint_format_version' AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION = 1 +AUTOEP_PLACEMENT_VERSION_KEY = 'version' +AUTOEP_PLACEMENT_VERSION = 1 +AUTOEP_PLACEMENT_NUM_EXPERTS = 'num_experts' +AUTOEP_PLACEMENT_EP_SIZE = 'ep_size' +AUTOEP_PLACEMENT_RANKS = 'ranks' +AUTOEP_PLACEMENT_RANK = 'rank' +AUTOEP_PLACEMENT_EXPERTS = 'experts' ######################################### # Universal Checkpoint EP keys diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py index 032a33e7e151..08e90aa09700 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/deepspeed/checkpoint/ds_to_universal.py @@ -50,7 +50,11 @@ AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS, AUTOEP_LAYERS_KEY, AUTOEP_LAYERS_KEY_LEGACY, + AUTOEP_EXPERT_PLACEMENT, AUTOEP_EXPERT_KEY_PREFIX, + AUTOEP_PLACEMENT_EXPERTS, + AUTOEP_PLACEMENT_RANK, + AUTOEP_PLACEMENT_RANKS, EP_IS_EXPERT_PARAM, EP_NUM_EXPERTS, EXPERT_PARAMETER_PATTERNS, @@ -61,6 +65,7 @@ validate_autoep_zero3_partitioned_metadata, ) from deepspeed.checkpoint.affine import ParamAffineMap, AFFINE_MAP_FORMAT_VERSION +from deepspeed.checkpoint.autoep_affine import autoep_placement_to_affine_map, validate_autoep_placement_descriptor def parse_arguments(): @@ -743,9 +748,20 @@ def _validate_autoep_expert_shapes(model_states_by_rank, metadata_by_rank): param_shapes = model_states_by_rank[rank][PARAM_SHAPES] zero_shape_names = {name for sub_group_shape in param_shapes for name in sub_group_shape} missing = set(expert_info) - zero_shape_names - if missing: + missing_nonempty = [] + for param_name in missing: + layer_info = expert_info[param_name] + placement = layer_info.get(AUTOEP_EXPERT_PLACEMENT) + if placement is None: + missing_nonempty.append(param_name) + continue + entries_by_rank = {entry[AUTOEP_PLACEMENT_RANK]: entry for entry in placement[AUTOEP_PLACEMENT_RANKS]} + rank_entry = entries_by_rank[layer_info['ep_rank']] + if rank_entry[AUTOEP_PLACEMENT_EXPERTS]: + missing_nonempty.append(param_name) + if missing_nonempty: raise RuntimeError(f"AutoEP expert parameters are missing from rank {rank} ZeRO param_shapes: " - f"{sorted(missing)}") + f"{sorted(missing_nonempty)}") frozen_shapes = model_states_by_rank[rank].get('frozen_param_shapes') or {} frozen_experts = set(expert_info).intersection(frozen_shapes) if frozen_experts: @@ -767,6 +783,34 @@ def _save_zero3_autoep_universal_tensor(output_dir, param_name, state_key, tenso ) +def _rebuild_zero3_autoep_rank_tensors(ep_tensors, placement, logical_shape, context): + if placement is None: + return torch.cat([ep_tensors[rank] for rank in sorted(ep_tensors)], dim=0) + + affine_map = autoep_placement_to_affine_map(placement, logical_shape) + expected_ranks = set(affine_map.shard_shapes) + actual_ranks = set(ep_tensors) + unexpected_ranks = actual_ranks - expected_ranks + missing_nonempty_ranks = { + rank + for rank in expected_ranks - actual_ranks if torch.Size(affine_map.shard_shapes[rank]).numel() != 0 + } + if unexpected_ranks or missing_nonempty_ranks: + raise RuntimeError( + f"Incomplete AutoEP EP-rank tensors for {context}: got {sorted(actual_ranks)}, " + f"expected nonempty ranks " + f"{sorted(rank for rank in expected_ranks if torch.Size(affine_map.shard_shapes[rank]).numel())}.") + for rank, tensor in ep_tensors.items(): + expected_shape = affine_map.shard_shapes[rank] + if tuple(tensor.shape) != expected_shape: + raise RuntimeError(f"AutoEP EP-rank tensor shape mismatch for {context}, rank {rank}: " + f"got {tuple(tensor.shape)}, expected {expected_shape}.") + try: + return affine_map.rebuild(ep_tensors) + except ValueError as exc: + raise RuntimeError(f"Failed to rebuild AutoEP universal tensor for {context}: {exc}") from exc + + def _consolidate_zero3_autoep_expert_states(output_dir, model_files, optim_files): model_rank_map, optim_rank_map = _validate_zero3_model_optim_rank_sets(model_files, optim_files) model_states_by_rank = { @@ -787,6 +831,8 @@ def _consolidate_zero3_autoep_expert_states(output_dir, model_files, optim_files num_experts_by_param = {} expected_dp_world_by_param_rank = {} expected_ep_ranks_by_param = {} + placement_by_param = {} + logical_shape_by_param = {} for rank, model_state in model_states_by_rank.items(): optim_state = optim_states_by_rank.get(rank) @@ -824,9 +870,29 @@ def _consolidate_zero3_autoep_expert_states(output_dir, model_files, optim_files if layer_info is not None: ep_rank = layer_info['ep_rank'] num_experts_by_param[param_name] = layer_info['num_experts'] + placement = layer_info.get(AUTOEP_EXPERT_PLACEMENT) + if placement is not None: + try: + validate_autoep_placement_descriptor(placement) + except (TypeError, ValueError) as exc: + raise RuntimeError(f"Invalid AutoEP expert placement for {param_name}: {exc}") from exc + existing_placement = placement_by_param.setdefault(param_name, placement) + if existing_placement != placement: + raise RuntimeError(f"AutoEP expert placement disagrees across ranks for {param_name}.") + logical_shape = (layer_info['num_experts'], ) + tuple(shape)[1:] + existing_logical_shape = logical_shape_by_param.setdefault(param_name, logical_shape) + if existing_logical_shape != logical_shape: + raise RuntimeError(f"AutoEP expert logical shape disagrees across ranks for {param_name}: " + f"{existing_logical_shape} != {logical_shape}.") + affine_map = autoep_placement_to_affine_map(placement, logical_shape) + expected_shape = affine_map.shard_shapes.get(ep_rank) + if expected_shape is None or tuple(shape) != expected_shape: + raise RuntimeError(f"AutoEP expert shard shape mismatch for {param_name}, EP rank " + f"{ep_rank}: got {tuple(shape)}, expected {expected_shape}.") expected_dp_world_by_param_rank[(param_name, ep_rank)] = layer_info['expert_data_parallel_world_size'] - expected_ep_ranks_by_param[param_name] = set(range(layer_info['ep_size'])) + expected_ep_ranks_by_param[param_name] = (set(affine_map.shard_shapes) if placement is not None + else set(range(layer_info['ep_size']))) for state_key, flat_tensor in flat_state.items(): if flat_tensor is None: raise RuntimeError(f"Missing optimizer state '{state_key}' for AutoEP expert " @@ -843,10 +909,18 @@ def _consolidate_zero3_autoep_expert_states(output_dir, model_files, optim_files for (param_name, state_key), ep_rank_fragments in grouped_by_param.items(): missing_ep_ranks = expected_ep_ranks_by_param[param_name] - set(ep_rank_fragments) + placement = placement_by_param.get(param_name) + logical_shape = logical_shape_by_param.get(param_name) + if placement is not None: + affine_map = autoep_placement_to_affine_map(placement, logical_shape) + missing_ep_ranks = { + rank + for rank in missing_ep_ranks if torch.Size(affine_map.shard_shapes[rank]).numel() != 0 + } if missing_ep_ranks: raise RuntimeError(f"Missing AutoEP universal fragments for {param_name}/{state_key} EP ranks: " f"{sorted(missing_ep_ranks)}") - ep_tensors = [] + ep_tensors = {} for ep_rank in sorted(ep_rank_fragments): fragments = sorted(ep_rank_fragments[ep_rank], key=lambda item: item[0]) expected_dp_world = expected_dp_world_by_param_rank[(param_name, ep_rank)] @@ -863,11 +937,12 @@ def _consolidate_zero3_autoep_expert_states(output_dir, model_files, optim_files raise RuntimeError(f"Inconsistent AutoEP expert fragment shapes for {param_name}/{state_key} " f"EP rank {ep_rank}") full_flat = torch.cat([fragment for _, fragment, _ in fragments], dim=0)[:_shape_numel(shape)] - ep_tensors.append(full_flat.view(shape)) + ep_tensors[ep_rank] = full_flat.view(shape) if not ep_tensors: continue - full_expert_tensor = torch.cat(ep_tensors, dim=0) + full_expert_tensor = _rebuild_zero3_autoep_rank_tensors(ep_tensors, placement, logical_shape, + f"{param_name}/{state_key}") if full_expert_tensor.shape[0] != num_experts_by_param[param_name]: raise RuntimeError(f"AutoEP universal tensor for {param_name}/{state_key} has wrong expert dimension: " f"got {full_expert_tensor.shape[0]}, expected {num_experts_by_param[param_name]}") diff --git a/deepspeed/checkpoint/universal_checkpoint.py b/deepspeed/checkpoint/universal_checkpoint.py index 33e0ee7e9d98..c07da4fc5c41 100644 --- a/deepspeed/checkpoint/universal_checkpoint.py +++ b/deepspeed/checkpoint/universal_checkpoint.py @@ -10,7 +10,11 @@ from typing import List, Tuple, Union from dataclasses import dataclass from .constants import (FP32_WEIGHT_KEY, PARAM, VOCAB_TENSOR, CAT_DIM, PARAM_N_SUB_PARAMS, SUB_PARAM_SHAPE, - EP_IS_EXPERT_PARAM, EP_NUM_EXPERTS, DS_AUTOTP_UC_META, UNIVERSAL_CHECKPOINT_VERSION_KEY) + EP_IS_EXPERT_PARAM, EP_NUM_EXPERTS, DS_AUTOEP_UC_META, DS_AUTOTP_UC_META, + AUTOEP_EXPERT_PLACEMENT, AUTOEP_PARAM_EP_RANK, AUTOEP_PARAM_LOCAL_EXPERTS, + AUTOEP_PARAM_LOGICAL_SHAPE, AUTOEP_PLACEMENT_EXPERTS, AUTOEP_PLACEMENT_RANK, + AUTOEP_PLACEMENT_RANKS, UNIVERSAL_CHECKPOINT_VERSION_KEY) +from .autoep_affine import autoep_placement_to_affine_map, extract_autoep_rank_tensor @dataclass @@ -31,6 +35,75 @@ def _get_param_uc_restore_meta(param): return getattr(param, DS_AUTOTP_UC_META, None) +def _resolve_autoep_partition(current_param, ckpt_dict, full_hp_param, ep_rank): + meta = getattr(current_param, DS_AUTOEP_UC_META, None) + if meta is None: + return None + if not isinstance(meta, dict): + raise RuntimeError(f"AutoEP universal checkpoint target metadata must be a dict, got {type(meta).__name__}.") + + required = { + AUTOEP_EXPERT_PLACEMENT, + AUTOEP_PARAM_LOGICAL_SHAPE, + AUTOEP_PARAM_EP_RANK, + AUTOEP_PARAM_LOCAL_EXPERTS, + } + missing = required - meta.keys() + if missing: + raise RuntimeError(f"AutoEP universal checkpoint target metadata is missing fields {sorted(missing)}.") + + etp_size = meta.get('expert_tensor_parallel_size', meta.get('etp_size', 1)) + if etp_size != 1: + raise NotImplementedError("Universal checkpoint restore for AutoEP expert tensor parallelism is not " + f"supported; got expert_tensor_parallel_size={etp_size}.") + + logical_shape_value = meta[AUTOEP_PARAM_LOGICAL_SHAPE] + if not isinstance(logical_shape_value, (tuple, list)): + raise RuntimeError("AutoEP universal checkpoint target logical_shape must be a tuple or list.") + logical_shape = tuple(logical_shape_value) + + checkpoint_num_experts = ckpt_dict.get(EP_NUM_EXPERTS) + if checkpoint_num_experts is None: + raise RuntimeError(f"AutoEP universal checkpoint is missing '{EP_NUM_EXPERTS}' metadata.") + placement = meta[AUTOEP_EXPERT_PLACEMENT] + placement_num_experts = placement.get('num_experts') if isinstance(placement, dict) else None + if logical_shape and logical_shape[0] != checkpoint_num_experts and placement_num_experts == logical_shape[0]: + ranks = placement.get(AUTOEP_PLACEMENT_RANKS, []) if isinstance(placement, dict) else [] + if ranks and checkpoint_num_experts % len(ranks) == 0: + checkpoint_local_experts = checkpoint_num_experts // len(ranks) + target_local_experts = len(meta[AUTOEP_PARAM_LOCAL_EXPERTS]) + raise ValueError("AutoEP expert shape mismatch: " + f"target_local_experts={target_local_experts}, " + f"checkpoint_local_experts={checkpoint_local_experts}.") + raise ValueError("AutoEP universal checkpoint expert count disagrees with target logical shape: " + f"checkpoint={checkpoint_num_experts}, target={logical_shape}.") + if tuple(full_hp_param.shape) != logical_shape: + raise RuntimeError("AutoEP universal checkpoint tensor shape disagrees with target logical shape: " + f"checkpoint={tuple(full_hp_param.shape)}, target={logical_shape}.") + if not logical_shape or logical_shape[0] != checkpoint_num_experts: + raise RuntimeError("AutoEP universal checkpoint expert count disagrees with target logical shape: " + f"checkpoint={checkpoint_num_experts}, target={logical_shape}.") + + try: + target_map = autoep_placement_to_affine_map(meta[AUTOEP_EXPERT_PLACEMENT], logical_shape) + if not isinstance(ep_rank, int) or isinstance(ep_rank, bool): + raise ValueError(f"ep_rank must be an integer, got {ep_rank!r}.") + metadata_ep_rank = meta[AUTOEP_PARAM_EP_RANK] + if metadata_ep_rank != ep_rank: + raise ValueError(f"target ep_rank {ep_rank} does not match parameter metadata ep_rank {metadata_ep_rank}.") + entries_by_rank = { + entry[AUTOEP_PLACEMENT_RANK]: entry + for entry in meta[AUTOEP_EXPERT_PLACEMENT][AUTOEP_PLACEMENT_RANKS] + } + expected_local_experts = entries_by_rank[ep_rank][AUTOEP_PLACEMENT_EXPERTS] + if meta[AUTOEP_PARAM_LOCAL_EXPERTS] != expected_local_experts: + raise ValueError("parameter metadata local_experts does not match the selected placement rank: " + f"{meta[AUTOEP_PARAM_LOCAL_EXPERTS]} != {expected_local_experts}.") + return extract_autoep_rank_tensor(full_hp_param, target_map, ep_rank) + except (AssertionError, KeyError, TypeError, ValueError) as exc: + raise RuntimeError(f"Invalid AutoEP universal checkpoint target metadata: {exc}") from exc + + def _narrow_sub_params(full_view, partition_dim, sub_dim_sizes, shard_widths, tp_rank, tp_world_size, uc_version): """Take this rank's piece of every sub-parameter and concatenate them back together. @@ -173,7 +246,11 @@ def load_hp_checkpoint_state(self, folder, tp_rank, tp_world_size, ep_rank=0, ep # Must happen BEFORE shape-match check so that after slicing, # full_hp_param.shape == self.shape triggers tp_rank=0, tp_world_size=1. is_expert_param = ckpt_dict.get(EP_IS_EXPERT_PARAM, False) - if is_expert_param and ep_size > 1: + autoep_hp_param = (_resolve_autoep_partition(self, ckpt_dict, full_hp_param, ep_rank) + if is_expert_param else None) + if autoep_hp_param is not None: + full_hp_param = autoep_hp_param + elif is_expert_param and ep_size > 1: ep_num_experts = ckpt_dict.get(EP_NUM_EXPERTS) assert ep_num_experts is not None, \ f"Expert param in {ckpt_file} missing '{EP_NUM_EXPERTS}' metadata" diff --git a/deepspeed/module_inject/auto_ep_layer.py b/deepspeed/module_inject/auto_ep_layer.py index 4b8fcec70b88..78ab72e55e44 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -19,6 +19,10 @@ import torch import torch.nn as nn import deepspeed.comm as dist +from deepspeed.checkpoint.autoep_affine import legacy_uniform_autoep_placement_descriptor +from deepspeed.checkpoint.constants import (AUTOEP_EXPERT_PLACEMENT, AUTOEP_PARAM_EP_RANK, AUTOEP_PARAM_LOCAL_EXPERTS, + AUTOEP_PARAM_LOGICAL_SHAPE, AUTOEP_PLACEMENT_EXPERTS, + AUTOEP_PLACEMENT_RANKS, DS_AUTOEP_UC_META) from deepspeed.module_inject.auto_ep_config import AutoEPConfig, MoELayerSpec, resolve_autoep_config_defaults from deepspeed.module_inject.auto_ep_folding import mark_autoep_folding_router_parameter from deepspeed.ops.triton_ops import autoep_fused_token_ops as fused_token_ops @@ -415,6 +419,8 @@ def __init__( self.ep_rank = ep_rank self.num_experts = spec.num_experts self.num_local_experts = spec.num_experts // ep_size + self.expert_placement_descriptor = legacy_uniform_autoep_placement_descriptor( + self.num_experts, self.num_local_experts, self.ep_size) self.hidden_size = spec.hidden_size self.ep_group_name = f"ep_size_{ep_size}" self.ep_group = None # Set by set_deepspeed_parallelism() @@ -507,6 +513,21 @@ def __init__( self.experts.w1.requires_grad_(w1_requires_grad) self.experts.w2.requires_grad_(w2_requires_grad) self.experts.w3.requires_grad_(w3_requires_grad) + local_experts = self.expert_placement_descriptor[AUTOEP_PLACEMENT_RANKS][ + self.ep_rank][AUTOEP_PLACEMENT_EXPERTS] + 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:]] + setattr( + param, + DS_AUTOEP_UC_META, + { + AUTOEP_EXPERT_PLACEMENT: self.expert_placement_descriptor, + AUTOEP_PARAM_LOGICAL_SHAPE: logical_shape, + AUTOEP_PARAM_EP_RANK: self.ep_rank, + AUTOEP_PARAM_LOCAL_EXPERTS: list(local_experts), + }, + ) self.shared_experts = getattr(source_module, spec.shared_experts_name, None) if spec.has_shared_experts else None diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 06f302f5011f..b7d6dbfa267d 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -66,6 +66,9 @@ DATA_PARALLEL_GROUP, GLOBAL_RANK, DDP_BFLOAT16, GRADIENT_ALLREDUCE_OP_MEAN from deepspeed.runtime.zero.config import ZeroStageEnum from deepspeed.checkpoint.constants import ( + AUTOEP_EXPERT_PLACEMENT, + AUTOEP_PLACEMENT_EXPERTS, + AUTOEP_PLACEMENT_RANKS, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY, @@ -4272,24 +4275,9 @@ def load_moe_state_dict(checkpoint_path, else: # Validate AutoEP metadata if present if autoep_layers is not None: - if not isinstance(autoep_layers, list): - raise RuntimeError( - f"ds_autoep_layers metadata is malformed: expected list, got {type(autoep_layers).__name__}") - seen_ids = set() - required_fields = { - 'moe_layer_id', 'module_path', 'num_experts', 'num_local_experts', 'ep_size', 'expert_key_prefix' - } - for entry in autoep_layers: - if not isinstance(entry, dict): - raise RuntimeError( - f"ds_autoep_layers entry is malformed: expected dict, got {type(entry).__name__}") - missing = required_fields - entry.keys() - if missing: - raise RuntimeError(f"ds_autoep_layers entry is invalid: missing fields {sorted(missing)}") - lid = entry['moe_layer_id'] - if lid in seen_ids: - raise RuntimeError(f"ds_autoep_layers metadata has duplicate moe_layer_id: {lid}") - seen_ids.add(lid) + DeepSpeedEngine._validate_autoep_zero3_partitioned_metadata(autoep_layers, + model=model, + require_partitioned=False) elif has_autoep_layers: logger.warning("Checkpoint does not contain ds_autoep_layers metadata. " "Loading AutoEP expert weights using best-effort module detection.") @@ -4601,13 +4589,17 @@ def _uses_autoep_zero3_partitioned_experts(autoep_layers): return any(is_autoep_zero3_partitioned_entry(entry) for entry in autoep_layers) @staticmethod - def _validate_autoep_zero3_partitioned_metadata(autoep_layers, model=None, require_partitioned=True): + def _validate_autoep_zero3_partitioned_metadata(autoep_layers, + model=None, + require_partitioned=True, + validate_runtime_placement=False): try: from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer as _AutoEPMoELayer except ImportError: _AutoEPMoELayer = None expected_expert_prefixes = None + expected_runtime_layers = None if _AutoEPMoELayer is not None and model is not None: expected_expert_prefixes = { module_name: f"{module_name}.experts" if module_name else "experts" @@ -4615,10 +4607,23 @@ def _validate_autoep_zero3_partitioned_metadata(autoep_layers, model=None, requi } if not expected_expert_prefixes: expected_expert_prefixes = None + if validate_runtime_placement: + expected_runtime_layers = {} + for module_name, module in model.named_modules(): + if not isinstance(module, _AutoEPMoELayer): + continue + rank_entry = module.expert_placement_descriptor[AUTOEP_PLACEMENT_RANKS][module.ep_rank] + expected_runtime_layers[module_name] = { + 'ep_rank': module.ep_rank, + 'local_experts': list(rank_entry[AUTOEP_PLACEMENT_EXPERTS]), + } + if not expected_runtime_layers: + expected_runtime_layers = None validate_autoep_zero3_partitioned_metadata(autoep_layers, require_partitioned=require_partitioned, expected_expert_prefixes=expected_expert_prefixes, + expected_runtime_layers=expected_runtime_layers, version_context="This DeepSpeed build") @staticmethod @@ -5290,6 +5295,8 @@ def autoep_expert_writer() -> bool: num_local_experts, 'ep_size': module.ep_size, + AUTOEP_EXPERT_PLACEMENT: + module.expert_placement_descriptor, 'expert_key_prefix': f"{module_prefix}experts", AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY: @@ -5354,6 +5361,12 @@ def autoep_expert_writer() -> bool: moe_layer_id += 1 + if autoep_layer_info: + DeepSpeedEngine._validate_autoep_zero3_partitioned_metadata(autoep_layer_info, + model=self.module, + require_partitioned=False, + validate_runtime_placement=True) + self._curr_ckpt_path = os.path.join(save_dir, tag) largest_group_name = groups._get_max_expert_size_name() diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index 0afd8c1b9f89..3aa3d44ec166 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -3529,11 +3529,17 @@ def load_hp_checkpoint_state(self, folder, key, param=None): def _slice_autoep_universal_expert_param(self, checkpoint_state, param): full_expert_tensor = checkpoint_state[PARAM] - checkpoint_num_experts = checkpoint_state.get(EP_NUM_EXPERTS, full_expert_tensor.shape[0]) group_name = getattr(param, "ds_zero_partition_group_name", None) if group_name is None: raise ValueError("AutoEP universal expert checkpoint target parameter is missing its EP group name") ep_rank = groups._get_expert_parallel_rank(group_name) + + from deepspeed.checkpoint.universal_checkpoint import _resolve_autoep_partition + affine_partition = _resolve_autoep_partition(param, checkpoint_state, full_expert_tensor, ep_rank) + if affine_partition is not None: + return affine_partition + + checkpoint_num_experts = checkpoint_state.get(EP_NUM_EXPERTS, full_expert_tensor.shape[0]) ep_world_size = groups._get_expert_parallel_world_size(group_name) if checkpoint_num_experts % ep_world_size != 0: raise ValueError("AutoEP universal expert checkpoint tensor cannot be evenly split across the target " diff --git a/tests/unit/checkpoint/test_autoep_affine.py b/tests/unit/checkpoint/test_autoep_affine.py new file mode 100644 index 000000000000..57aa5a9251e3 --- /dev/null +++ b/tests/unit/checkpoint/test_autoep_affine.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team + +import json + +import pytest +import torch + +from deepspeed.checkpoint.affine import ParamAffineMap +from deepspeed.checkpoint.autoep_affine import (AUTOEP_PLACEMENT_VERSION, autoep_placement_to_affine_map, + extract_autoep_rank_tensor, legacy_uniform_autoep_placement_descriptor, + make_autoep_placement_descriptor, validate_autoep_placement_descriptor) + + +def _oracle_shards(full_tensor, experts_by_rank): + return { + rank: + torch.stack([full_tensor[expert_id] + for expert_id in experts]) if experts else full_tensor.new_empty((0, ) + + tuple(full_tensor.shape[1:])) + for rank, experts in enumerate(experts_by_rank) + } + + +@pytest.mark.parametrize('num_experts, ep_size', [(8, 2), (8, 4)]) +def test_uniform_placement_lowers_and_round_trips(num_experts, ep_size): + num_local_experts = num_experts // ep_size + descriptor = legacy_uniform_autoep_placement_descriptor(num_experts, num_local_experts, ep_size) + affine_map = autoep_placement_to_affine_map(descriptor, (num_experts, 3, 2)) + full_tensor = torch.arange(num_experts * 6, dtype=torch.float32).reshape(num_experts, 3, 2) + experts_by_rank = [ + list(range(rank * num_local_experts, (rank + 1) * num_local_experts)) for rank in range(ep_size) + ] + expected_shards = _oracle_shards(full_tensor, experts_by_rank) + + assert affine_map.shard_shapes == {rank: (num_local_experts, 3, 2) for rank in range(ep_size)} + assert all(len(pieces) == 1 for pieces in affine_map.pieces_by_rank.values()) + for rank in range(ep_size): + assert torch.equal(affine_map.extract(full_tensor, rank), expected_shards[rank]) + assert torch.equal(affine_map.rebuild(expected_shards), full_tensor) + + +def test_non_contiguous_non_uniform_placement_preserves_local_order(): + experts_by_rank = [[4, 1, 5], [0], [3, 2]] + descriptor = make_autoep_placement_descriptor(6, experts_by_rank) + affine_map = autoep_placement_to_affine_map(descriptor, (6, 2)) + full_tensor = torch.tensor([[10, 11], [20, 21], [30, 31], [40, 41], [50, 51], [60, 61]]) + expected_shards = _oracle_shards(full_tensor, experts_by_rank) + + assert affine_map.shard_shapes == {0: (3, 2), 1: (1, 2), 2: (2, 2)} + assert torch.equal(affine_map.extract(full_tensor, 0), torch.tensor([[50, 51], [20, 21], [60, 61]])) + assert torch.equal(affine_map.rebuild(expected_shards), full_tensor) + + +def test_replication_records_exact_holders_and_stops_piece_merging(): + experts_by_rank = [[0, 1, 2], [1, 2, 3]] + descriptor = make_autoep_placement_descriptor(4, experts_by_rank) + affine_map = autoep_placement_to_affine_map(descriptor, (4, 2)) + full_tensor = torch.arange(8).reshape(4, 2) + + assert [piece.locations + for piece in affine_map.pieces_by_rank[0]] == [frozenset({0}), + frozenset({0, 1}), + frozenset({0, 1})] + assert [piece.shape for piece in affine_map.pieces_by_rank[0]] == [(1, 2), (1, 2), (1, 2)] + assert [piece.locations + for piece in affine_map.pieces_by_rank[1]] == [frozenset({0, 1}), + frozenset({0, 1}), + frozenset({1})] + assert torch.equal(affine_map.rebuild(_oracle_shards(full_tensor, experts_by_rank)), full_tensor) + + +def test_replication_disagreement_with_different_local_order_fails(): + experts_by_rank = [[0, 1, 3, 4], [2, 4, 3]] + descriptor = make_autoep_placement_descriptor(5, experts_by_rank) + affine_map = autoep_placement_to_affine_map(descriptor, (5, 2)) + full_tensor = torch.arange(10, dtype=torch.float32).reshape(5, 2) + shards = _oracle_shards(full_tensor, experts_by_rank) + shards[1][1].add_(1) + + with pytest.raises(ValueError, match='different data'): + affine_map.rebuild(shards) + + +def test_empty_rank_shard_is_representable(): + descriptor = make_autoep_placement_descriptor(3, [[2, 0], [], [1]]) + affine_map = autoep_placement_to_affine_map(descriptor, (3, 4)) + full_tensor = torch.arange(12).reshape(3, 4) + + assert affine_map.shard_shapes[1] == (0, 4) + assert affine_map.pieces_by_rank[1] == [] + assert extract_autoep_rank_tensor(full_tensor, affine_map, 1).shape == (0, 4) + assert torch.equal(affine_map.rebuild(_oracle_shards(full_tensor, [[2, 0], [], [1]])), full_tensor) + shards_without_empty_rank = _oracle_shards(full_tensor, [[2, 0], [], [1]]) + del shards_without_empty_rank[1] + assert torch.equal(affine_map.rebuild(shards_without_empty_rank), full_tensor) + + +@pytest.mark.parametrize('descriptor, match', [ + ({ + 'version': AUTOEP_PLACEMENT_VERSION, + 'num_experts': 2, + 'ep_size': 2, + 'ranks': [{ + 'rank': 0, + 'experts': [0] + }], + }, 'exactly 2'), + ({ + 'version': AUTOEP_PLACEMENT_VERSION, + 'num_experts': 2, + 'ep_size': 2, + 'ranks': [{ + 'rank': 0, + 'experts': [0] + }, { + 'rank': 0, + 'experts': [1] + }], + }, 'rank 0 appears more than once'), + ({ + 'version': AUTOEP_PLACEMENT_VERSION, + 'num_experts': 2, + 'ep_size': 1, + 'ranks': [{ + 'rank': 0, + 'experts': [0, 0, 1] + }], + }, 'duplicate expert ID 0'), + ({ + 'version': AUTOEP_PLACEMENT_VERSION, + 'num_experts': 3, + 'ep_size': 2, + 'ranks': [{ + 'rank': 0, + 'experts': [0] + }, { + 'rank': 1, + 'experts': [2] + }], + }, 'does not cover global expert IDs \\[1\\]'), + ({ + 'version': AUTOEP_PLACEMENT_VERSION, + 'num_experts': 2, + 'ep_size': 1, + 'ranks': [{ + 'rank': 0, + 'experts': [0, 2] + }], + }, 'outside \\[0, 2\\)'), +]) +def test_malformed_placement_is_rejected(descriptor, match): + with pytest.raises(ValueError, match=match): + validate_autoep_placement_descriptor(descriptor) + + +@pytest.mark.parametrize('num_experts, num_local_experts, ep_size', [(8, 3, 2), (0, 1, 1), (4, 0, 2)]) +def test_inconsistent_legacy_metadata_is_rejected(num_experts, num_local_experts, ep_size): + with pytest.raises(ValueError): + legacy_uniform_autoep_placement_descriptor(num_experts, num_local_experts, ep_size) + + +def test_descriptor_and_affine_map_serialization_round_trip(): + descriptor = make_autoep_placement_descriptor(4, [[2, 0], [1, 2, 3]]) + serialized_descriptor = json.loads(json.dumps(descriptor)) + validate_autoep_placement_descriptor(serialized_descriptor) + affine_map = autoep_placement_to_affine_map(serialized_descriptor, (4, 3)) + restored_map = ParamAffineMap.from_dict(json.loads(json.dumps(affine_map.to_dict()))) + full_tensor = torch.randn(4, 3) + shards = _oracle_shards(full_tensor, [[2, 0], [1, 2, 3]]) + + assert serialized_descriptor == descriptor + assert restored_map.to_dict() == affine_map.to_dict() + assert torch.equal(restored_map.rebuild(shards), full_tensor) + + +def test_extract_target_rank_uses_target_packed_order(): + descriptor = make_autoep_placement_descriptor(5, [[0, 2], [4, 1, 3]]) + target_map = autoep_placement_to_affine_map(descriptor, (5, 2, 2)) + universal_tensor = torch.arange(20).reshape(5, 2, 2) + expected = torch.stack([universal_tensor[4], universal_tensor[1], universal_tensor[3]]) + + actual = extract_autoep_rank_tensor(universal_tensor, target_map, 1) + + assert torch.equal(actual, expected) diff --git a/tests/unit/checkpoint/test_autoep_affine_universal.py b/tests/unit/checkpoint/test_autoep_affine_universal.py new file mode 100644 index 000000000000..885998a9a030 --- /dev/null +++ b/tests/unit/checkpoint/test_autoep_affine_universal.py @@ -0,0 +1,494 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team + +import os +from types import SimpleNamespace + +import pytest +import torch + +from deepspeed.checkpoint.autoep_affine import make_autoep_placement_descriptor +from deepspeed.checkpoint.autoep_universal import (consolidate_autoep_zero12_expert_states, + get_autoep_zero12_expert_param_info) +from deepspeed.checkpoint.constants import ( + AUTOEP_EXPERT_PLACEMENT, AUTOEP_LAYERS_KEY, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY, + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY, + AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT, AUTOEP_PARAM_LOCAL_EXPERTS, AUTOEP_PLACEMENT_EXPERTS, + AUTOEP_PLACEMENT_RANK, AUTOEP_PLACEMENT_RANKS, DS_AUTOEP_UC_META, DS_AUTOTP_UC_META, EP_IS_EXPERT_PARAM, + EP_NUM_EXPERTS, OPTIMIZER_STATE_DICT, PARAM, PARAM_SHAPES) +from deepspeed.checkpoint.ds_to_universal import (_consolidate_zero3_autoep_expert_states, + _rebuild_zero3_autoep_rank_tensors) +from deepspeed.checkpoint.universal_checkpoint import _resolve_autoep_partition, load_hp_checkpoint_state +from deepspeed.runtime.zero import stage3 +from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 + + +def _write_zero12_fragments(temp_dir, param_name, rank_tensors): + for state_name in ('fp32', 'exp_avg', 'exp_avg_sq'): + for dp_rank, tensor in rank_tensors.items(): + path = os.path.join(temp_dir, param_name, "0", f"{state_name}.{dp_rank:02d}") + os.makedirs(os.path.dirname(path), exist_ok=True) + torch.save(tensor.flatten().float(), path) + for dp_rank in rank_tensors: + torch.save(7, os.path.join(temp_dir, param_name, "0", f"step.{dp_rank:02d}")) + + +def _load_consolidated(output_dir, param_name, state_name='fp32'): + state = torch.load(os.path.join(output_dir, "zero", param_name, f"{state_name}.pt"), + map_location='cpu', + weights_only=False) + assert state[EP_IS_EXPERT_PARAM] + return state + + +def test_zero12_affine_consolidation_supports_nonuniform_noncontiguous_placement(tmp_path): + param_name = "layer.experts.w1" + placement = make_autoep_placement_descriptor(4, [[2, 0, 3], [1]]) + full = torch.arange(8, dtype=torch.float32).view(4, 2) + rank_tensors = { + 0: torch.stack([full[2], full[0], full[3]]), + 1: full[1:2], + } + _write_zero12_fragments(str(tmp_path / "fragments"), param_name, rank_tensors) + info = { + param_name: { + 'num_experts': 4, + 'num_local_experts': 2, + 'ep_size': 2, + 'expert_placement': placement, + } + } + + consolidate_autoep_zero12_expert_states(str(tmp_path / "fragments"), str(tmp_path / "output"), info, + {param_name: (2, 2)}, 2, 1, False) + + for state_name in ('fp32', 'exp_avg', 'exp_avg_sq'): + state = _load_consolidated(str(tmp_path / "output"), param_name, state_name) + assert state[EP_NUM_EXPERTS] == 4 + assert torch.equal(state[PARAM], full) + + +def test_zero12_metadata_accepts_explicit_empty_rank(): + placement = make_autoep_placement_descriptor(2, [[], [0, 1]]) + metadata = [{ + 'moe_layer_id': 0, + 'module_path': 'layer', + 'num_experts': 2, + 'num_local_experts': 0, + 'ep_size': 2, + 'ep_rank': 0, + 'expert_key_prefix': 'layer.experts', + AUTOEP_EXPERT_PLACEMENT: placement, + }] + + param_info = get_autoep_zero12_expert_param_info(metadata) + + assert param_info['layer.experts.w1']['num_local_experts'] == 0 + + +def test_zero12_uniform_affine_consolidation_matches_legacy_order(tmp_path): + param_name = "layer.experts.w1" + placement = make_autoep_placement_descriptor(4, [[0, 1], [2, 3]]) + full = torch.arange(8, dtype=torch.float32).view(4, 2) + _write_zero12_fragments(str(tmp_path / "fragments"), param_name, {0: full[:2], 1: full[2:]}) + info = { + param_name: { + 'num_experts': 4, + 'num_local_experts': 2, + 'ep_size': 2, + 'expert_placement': placement, + } + } + + consolidate_autoep_zero12_expert_states(str(tmp_path / "fragments"), str(tmp_path / "output"), info, + {param_name: (2, 2)}, 2, 1, False) + + assert torch.equal(_load_consolidated(str(tmp_path / "output"), param_name)[PARAM], full) + + +def test_zero12_affine_consolidation_validates_replicated_experts(tmp_path): + param_name = "layer.experts.w2" + placement = make_autoep_placement_descriptor(3, [[0, 1], [1, 2]]) + full = torch.arange(6, dtype=torch.float32).view(3, 2) + rank_tensors = { + 0: torch.stack([full[0], full[1]]), + 1: torch.stack([full[1] + 1, full[2]]), + } + _write_zero12_fragments(str(tmp_path / "fragments"), param_name, rank_tensors) + info = { + param_name: { + 'num_experts': 3, + 'num_local_experts': 2, + 'ep_size': 2, + 'expert_placement': placement, + } + } + + with pytest.raises(ValueError, match="different data"): + consolidate_autoep_zero12_expert_states(str(tmp_path / "fragments"), str(tmp_path / "output"), info, + {param_name: (2, 2)}, 2, 1, False) + + +def test_zero12_legacy_consolidation_still_concatenates_ep_ranks(tmp_path): + param_name = "layer.experts.w3" + full = torch.arange(8, dtype=torch.float32).view(4, 2) + _write_zero12_fragments(str(tmp_path / "fragments"), param_name, {0: full[:2], 1: full[2:]}) + info = { + param_name: { + 'num_experts': 4, + 'num_local_experts': 2, + 'ep_size': 2, + 'expert_placement': None, + } + } + + consolidate_autoep_zero12_expert_states(str(tmp_path / "fragments"), str(tmp_path / "output"), info, + {param_name: (2, 2)}, 2, 1, False) + + assert torch.equal(_load_consolidated(str(tmp_path / "output"), param_name)[PARAM], full) + + +def test_zero3_affine_rebuild_checks_ranks_shapes_and_replicas(): + placement = make_autoep_placement_descriptor(3, [[2, 0], [1, 2]]) + full = torch.arange(6, dtype=torch.float32).view(3, 2) + shards = {0: torch.stack([full[2], full[0]]), 1: torch.stack([full[1], full[2]])} + + assert torch.equal(_rebuild_zero3_autoep_rank_tensors(shards, placement, full.shape, "test"), full) + with pytest.raises(RuntimeError, match="expected nonempty ranks \\[0, 1\\]"): + _rebuild_zero3_autoep_rank_tensors({0: shards[0]}, placement, full.shape, "test") + with pytest.raises(RuntimeError, match="different data"): + _rebuild_zero3_autoep_rank_tensors({ + 0: shards[0], + 1: shards[1] + torch.tensor([[0, 0], [1, 0]]) + }, placement, full.shape, "test") + + +def test_zero3_affine_rebuild_allows_omitted_empty_rank(): + placement = make_autoep_placement_descriptor(3, [[2, 0, 1], []]) + full = torch.arange(6, dtype=torch.float32).view(3, 2) + shards = {0: torch.stack([full[2], full[0], full[1]])} + + assert torch.equal(_rebuild_zero3_autoep_rank_tensors(shards, placement, full.shape, "test"), full) + + +def test_zero3_partition_native_consolidation_uses_affine_placement(tmp_path): + param_names = [f"layer.experts.w{index}" for index in range(1, 4)] + placement = make_autoep_placement_descriptor(4, [[2, 0, 3], [1]]) + full = torch.arange(8, dtype=torch.float32).view(4, 2) + rank_tensors = { + 0: torch.stack([full[2], full[0], full[3]]), + 1: full[1:2], + } + model_files = [] + optim_files = [] + for rank, local_tensor in rank_tensors.items(): + layer_info = { + 'moe_layer_id': 0, + 'module_path': 'layer', + 'num_experts': 4, + 'num_local_experts': local_tensor.shape[0], + 'ep_size': 2, + 'expert_key_prefix': 'layer.experts', + AUTOEP_EXPERT_PLACEMENT: placement, + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY: AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT, + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY: AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION, + 'ep_group_name': 'ep_size_2', + 'ep_rank': rank, + 'expert_data_parallel_rank': 0, + 'expert_data_parallel_world_size': 1, + 'global_expert_start': 0, + 'global_expert_end': 0, + } + model_file = tmp_path / f"zero_pp_rank_{rank}_mp_rank_00_model_states.pt" + torch.save( + { + AUTOEP_LAYERS_KEY: [layer_info], + PARAM_SHAPES: [{ + param_name: local_tensor.shape + } for param_name in param_names], + }, model_file) + model_files.append(str(model_file)) + + flat = local_tensor.flatten() + optim_file = tmp_path / f"bf16_zero_pp_rank_{rank}_mp_rank_00_optim_states.pt" + torch.save( + { + OPTIMIZER_STATE_DICT: { + 'ds_zero_partition_groups': [{ + 'partition_count': 1, + 'partition_rank': 0, + } for _ in param_names], + 'optimizer_state_dict': { + 'state': { + index: { + 'exp_avg': flat.clone(), + 'exp_avg_sq': flat.clone(), + } + for index in range(len(param_names)) + } + }, + 'fp32_flat_groups': [flat.clone() for _ in param_names], + } + }, optim_file) + optim_files.append(str(optim_file)) + + output_dir = tmp_path / "universal" + _consolidate_zero3_autoep_expert_states(str(output_dir), model_files, optim_files) + + for param_name in param_names: + for state_name in ('fp32', 'exp_avg', 'exp_avg_sq'): + state = _load_consolidated(str(output_dir), param_name, state_name) + assert torch.equal(state[PARAM], full) + + +def test_zero3_partition_native_consolidation_allows_rank_without_fragments(tmp_path): + param_names = [f"layer.experts.w{index}" for index in range(1, 4)] + placement = make_autoep_placement_descriptor(3, [[2, 0, 1], []]) + full = torch.arange(6, dtype=torch.float32).view(3, 2) + model_files = [] + optim_files = [] + for rank in range(2): + local_tensor = torch.stack([full[2], full[0], full[1]]) if rank == 0 else None + layer_info = { + 'moe_layer_id': 0, + 'module_path': 'layer', + 'num_experts': 3, + 'num_local_experts': 3 if rank == 0 else 0, + 'ep_size': 2, + 'expert_key_prefix': 'layer.experts', + AUTOEP_EXPERT_PLACEMENT: placement, + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY: AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT, + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY: AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION, + 'ep_group_name': 'ep_size_2', + 'ep_rank': rank, + 'expert_data_parallel_rank': 0, + 'expert_data_parallel_world_size': 1, + 'global_expert_start': 0, + 'global_expert_end': 0, + } + model_file = tmp_path / f"zero_pp_rank_{rank}_mp_rank_00_model_states.pt" + param_shapes = [] if local_tensor is None else [{param_name: local_tensor.shape} for param_name in param_names] + torch.save({ + AUTOEP_LAYERS_KEY: [layer_info], + PARAM_SHAPES: param_shapes, + }, model_file) + model_files.append(str(model_file)) + + optim_file = tmp_path / f"bf16_zero_pp_rank_{rank}_mp_rank_00_optim_states.pt" + if local_tensor is None: + optimizer_state = { + 'ds_zero_partition_groups': [], + 'optimizer_state_dict': { + 'state': {} + }, + 'fp32_flat_groups': [], + } + else: + flat = local_tensor.flatten() + optimizer_state = { + 'ds_zero_partition_groups': [{ + 'partition_count': 1, + 'partition_rank': 0, + } for _ in param_names], + 'optimizer_state_dict': { + 'state': { + index: { + 'exp_avg': flat.clone(), + 'exp_avg_sq': flat.clone(), + } + for index in range(len(param_names)) + } + }, + 'fp32_flat_groups': [flat.clone() for _ in param_names], + } + torch.save({OPTIMIZER_STATE_DICT: optimizer_state}, optim_file) + optim_files.append(str(optim_file)) + + output_dir = tmp_path / "universal" + _consolidate_zero3_autoep_expert_states(str(output_dir), model_files, optim_files) + + for param_name in param_names: + for state_name in ('fp32', 'exp_avg', 'exp_avg_sq'): + state = _load_consolidated(str(output_dir), param_name, state_name) + assert torch.equal(state[PARAM], full) + + +def _target_param(placement, logical_shape, ep_rank): + entries_by_rank = {entry[AUTOEP_PLACEMENT_RANK]: entry for entry in placement[AUTOEP_PLACEMENT_RANKS]} + param = torch.nn.Parameter(torch.empty(1)) + setattr( + param, + DS_AUTOEP_UC_META, + { + AUTOEP_EXPERT_PLACEMENT: placement, + 'logical_shape': list(logical_shape), + 'ep_rank': ep_rank, + AUTOEP_PARAM_LOCAL_EXPERTS: list(entries_by_rank[ep_rank][AUTOEP_PLACEMENT_EXPERTS]), + }, + ) + return param + + +def test_autoep_target_extraction_precedes_tp_and_stage3_partitioning(monkeypatch): + placement = make_autoep_placement_descriptor(5, [[0, 3], [4, 1, 2]]) + full = torch.arange(20, dtype=torch.float32).view(5, 2, 2) + checkpoint_state = { + PARAM: full, + EP_IS_EXPERT_PARAM: True, + EP_NUM_EXPERTS: 5, + } + param = _target_param(placement, full.shape, 1) + expected = torch.stack([full[4], full[1], full[2]]) + + param.ds_zero_partition_group_name = "ep" + monkeypatch.setattr(stage3.groups, "_get_expert_parallel_rank", lambda _: 1) + + assert torch.equal(_resolve_autoep_partition(param, checkpoint_state, full, 1), expected) + assert torch.equal( + DeepSpeedZeroOptimizer_Stage3._slice_autoep_universal_expert_param(None, checkpoint_state, param), expected) + + +def test_zero12_restore_loads_target_packed_order(tmp_path): + placement = make_autoep_placement_descriptor(5, [[0, 3], [4, 1, 2]]) + full = torch.arange(10, dtype=torch.float32).view(5, 2) + expected = torch.stack([full[4], full[1], full[2]]) + checkpoint_dir = tmp_path / "checkpoint" + checkpoint_dir.mkdir() + torch.save({ + PARAM: full, + EP_IS_EXPERT_PARAM: True, + EP_NUM_EXPERTS: 5, + }, checkpoint_dir / "fp32.pt") + + param = torch.nn.Parameter(torch.empty_like(expected)) + setattr( + param, + DS_AUTOEP_UC_META, + { + AUTOEP_EXPERT_PLACEMENT: placement, + 'logical_shape': list(full.shape), + 'ep_rank': 1, + AUTOEP_PARAM_LOCAL_EXPERTS: [4, 1, 2], + }, + ) + destination = torch.empty(expected.numel(), dtype=torch.float32) + param._hp_mapping = SimpleNamespace( + optim_fragment={}, + lp_fragment_address=SimpleNamespace(start=0, numel=expected.numel()), + get_hp_fragment=lambda: destination, + ) + + load_hp_checkpoint_state(param, str(checkpoint_dir), tp_rank=0, tp_world_size=1, ep_rank=1, ep_size=2) + + assert torch.equal(destination.view_as(expected), expected) + + +def test_zero12_restore_applies_autoep_before_independent_autotp(tmp_path): + placement = make_autoep_placement_descriptor(4, [[2, 0], [3, 1]]) + full = torch.arange(16, dtype=torch.float32).view(4, 4) + ep_local = torch.stack([full[2], full[0]]) + checkpoint_dir = tmp_path / "checkpoint" + checkpoint_dir.mkdir() + torch.save({ + PARAM: full, + EP_IS_EXPERT_PARAM: True, + EP_NUM_EXPERTS: 4, + }, checkpoint_dir / "fp32.pt") + + param = _target_param(placement, full.shape, 0) + setattr( + param, + DS_AUTOTP_UC_META, + { + 'logical_shape': list(ep_local.shape), + 'partition_dim': 1, + 'partition_sizes': [2, 2], + }, + ) + destination = torch.empty(4, dtype=torch.float32) + param._hp_mapping = SimpleNamespace( + optim_fragment={}, + lp_fragment_address=SimpleNamespace(start=0, numel=4), + get_hp_fragment=lambda: destination, + ) + + load_hp_checkpoint_state(param, str(checkpoint_dir), tp_rank=1, tp_world_size=2, ep_rank=0, ep_size=2) + + assert torch.equal(destination.view(2, 2), ep_local[:, 2:]) + + +def test_zero3_restore_extracts_ep_rank_before_zero_partition(tmp_path, monkeypatch): + placement = make_autoep_placement_descriptor(5, [[0, 3], [4, 1, 2]]) + full = torch.arange(10, dtype=torch.float32).view(5, 2) + expected_local = torch.stack([full[4], full[1], full[2]]) + checkpoint_dir = tmp_path / "checkpoint" + checkpoint_dir.mkdir() + torch.save({ + PARAM: full, + EP_IS_EXPERT_PARAM: True, + EP_NUM_EXPERTS: 5, + }, checkpoint_dir / "fp32.pt") + param = _target_param(placement, full.shape, 1) + param.ds_zero_partition_group_name = "ep" + partition_group = object() + optimizer = SimpleNamespace( + _get_param_partition_group=lambda _: partition_group, + _slice_autoep_universal_expert_param=lambda checkpoint_state, target_param: DeepSpeedZeroOptimizer_Stage3. + _slice_autoep_universal_expert_param(None, checkpoint_state, target_param), + ) + monkeypatch.setattr(stage3.dist, "get_rank", lambda group=None: 1) + monkeypatch.setattr(stage3.dist, "get_world_size", lambda group=None: 2) + monkeypatch.setattr(stage3.groups, "_get_expert_parallel_rank", lambda _: 1) + + actual = DeepSpeedZeroOptimizer_Stage3.load_hp_checkpoint_state(optimizer, str(checkpoint_dir), "fp32", param) + + assert torch.equal(actual, expected_local.flatten()[3:]) + + +@pytest.mark.parametrize( + 'metadata, match', + [ + ({}, "missing fields"), + ({ + AUTOEP_EXPERT_PLACEMENT: make_autoep_placement_descriptor(2, [[0], [1]]), + 'logical_shape': [3, 2], + 'ep_rank': 0, + AUTOEP_PARAM_LOCAL_EXPERTS: [0], + }, "tensor shape disagrees"), + ({ + AUTOEP_EXPERT_PLACEMENT: make_autoep_placement_descriptor(2, [[0], [1]]), + 'logical_shape': [2, 2], + 'ep_rank': 2, + AUTOEP_PARAM_LOCAL_EXPERTS: [], + }, "Invalid AutoEP"), + ], +) +def test_autoep_target_metadata_fails_loudly(metadata, match): + full = torch.arange(4, dtype=torch.float32).view(2, 2) + checkpoint_state = {PARAM: full, EP_IS_EXPERT_PARAM: True, EP_NUM_EXPERTS: 2} + param = torch.nn.Parameter(torch.empty(1)) + setattr(param, DS_AUTOEP_UC_META, metadata) + + with pytest.raises(RuntimeError, match=match): + _resolve_autoep_partition(param, checkpoint_state, full, metadata.get('ep_rank', 0)) + + +def test_autoep_target_metadata_rejects_caller_ep_rank_mismatch(): + placement = make_autoep_placement_descriptor(2, [[0], [1]]) + full = torch.arange(4, dtype=torch.float32).view(2, 2) + checkpoint_state = {PARAM: full, EP_IS_EXPERT_PARAM: True, EP_NUM_EXPERTS: 2} + param = _target_param(placement, full.shape, 0) + + with pytest.raises(RuntimeError, match="target ep_rank 1 does not match parameter metadata ep_rank 0"): + _resolve_autoep_partition(param, checkpoint_state, full, 1) + + +def test_autoep_target_metadata_rejects_local_expert_list_mismatch(): + placement = make_autoep_placement_descriptor(3, [[2, 0], [1]]) + full = torch.arange(6, dtype=torch.float32).view(3, 2) + checkpoint_state = {PARAM: full, EP_IS_EXPERT_PARAM: True, EP_NUM_EXPERTS: 3} + param = _target_param(placement, full.shape, 0) + getattr(param, DS_AUTOEP_UC_META)[AUTOEP_PARAM_LOCAL_EXPERTS] = [0, 2] + + with pytest.raises(RuntimeError, match="local_experts does not match"): + _resolve_autoep_partition(param, checkpoint_state, full, 0) diff --git a/tests/unit/checkpoint/test_autoep_zero3_metadata.py b/tests/unit/checkpoint/test_autoep_zero3_metadata.py new file mode 100644 index 000000000000..38b4d1dab341 --- /dev/null +++ b/tests/unit/checkpoint/test_autoep_zero3_metadata.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Contract tests for AutoEP checkpoint placement metadata.""" + +import copy + +import pytest + +from deepspeed.checkpoint.autoep_affine import make_autoep_placement_descriptor +from deepspeed.checkpoint.autoep_zero3_metadata import validate_autoep_zero3_partitioned_metadata +from deepspeed.checkpoint.constants import (AUTOEP_EXPERT_PLACEMENT, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY, + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION, + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY, + AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT) + + +def _metadata_entry(): + return { + "moe_layer_id": 0, + "module_path": "model.layers.0.mlp", + "num_experts": 4, + "num_local_experts": 2, + "ep_size": 2, + "expert_key_prefix": "model.layers.0.mlp.experts", + "ep_rank": 1, + } + + +def _partitioned_metadata_entry(): + entry = _metadata_entry() + entry.update({ + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY: AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT, + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY: AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION, + "ep_group_name": "ep_size_2", + "expert_data_parallel_rank": 0, + "expert_data_parallel_world_size": 1, + "global_expert_start": 2, + "global_expert_end": 4, + }) + return entry + + +def test_legacy_metadata_synthesizes_uniform_placement(): + placements = validate_autoep_zero3_partitioned_metadata([_metadata_entry()], require_partitioned=False) + + assert placements == [{ + "version": 1, + "num_experts": 4, + "ep_size": 2, + "ranks": [{ + "rank": 0, + "experts": [0, 1] + }, { + "rank": 1, + "experts": [2, 3] + }], + }] + + +@pytest.mark.parametrize( + "placement", + [ + { + "version": 999, + "num_experts": 4, + "ep_size": 2, + "ranks": [], + }, + make_autoep_placement_descriptor(6, [[0, 1, 2], [3, 4, 5]]), + make_autoep_placement_descriptor(4, [[0, 1], [2, 3], []]), + ], +) +def test_new_invalid_or_mismatched_placement_is_rejected(placement): + entry = _metadata_entry() + entry[AUTOEP_EXPERT_PLACEMENT] = placement + + with pytest.raises(RuntimeError, match="placement metadata"): + validate_autoep_zero3_partitioned_metadata([entry], require_partitioned=False) + + +def test_rank_local_placement_must_match_runtime_packing(): + entry = _metadata_entry() + entry[AUTOEP_EXPERT_PLACEMENT] = make_autoep_placement_descriptor(4, [[0, 2], [1, 3]]) + expected_runtime_layers = { + "model.layers.0.mlp": { + "ep_rank": 1, + "local_experts": [2, 3], + } + } + + with pytest.raises(RuntimeError, match="expert order"): + validate_autoep_zero3_partitioned_metadata([entry], + require_partitioned=False, + expected_runtime_layers=expected_runtime_layers) + + +def test_rank_local_expert_count_must_match_layer_entry(): + entry = copy.deepcopy(_metadata_entry()) + entry["num_local_experts"] = 1 + entry[AUTOEP_EXPERT_PLACEMENT] = make_autoep_placement_descriptor(4, [[0, 1], [2, 3]]) + + with pytest.raises(RuntimeError, match="expert count"): + validate_autoep_zero3_partitioned_metadata([entry], require_partitioned=False) + + +def test_partitioned_metadata_with_explicit_uneven_noncontiguous_placement_is_accepted(): + entry = _partitioned_metadata_entry() + entry["num_experts"] = 5 + entry["num_local_experts"] = 3 + entry["global_expert_start"] = 99 + entry["global_expert_end"] = 102 + entry[AUTOEP_EXPERT_PLACEMENT] = make_autoep_placement_descriptor(5, [[1, 3], [4, 0, 2]]) + + placements = validate_autoep_zero3_partitioned_metadata([entry]) + + assert placements == [entry[AUTOEP_EXPERT_PLACEMENT]] + + +def test_partitioned_legacy_metadata_still_requires_uniform_contiguous_ranges(): + entry = _partitioned_metadata_entry() + entry["global_expert_start"] = 1 + + with pytest.raises(RuntimeError, match="global expert range"): + validate_autoep_zero3_partitioned_metadata([entry]) diff --git a/tests/unit/v1/moe/test_autoep_unit.py b/tests/unit/v1/moe/test_autoep_unit.py index e0e15c3db6ad..63f3b3f8db26 100644 --- a/tests/unit/v1/moe/test_autoep_unit.py +++ b/tests/unit/v1/moe/test_autoep_unit.py @@ -520,6 +520,92 @@ def test_autoep_layer_marks_zero3_param_placement_families(self): for param in autoep_layer.router.parameters(): assert param.ds_zero_placement_family == "replicated" + def test_autoep_layer_emits_uniform_placement_and_param_restore_metadata(self): + from deepspeed.checkpoint.constants import ( + AUTOEP_EXPERT_PLACEMENT, + AUTOEP_PARAM_EP_RANK, + AUTOEP_PARAM_LOCAL_EXPERTS, + AUTOEP_PARAM_LOGICAL_SHAPE, + DS_AUTOEP_UC_META, + ) + + source = MockMoEBlock(num_experts=4, ffn_hidden=128, hidden_size=64) + layer = AutoEPMoELayer(_make_spec(), + source, + ep_size=2, + ep_rank=1, + config=_runtime_config(enabled=True, autoep_size=2)) + peer_layer = AutoEPMoELayer(_make_spec(), + source, + ep_size=2, + ep_rank=0, + config=_runtime_config(enabled=True, autoep_size=2)) + expected_placement = { + "version": 1, + "num_experts": 4, + "ep_size": 2, + "ranks": [{ + "rank": 0, + "experts": [0, 1] + }, { + "rank": 1, + "experts": [2, 3] + }], + } + + assert layer.expert_placement_descriptor == expected_placement + assert peer_layer.expert_placement_descriptor == expected_placement + for param in layer.experts.parameters(): + restore_metadata = getattr(param, DS_AUTOEP_UC_META) + assert restore_metadata[AUTOEP_EXPERT_PLACEMENT] == expected_placement + assert restore_metadata[AUTOEP_PARAM_LOGICAL_SHAPE] == [4, *param.shape[1:]] + assert restore_metadata[AUTOEP_PARAM_EP_RANK] == 1 + assert restore_metadata[AUTOEP_PARAM_LOCAL_EXPERTS] == [2, 3] + + for param in layer.router.parameters(): + assert not hasattr(param, DS_AUTOEP_UC_META) + + def test_checkpoint_source_placement_does_not_need_to_match_target_topology(self): + from deepspeed.checkpoint.autoep_affine import make_autoep_placement_descriptor + from deepspeed.checkpoint.constants import ( + AUTOEP_EXPERT_PLACEMENT, + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY, + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION, + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY, + AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT, + ) + + source = MockMoEBlock(num_experts=4, ffn_hidden=128, hidden_size=64) + target_layer = AutoEPMoELayer(_make_spec(), + source, + ep_size=4, + ep_rank=0, + config=_runtime_config(enabled=True, autoep_size=4)) + target_model = nn.Sequential(target_layer) + source_metadata = [{ + "moe_layer_id": 0, + "module_path": "0", + "num_experts": 4, + "num_local_experts": 2, + "ep_size": 2, + AUTOEP_EXPERT_PLACEMENT: make_autoep_placement_descriptor(4, [[0, 1], [2, 3]]), + "expert_key_prefix": "0.experts", + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY: AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT, + AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY: AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION, + "ep_group_name": "ep_size_2", + "ep_rank": 0, + "expert_data_parallel_rank": 0, + "expert_data_parallel_world_size": 1, + "global_expert_start": 0, + "global_expert_end": 2, + }] + + DeepSpeedEngine._validate_autoep_zero3_partitioned_metadata(source_metadata, model=target_model) + with pytest.raises(RuntimeError, match="expert order"): + DeepSpeedEngine._validate_autoep_zero3_partitioned_metadata(source_metadata, + model=target_model, + validate_runtime_placement=True) + def test_zero3_checkpoint_metadata_includes_partition_group_ranks(self): optimizer = object.__new__(DeepSpeedZeroOptimizer_Stage3) param = nn.Parameter(torch.empty(1)) From 2ef509386b8abee44b928d75bc8b54fbd58299d6 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Tue, 22 Sep 2026 09:10:38 +0000 Subject: [PATCH 2/6] Align AutoEP checkpoint maps with AffineIR Persist versioned per-parameter affine maps in AutoEP checkpoint metadata and use them for conversion and restore, while retaining placement descriptors for provenance and legacy checkpoints. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- deepspeed/checkpoint/affine_ir_spec.md | 16 +++--- deepspeed/checkpoint/autoep_affine.py | 32 ++++++++++- deepspeed/checkpoint/autoep_universal.py | 20 ++++++- deepspeed/checkpoint/constants.py | 1 + deepspeed/checkpoint/ds_to_universal.py | 60 +++++++++++++------- deepspeed/checkpoint/universal_checkpoint.py | 4 +- deepspeed/module_inject/auto_ep_layer.py | 11 ++-- deepspeed/runtime/engine.py | 17 ++++++ tests/unit/checkpoint/test_autoep_affine.py | 43 +++++++++++++- tests/unit/v1/moe/test_autoep_unit.py | 2 + 10 files changed, 166 insertions(+), 40 deletions(-) diff --git a/deepspeed/checkpoint/affine_ir_spec.md b/deepspeed/checkpoint/affine_ir_spec.md index 993163d6cf15..133a015d2073 100644 --- a/deepspeed/checkpoint/affine_ir_spec.md +++ b/deepspeed/checkpoint/affine_ir_spec.md @@ -568,13 +568,15 @@ 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 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. +**8.4 AutoEP, ZeRO, and offload placement.** AutoEP persists a versioned +`ParamAffineMap` for each expert parameter in addition to a versioned placement +descriptor in EP-local rank coordinates. The persisted map is the geometry execution +contract for conversion and restore; descriptor lowering is retained only as a legacy +checkpoint fallback. The descriptor records each rank's ordered global expert IDs, +including uneven, non-contiguous, replicated, and empty placements, so it remains useful +for placement provenance and validation rather than 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. diff --git a/deepspeed/checkpoint/autoep_affine.py b/deepspeed/checkpoint/autoep_affine.py index cd615c8c43de..2fa48767acf4 100644 --- a/deepspeed/checkpoint/autoep_affine.py +++ b/deepspeed/checkpoint/autoep_affine.py @@ -7,8 +7,9 @@ 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, +from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION, AffinePiece, ParamAffineMap +from deepspeed.checkpoint.constants import (AFFINE_MAP, AFFINE_MAP_PARAMS, AFFINE_MAP_VERSION, AUTOEP_AFFINE_MAPS, + AUTOEP_PLACEMENT_EP_SIZE, AUTOEP_PLACEMENT_EXPERTS, AUTOEP_PLACEMENT_NUM_EXPERTS, AUTOEP_PLACEMENT_RANK, AUTOEP_PLACEMENT_RANKS, AUTOEP_PLACEMENT_VERSION, AUTOEP_PLACEMENT_VERSION_KEY) @@ -19,6 +20,7 @@ 'validate_autoep_placement_descriptor', 'legacy_uniform_autoep_placement_descriptor', 'autoep_placement_to_affine_map', + 'autoep_metadata_to_affine_map', 'extract_autoep_rank_tensor', ] @@ -136,6 +138,32 @@ def autoep_placement_to_affine_map(descriptor, logical_shape): return affine_map +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) + if serialized_map is None and param_name is not None: + affine_maps = metadata.get(AUTOEP_AFFINE_MAPS, {}) + if affine_maps: + if affine_maps.get(AFFINE_MAP_VERSION) != AFFINE_MAP_FORMAT_VERSION: + raise ValueError(f"Unsupported AutoEP affine map format version " + f"{affine_maps.get(AFFINE_MAP_VERSION)!r}.") + serialized_map = affine_maps.get(AFFINE_MAP_PARAMS, {}).get(param_name) + + if serialized_map is not None: + affine_map = ParamAffineMap.from_dict(serialized_map) + expected_shape = tuple(int(dim) for dim in logical_shape) + if affine_map.logical_shape != expected_shape: + raise ValueError(f"AutoEP affine map logical shape {affine_map.logical_shape} does not match " + f"parameter logical shape {expected_shape}.") + affine_map.validate() + return affine_map + + placement = metadata.get('expert_placement') + if placement is None: + raise ValueError("AutoEP metadata has neither an affine map nor an expert placement descriptor.") + return autoep_placement_to_affine_map(placement, logical_shape) + + 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): diff --git a/deepspeed/checkpoint/autoep_universal.py b/deepspeed/checkpoint/autoep_universal.py index c4d82b428f3a..fd386bcf908b 100644 --- a/deepspeed/checkpoint/autoep_universal.py +++ b/deepspeed/checkpoint/autoep_universal.py @@ -13,8 +13,10 @@ import torch from deepspeed.utils import logger +from .affine import AFFINE_MAP_FORMAT_VERSION from .constants import ( + AUTOEP_AFFINE_MAPS, AUTOEP_EP_SIZE, AUTOEP_EXPERT_PLACEMENT, AUTOEP_EXPERT_KEY_PREFIX, @@ -26,6 +28,8 @@ AUTOEP_PLACEMENT_RANK, AUTOEP_PLACEMENT_RANKS, AUTOEP_ZERO12_REQUIRED_FIELDS, + AFFINE_MAP_PARAMS, + AFFINE_MAP_VERSION, PARAM, CAT_DIM, EP_IS_EXPERT_PARAM, @@ -46,7 +50,7 @@ FOLDING_FAMILY, FOLDING_PARAM_FAMILIES, ) -from .autoep_affine import autoep_placement_to_affine_map, validate_autoep_placement_descriptor +from .autoep_affine import autoep_metadata_to_affine_map, validate_autoep_placement_descriptor def make_folding_metadata(*, @@ -284,7 +288,17 @@ def get_autoep_zero12_expert_param_info(autoep_layers_metadata): param_name = f"{prefix}.{weight_name}" if param_name in param_info: raise RuntimeError(f"Duplicate AutoEP expert parameter metadata for {param_name}.") - param_info[param_name] = normalized + param_metadata = dict(normalized) + affine_maps = layer_info.get(AUTOEP_AFFINE_MAPS, {}) + if affine_maps: + try: + if affine_maps.get(AFFINE_MAP_VERSION) != AFFINE_MAP_FORMAT_VERSION: + raise RuntimeError(f"Unsupported AutoEP affine map format version " + f"{affine_maps.get(AFFINE_MAP_VERSION)!r}.") + param_metadata['affine_map'] = affine_maps[AFFINE_MAP_PARAMS][param_name] + except KeyError as exc: + raise RuntimeError(f"AutoEP layer metadata is missing affine map for {param_name}.") from exc + param_info[param_name] = param_metadata return param_info @@ -331,7 +345,7 @@ def consolidate_autoep_zero12_expert_states(temp_dir, output_dir, expert_param_i 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) + affine_map = autoep_metadata_to_affine_map(metadata, 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={metadata['num_local_experts']}.") diff --git a/deepspeed/checkpoint/constants.py b/deepspeed/checkpoint/constants.py index 230f22fabc1e..9285be03842b 100644 --- a/deepspeed/checkpoint/constants.py +++ b/deepspeed/checkpoint/constants.py @@ -114,6 +114,7 @@ AUTOEP_NUM_LOCAL_EXPERTS = 'num_local_experts' AUTOEP_EP_SIZE = 'ep_size' AUTOEP_EXPERT_PLACEMENT = 'expert_placement' +AUTOEP_AFFINE_MAPS = 'affine_maps' DS_AUTOEP_UC_META = 'ds_autoep_uc_meta' AUTOEP_PARAM_LOGICAL_SHAPE = 'logical_shape' AUTOEP_PARAM_EP_RANK = 'ep_rank' diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py index 08e90aa09700..825705d17101 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/deepspeed/checkpoint/ds_to_universal.py @@ -65,7 +65,7 @@ validate_autoep_zero3_partitioned_metadata, ) from deepspeed.checkpoint.affine import ParamAffineMap, AFFINE_MAP_FORMAT_VERSION -from deepspeed.checkpoint.autoep_affine import autoep_placement_to_affine_map, validate_autoep_placement_descriptor +from deepspeed.checkpoint.autoep_affine import autoep_metadata_to_affine_map, validate_autoep_placement_descriptor def parse_arguments(): @@ -496,6 +496,11 @@ def consume_superseded_patterns(name_): param = torch.cat(merged_chunks, dim=partition_dim) ckpt_dict[SUB_PARAM_SHAPE] = normalized_subparam_shape else: + if not slices: + # ZeRO does not materialize optimizer state for parameters that have + # never received a gradient. Their model state is copied separately + # below, so there is no optimizer tensor to merge here. + continue cat_dim = 1 if get_matched_pattern(parameters_with_row_parallelism, name) else 0 # print(f"merge {name} with CAT DIM: {cat_dim}") param = torch.cat(slices, dim=cat_dim) @@ -783,11 +788,17 @@ def _save_zero3_autoep_universal_tensor(output_dir, param_name, state_key, tenso ) -def _rebuild_zero3_autoep_rank_tensors(ep_tensors, placement, logical_shape, context): - if placement is None: +def _rebuild_zero3_autoep_rank_tensors(ep_tensors, affine_map_or_placement, logical_shape, context): + if affine_map_or_placement is None: return torch.cat([ep_tensors[rank] for rank in sorted(ep_tensors)], dim=0) - affine_map = autoep_placement_to_affine_map(placement, logical_shape) + if hasattr(affine_map_or_placement, 'shard_shapes'): + affine_map = affine_map_or_placement + else: + affine_map = autoep_metadata_to_affine_map( + {AUTOEP_EXPERT_PLACEMENT: affine_map_or_placement}, + logical_shape, + ) expected_ranks = set(affine_map.shard_shapes) actual_ranks = set(ep_tensors) unexpected_ranks = actual_ranks - expected_ranks @@ -833,6 +844,7 @@ def _consolidate_zero3_autoep_expert_states(output_dir, model_files, optim_files expected_ep_ranks_by_param = {} placement_by_param = {} logical_shape_by_param = {} + affine_map_by_param = {} for rank, model_state in model_states_by_rank.items(): optim_state = optim_states_by_rank.get(rank) @@ -884,7 +896,10 @@ def _consolidate_zero3_autoep_expert_states(output_dir, model_files, optim_files if existing_logical_shape != logical_shape: raise RuntimeError(f"AutoEP expert logical shape disagrees across ranks for {param_name}: " f"{existing_logical_shape} != {logical_shape}.") - affine_map = autoep_placement_to_affine_map(placement, logical_shape) + affine_map = autoep_metadata_to_affine_map(layer_info, logical_shape) + existing_map = affine_map_by_param.setdefault(param_name, affine_map) + if existing_map.to_dict() != affine_map.to_dict(): + raise RuntimeError(f"AutoEP affine map disagrees across ranks for {param_name}.") expected_shape = affine_map.shard_shapes.get(ep_rank) if expected_shape is None or tuple(shape) != expected_shape: raise RuntimeError(f"AutoEP expert shard shape mismatch for {param_name}, EP rank " @@ -909,10 +924,8 @@ def _consolidate_zero3_autoep_expert_states(output_dir, model_files, optim_files for (param_name, state_key), ep_rank_fragments in grouped_by_param.items(): missing_ep_ranks = expected_ep_ranks_by_param[param_name] - set(ep_rank_fragments) - placement = placement_by_param.get(param_name) - logical_shape = logical_shape_by_param.get(param_name) - if placement is not None: - affine_map = autoep_placement_to_affine_map(placement, logical_shape) + affine_map = affine_map_by_param.get(param_name) + if affine_map is not None: missing_ep_ranks = { rank for rank in missing_ep_ranks if torch.Size(affine_map.shard_shapes[rank]).numel() != 0 @@ -941,7 +954,8 @@ def _consolidate_zero3_autoep_expert_states(output_dir, model_files, optim_files if not ep_tensors: continue - full_expert_tensor = _rebuild_zero3_autoep_rank_tensors(ep_tensors, placement, logical_shape, + full_expert_tensor = _rebuild_zero3_autoep_rank_tensors(ep_tensors, affine_map, + logical_shape_by_param.get(param_name), f"{param_name}/{state_key}") if full_expert_tensor.shape[0] != num_experts_by_param[param_name]: raise RuntimeError(f"AutoEP universal tensor for {param_name}/{state_key} has wrong expert dimension: " @@ -1175,10 +1189,7 @@ def main(args): slice_shapes = _group_per_tp_shapes(slice_shapes_by_tp, ds_checkpoint.pp_degree, ds_checkpoint.tp_degree) temp_dir = os.path.join(args.output_folder, 'tmp') - from deepspeed.checkpoint.autoep_universal import ( - consolidate_autoep_zero12_expert_states, - get_autoep_zero12_expert_param_info, - ) + from deepspeed.checkpoint.autoep_universal import get_autoep_zero12_expert_param_info expert_files = glob.glob(os.path.join(args.input_folder, 'layer_*_expert_*_model_states.pt')) autoep_expert_file_type = _classify_autoep_expert_file_consolidation(autoep_metadata, expert_files) @@ -1196,12 +1207,23 @@ def main(args): temp_dir, exclude_param_names=set(autoep_expert_param_info)) + # ZeRO may not create optimizer fragments for parameters that have not + # received a gradient. Universal restore still needs their model values. + zero_output_folder = os.path.join(args.output_folder, "zero") + for model_file in ds_checkpoint.mp_rank_files: + model_state = torch.load(model_file, map_location="cpu", weights_only=False) + for name, tensor in model_state.get("module", {}).items(): + param_dir = os.path.join(zero_output_folder, name) + fp32_path = os.path.join(param_dir, "fp32.pt") + if os.path.isfile(fp32_path): + continue + os.makedirs(param_dir, exist_ok=True) + torch.save({PARAM: tensor.float()}, fp32_path) + if autoep_expert_file_type == 'autoep': - print('*** 2.5. Consolidating AutoEP ZeRO-1/2 expert states') - autoep_slice_shapes = {name: per_tp_shapes[0] for name, per_tp_shapes in slice_shapes.items()} - consolidate_autoep_zero12_expert_states(temp_dir, args.output_folder, autoep_expert_param_info, - autoep_slice_shapes, ds_checkpoint.dp_degree, - ds_checkpoint.tp_degree, use_data_before_expert_parallel) + print('*** 2.5. Consolidating AutoEP expert model states') + from deepspeed.checkpoint.autoep_universal import consolidate_autoep_expert_files + consolidate_autoep_expert_files(args.input_folder, args.output_folder, autoep_metadata) print(f' Consolidated {len(autoep_metadata)} AutoEP layer(s)') elif autoep_expert_file_type == 'native_moe': print(f' Found {len(expert_files)} expert checkpoint file(s) but no AutoEP metadata; ' diff --git a/deepspeed/checkpoint/universal_checkpoint.py b/deepspeed/checkpoint/universal_checkpoint.py index c07da4fc5c41..c0ea95931b6a 100644 --- a/deepspeed/checkpoint/universal_checkpoint.py +++ b/deepspeed/checkpoint/universal_checkpoint.py @@ -14,7 +14,7 @@ AUTOEP_EXPERT_PLACEMENT, AUTOEP_PARAM_EP_RANK, AUTOEP_PARAM_LOCAL_EXPERTS, AUTOEP_PARAM_LOGICAL_SHAPE, AUTOEP_PLACEMENT_EXPERTS, AUTOEP_PLACEMENT_RANK, AUTOEP_PLACEMENT_RANKS, UNIVERSAL_CHECKPOINT_VERSION_KEY) -from .autoep_affine import autoep_placement_to_affine_map, extract_autoep_rank_tensor +from .autoep_affine import autoep_metadata_to_affine_map, extract_autoep_rank_tensor @dataclass @@ -85,7 +85,7 @@ def _resolve_autoep_partition(current_param, ckpt_dict, full_hp_param, ep_rank): f"checkpoint={checkpoint_num_experts}, target={logical_shape}.") try: - target_map = autoep_placement_to_affine_map(meta[AUTOEP_EXPERT_PLACEMENT], logical_shape) + target_map = autoep_metadata_to_affine_map(meta, logical_shape) if not isinstance(ep_rank, int) or isinstance(ep_rank, bool): raise ValueError(f"ep_rank must be an integer, got {ep_rank!r}.") metadata_ep_rank = meta[AUTOEP_PARAM_EP_RANK] diff --git a/deepspeed/module_inject/auto_ep_layer.py b/deepspeed/module_inject/auto_ep_layer.py index 2d02a54277b0..c47a262ab562 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -21,10 +21,11 @@ import torch.nn as nn import deepspeed.comm as dist from deepspeed.accelerator import get_accelerator -from deepspeed.checkpoint.autoep_affine import legacy_uniform_autoep_placement_descriptor -from deepspeed.checkpoint.constants import (AUTOEP_EXPERT_PLACEMENT, AUTOEP_PARAM_EP_RANK, AUTOEP_PARAM_LOCAL_EXPERTS, - AUTOEP_PARAM_LOGICAL_SHAPE, AUTOEP_PLACEMENT_EXPERTS, - AUTOEP_PLACEMENT_RANKS, DS_AUTOEP_UC_META) +from deepspeed.checkpoint.autoep_affine import (autoep_placement_to_affine_map, + legacy_uniform_autoep_placement_descriptor) +from deepspeed.checkpoint.constants import (AFFINE_MAP, AUTOEP_EXPERT_PLACEMENT, AUTOEP_PARAM_EP_RANK, + AUTOEP_PARAM_LOCAL_EXPERTS, AUTOEP_PARAM_LOGICAL_SHAPE, + AUTOEP_PLACEMENT_EXPERTS, AUTOEP_PLACEMENT_RANKS, DS_AUTOEP_UC_META) from deepspeed.module_inject.auto_ep_config import AutoEPConfig, MoELayerSpec, resolve_autoep_config_defaults from deepspeed.module_inject.auto_ep_folding import mark_autoep_folding_router_parameter from deepspeed.ops.triton_ops import autoep_fused_token_ops as fused_token_ops @@ -607,11 +608,13 @@ def __init__( 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:]] + affine_map = autoep_placement_to_affine_map(self.expert_placement_descriptor, logical_shape) setattr( param, DS_AUTOEP_UC_META, { AUTOEP_EXPERT_PLACEMENT: self.expert_placement_descriptor, + AFFINE_MAP: affine_map.to_dict(), AUTOEP_PARAM_LOGICAL_SHAPE: logical_shape, AUTOEP_PARAM_EP_RANK: self.ep_rank, AUTOEP_PARAM_LOCAL_EXPERTS: list(local_experts), diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 94663e2ec069..3c24aaef1865 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -66,6 +66,10 @@ DATA_PARALLEL_GROUP, GLOBAL_RANK, DDP_BFLOAT16, GRADIENT_ALLREDUCE_OP_MEAN from deepspeed.runtime.zero.config import ZeroStageEnum from deepspeed.checkpoint.constants import ( + AFFINE_MAP, + AFFINE_MAP_PARAMS, + AFFINE_MAP_VERSION, + AUTOEP_AFFINE_MAPS, AUTOEP_EXPERT_PLACEMENT, AUTOEP_PLACEMENT_EXPERTS, AUTOEP_PLACEMENT_RANKS, @@ -83,11 +87,13 @@ UNIVERSAL_CHECKPOINT_INFO, UNIVERSAL_CHECKPOINT_VERSION_KEY, UNIVERSAL_CHECKPOINT_VERSION_VALUE, + DS_AUTOEP_UC_META, ) from deepspeed.checkpoint.autoep_zero3_metadata import ( is_autoep_zero3_partitioned_entry, validate_autoep_zero3_partitioned_metadata, ) +from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION from deepspeed.checkpoint.utils import clone_tensors_for_torch_save from deepspeed.checkpoint.ds_to_universal import dp_index_to_str from deepspeed.runtime.sparse_tensor import SparseTensor @@ -5063,6 +5069,13 @@ def autoep_expert_writer() -> bool: exp_dp_rank = groups._get_expert_data_parallel_rank(group_name) module_prefix = f"{n_module}." if n_module else "" expert_params = [getattr(module.experts, wname) for wname in ('w1', 'w2', 'w3')] + expert_affine_maps = {} + for wname, param in zip(('w1', 'w2', 'w3'), expert_params): + param_metadata = getattr(param, DS_AUTOEP_UC_META, None) + if not isinstance(param_metadata, dict) or AFFINE_MAP not in param_metadata: + raise RuntimeError(f"AutoEP expert parameter {module_prefix}experts.{wname} " + "is missing save-time affine map metadata.") + expert_affine_maps[f"{module_prefix}experts.{wname}"] = param_metadata[AFFINE_MAP] if self.zero_optimization_partition_weights(): frozen_expert_names = [ f"{module_prefix}experts.{wname}" for wname, param in zip(('w1', 'w2', 'w3'), expert_params) @@ -5087,6 +5100,10 @@ def autoep_expert_writer() -> bool: module.ep_size, AUTOEP_EXPERT_PLACEMENT: module.expert_placement_descriptor, + AUTOEP_AFFINE_MAPS: { + AFFINE_MAP_VERSION: AFFINE_MAP_FORMAT_VERSION, + AFFINE_MAP_PARAMS: expert_affine_maps, + }, 'expert_key_prefix': f"{module_prefix}experts", AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY: diff --git a/tests/unit/checkpoint/test_autoep_affine.py b/tests/unit/checkpoint/test_autoep_affine.py index 57aa5a9251e3..b7fe8885b27c 100644 --- a/tests/unit/checkpoint/test_autoep_affine.py +++ b/tests/unit/checkpoint/test_autoep_affine.py @@ -6,10 +6,13 @@ import pytest import torch -from deepspeed.checkpoint.affine import ParamAffineMap -from deepspeed.checkpoint.autoep_affine import (AUTOEP_PLACEMENT_VERSION, autoep_placement_to_affine_map, - extract_autoep_rank_tensor, legacy_uniform_autoep_placement_descriptor, +from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION, ParamAffineMap +from deepspeed.checkpoint.autoep_affine import (AUTOEP_PLACEMENT_VERSION, autoep_metadata_to_affine_map, + autoep_placement_to_affine_map, 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) def _oracle_shards(full_tensor, experts_by_rank): @@ -52,6 +55,40 @@ def test_non_contiguous_non_uniform_placement_preserves_local_order(): assert torch.equal(affine_map.rebuild(expected_shards), full_tensor) +def test_persisted_autoep_affine_map_takes_precedence_over_placement_lowering(): + logical_shape = (4, 2) + persisted_placement = make_autoep_placement_descriptor(4, [[3, 1], [2, 0]]) + legacy_placement = make_autoep_placement_descriptor(4, [[0, 1], [2, 3]]) + persisted_map = autoep_placement_to_affine_map(persisted_placement, logical_shape) + + metadata = { + AUTOEP_EXPERT_PLACEMENT: legacy_placement, + AFFINE_MAP: persisted_map.to_dict(), + } + loaded_map = autoep_metadata_to_affine_map(metadata, logical_shape) + full_tensor = torch.arange(8, dtype=torch.float32).reshape(logical_shape) + + assert torch.equal(loaded_map.extract(full_tensor, 0), torch.tensor([[6., 7.], [2., 3.]])) + + +def test_versioned_autoep_map_collection_loads_by_parameter_name(): + logical_shape = (2, 2) + descriptor = make_autoep_placement_descriptor(2, [[1], [0]]) + persisted_map = autoep_placement_to_affine_map(descriptor, logical_shape) + metadata = { + AUTOEP_AFFINE_MAPS: { + AFFINE_MAP_VERSION: AFFINE_MAP_FORMAT_VERSION, + AFFINE_MAP_PARAMS: { + 'experts.w1': persisted_map.to_dict() + }, + }, + } + + loaded_map = autoep_metadata_to_affine_map(metadata, logical_shape, 'experts.w1') + + assert loaded_map.to_dict() == persisted_map.to_dict() + + def test_replication_records_exact_holders_and_stops_piece_merging(): experts_by_rank = [[0, 1, 2], [1, 2, 3]] descriptor = make_autoep_placement_descriptor(4, experts_by_rank) diff --git a/tests/unit/v1/moe/test_autoep_unit.py b/tests/unit/v1/moe/test_autoep_unit.py index 6f93d6aaadae..575f72ab2aeb 100644 --- a/tests/unit/v1/moe/test_autoep_unit.py +++ b/tests/unit/v1/moe/test_autoep_unit.py @@ -546,6 +546,7 @@ def test_autoep_layer_marks_zero3_param_placement_families(self): def test_autoep_layer_emits_uniform_placement_and_param_restore_metadata(self): from deepspeed.checkpoint.constants import ( + AFFINE_MAP, AUTOEP_EXPERT_PLACEMENT, AUTOEP_PARAM_EP_RANK, AUTOEP_PARAM_LOCAL_EXPERTS, @@ -585,6 +586,7 @@ def test_autoep_layer_emits_uniform_placement_and_param_restore_metadata(self): assert restore_metadata[AUTOEP_PARAM_LOGICAL_SHAPE] == [4, *param.shape[1:]] assert restore_metadata[AUTOEP_PARAM_EP_RANK] == 1 assert restore_metadata[AUTOEP_PARAM_LOCAL_EXPERTS] == [2, 3] + assert restore_metadata[AFFINE_MAP]["logical_shape"] == [4, *param.shape[1:]] for param in layer.router.parameters(): assert not hasattr(param, DS_AUTOEP_UC_META) From 4012bfc13c56dc26b400c99b7fd1d7d7c854f3a6 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Tue, 22 Sep 2026 15:48:45 +0000 Subject: [PATCH 3/6] Validate AutoEP affine placement metadata Use explicit placement rank IDs for local expert lookup and reject checkpoints whose persisted affine map conflicts with placement provenance. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- deepspeed/checkpoint/affine_ir_spec.md | 8 +++-- deepspeed/checkpoint/autoep_affine.py | 23 ++++++++++--- deepspeed/module_inject/auto_ep_layer.py | 8 ++--- deepspeed/runtime/engine.py | 7 ++-- tests/unit/checkpoint/test_autoep_affine.py | 38 +++++++++++++++------ 5 files changed, 57 insertions(+), 27 deletions(-) diff --git a/deepspeed/checkpoint/affine_ir_spec.md b/deepspeed/checkpoint/affine_ir_spec.md index 133a015d2073..52413c6b39d6 100644 --- a/deepspeed/checkpoint/affine_ir_spec.md +++ b/deepspeed/checkpoint/affine_ir_spec.md @@ -574,9 +574,11 @@ descriptor in EP-local rank coordinates. The persisted map is the geometry execu contract for conversion and restore; descriptor lowering is retained only as a legacy checkpoint fallback. The descriptor records each rank's ordered global expert IDs, including uneven, non-contiguous, replicated, and empty placements, so it remains useful -for placement provenance and validation rather than scheduling policy. ZeRO and EDP -fragments remain outside the map: callers first normalize storage to one logical packed -expert tensor per EP rank. +for placement provenance and validation rather than scheduling policy. Rank entries are +identified by their explicit rank IDs, not their list positions. When a checkpoint carries +both forms, lowering the descriptor must reproduce the persisted map exactly; a conflict +invalidates the checkpoint. 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. diff --git a/deepspeed/checkpoint/autoep_affine.py b/deepspeed/checkpoint/autoep_affine.py index 2fa48767acf4..f30311803368 100644 --- a/deepspeed/checkpoint/autoep_affine.py +++ b/deepspeed/checkpoint/autoep_affine.py @@ -9,15 +9,16 @@ from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION, AffinePiece, ParamAffineMap from deepspeed.checkpoint.constants import (AFFINE_MAP, AFFINE_MAP_PARAMS, AFFINE_MAP_VERSION, AUTOEP_AFFINE_MAPS, - AUTOEP_PLACEMENT_EP_SIZE, AUTOEP_PLACEMENT_EXPERTS, - AUTOEP_PLACEMENT_NUM_EXPERTS, AUTOEP_PLACEMENT_RANK, - AUTOEP_PLACEMENT_RANKS, AUTOEP_PLACEMENT_VERSION, + AUTOEP_EXPERT_PLACEMENT, 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', + 'autoep_experts_for_rank', 'legacy_uniform_autoep_placement_descriptor', 'autoep_placement_to_affine_map', 'autoep_metadata_to_affine_map', @@ -89,6 +90,15 @@ def validate_autoep_placement_descriptor(descriptor): raise ValueError(f'AutoEP placement descriptor does not cover global expert IDs {missing}.') +def autoep_experts_for_rank(descriptor, rank): + """Return expert IDs in one rank's local packed order.""" + validate_autoep_placement_descriptor(descriptor) + entries_by_rank = {entry[AUTOEP_PLACEMENT_RANK]: entry for entry in descriptor[AUTOEP_PLACEMENT_RANKS]} + if not isinstance(rank, int) or isinstance(rank, bool) or rank not in entries_by_rank: + raise ValueError(f'AutoEP placement rank {rank!r} is outside [0, {len(entries_by_rank)}).') + return entries_by_rank[rank][AUTOEP_PLACEMENT_EXPERTS] + + 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') @@ -156,9 +166,14 @@ def autoep_metadata_to_affine_map(metadata, logical_shape, param_name=None): raise ValueError(f"AutoEP affine map logical shape {affine_map.logical_shape} does not match " f"parameter logical shape {expected_shape}.") affine_map.validate() + placement = metadata.get(AUTOEP_EXPERT_PLACEMENT) + if placement is not None: + expected_map = autoep_placement_to_affine_map(placement, expected_shape) + if affine_map.to_dict() != expected_map.to_dict(): + raise ValueError("AutoEP affine map disagrees with the expert placement descriptor.") return affine_map - placement = metadata.get('expert_placement') + placement = metadata.get(AUTOEP_EXPERT_PLACEMENT) if placement is None: raise ValueError("AutoEP metadata has neither an affine map nor an expert placement descriptor.") return autoep_placement_to_affine_map(placement, logical_shape) diff --git a/deepspeed/module_inject/auto_ep_layer.py b/deepspeed/module_inject/auto_ep_layer.py index c47a262ab562..42c88a079a02 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -21,11 +21,10 @@ import torch.nn as nn import deepspeed.comm as dist from deepspeed.accelerator import get_accelerator -from deepspeed.checkpoint.autoep_affine import (autoep_placement_to_affine_map, +from deepspeed.checkpoint.autoep_affine import (autoep_experts_for_rank, autoep_placement_to_affine_map, legacy_uniform_autoep_placement_descriptor) from deepspeed.checkpoint.constants import (AFFINE_MAP, AUTOEP_EXPERT_PLACEMENT, AUTOEP_PARAM_EP_RANK, - AUTOEP_PARAM_LOCAL_EXPERTS, AUTOEP_PARAM_LOGICAL_SHAPE, - AUTOEP_PLACEMENT_EXPERTS, AUTOEP_PLACEMENT_RANKS, DS_AUTOEP_UC_META) + 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 from deepspeed.module_inject.auto_ep_folding import mark_autoep_folding_router_parameter from deepspeed.ops.triton_ops import autoep_fused_token_ops as fused_token_ops @@ -603,8 +602,7 @@ def __init__( self.experts.w1.requires_grad_(w1_requires_grad) self.experts.w2.requires_grad_(w2_requires_grad) self.experts.w3.requires_grad_(w3_requires_grad) - local_experts = self.expert_placement_descriptor[AUTOEP_PLACEMENT_RANKS][ - self.ep_rank][AUTOEP_PLACEMENT_EXPERTS] + 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:]] diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 3c24aaef1865..0f7c97f13a56 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -71,8 +71,6 @@ AFFINE_MAP_VERSION, AUTOEP_AFFINE_MAPS, AUTOEP_EXPERT_PLACEMENT, - AUTOEP_PLACEMENT_EXPERTS, - AUTOEP_PLACEMENT_RANKS, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY, @@ -94,6 +92,7 @@ validate_autoep_zero3_partitioned_metadata, ) from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION +from deepspeed.checkpoint.autoep_affine import autoep_experts_for_rank from deepspeed.checkpoint.utils import clone_tensors_for_torch_save from deepspeed.checkpoint.ds_to_universal import dp_index_to_str from deepspeed.runtime.sparse_tensor import SparseTensor @@ -4437,10 +4436,10 @@ def _validate_autoep_zero3_partitioned_metadata(autoep_layers, for module_name, module in model.named_modules(): if not isinstance(module, _AutoEPMoELayer): continue - rank_entry = module.expert_placement_descriptor[AUTOEP_PLACEMENT_RANKS][module.ep_rank] expected_runtime_layers[module_name] = { 'ep_rank': module.ep_rank, - 'local_experts': list(rank_entry[AUTOEP_PLACEMENT_EXPERTS]), + 'local_experts': + list(autoep_experts_for_rank(module.expert_placement_descriptor, module.ep_rank)), } if not expected_runtime_layers: expected_runtime_layers = None diff --git a/tests/unit/checkpoint/test_autoep_affine.py b/tests/unit/checkpoint/test_autoep_affine.py index b7fe8885b27c..e73e6fb4f352 100644 --- a/tests/unit/checkpoint/test_autoep_affine.py +++ b/tests/unit/checkpoint/test_autoep_affine.py @@ -7,9 +7,9 @@ import torch from deepspeed.checkpoint.affine import AFFINE_MAP_FORMAT_VERSION, ParamAffineMap -from deepspeed.checkpoint.autoep_affine import (AUTOEP_PLACEMENT_VERSION, autoep_metadata_to_affine_map, - autoep_placement_to_affine_map, extract_autoep_rank_tensor, - legacy_uniform_autoep_placement_descriptor, +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, 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) @@ -55,20 +55,36 @@ def test_non_contiguous_non_uniform_placement_preserves_local_order(): assert torch.equal(affine_map.rebuild(expected_shards), full_tensor) -def test_persisted_autoep_affine_map_takes_precedence_over_placement_lowering(): +@pytest.mark.parametrize('use_collection', [False, True]) +def test_persisted_autoep_affine_map_must_match_placement_descriptor(use_collection): logical_shape = (4, 2) persisted_placement = make_autoep_placement_descriptor(4, [[3, 1], [2, 0]]) legacy_placement = make_autoep_placement_descriptor(4, [[0, 1], [2, 3]]) persisted_map = autoep_placement_to_affine_map(persisted_placement, logical_shape) - metadata = { - AUTOEP_EXPERT_PLACEMENT: legacy_placement, - AFFINE_MAP: persisted_map.to_dict(), - } - loaded_map = autoep_metadata_to_affine_map(metadata, logical_shape) - full_tensor = torch.arange(8, dtype=torch.float32).reshape(logical_shape) + metadata = {AUTOEP_EXPERT_PLACEMENT: legacy_placement} + if use_collection: + metadata[AUTOEP_AFFINE_MAPS] = { + AFFINE_MAP_VERSION: AFFINE_MAP_FORMAT_VERSION, + AFFINE_MAP_PARAMS: { + 'experts.w1': persisted_map.to_dict() + }, + } + else: + metadata[AFFINE_MAP] = persisted_map.to_dict() + with pytest.raises(ValueError, match='disagrees with the expert placement'): + autoep_metadata_to_affine_map(metadata, logical_shape, 'experts.w1') + + +def test_explicit_rank_ids_define_entries_independently_of_list_order(): + descriptor = make_autoep_placement_descriptor(4, [[0, 2], [3, 1]]) + descriptor['ranks'].reverse() - assert torch.equal(loaded_map.extract(full_tensor, 0), torch.tensor([[6., 7.], [2., 3.]])) + assert autoep_experts_for_rank(descriptor, 0) == [0, 2] + assert autoep_experts_for_rank(descriptor, 1) == [3, 1] + affine_map = autoep_placement_to_affine_map(descriptor, (4, 2)) + full_tensor = torch.arange(8).reshape(4, 2) + assert torch.equal(affine_map.extract(full_tensor, 0), torch.stack([full_tensor[0], full_tensor[2]])) def test_versioned_autoep_map_collection_loads_by_parameter_name(): From a924776c372c72bf595431c3462330324cef8442 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Wed, 23 Sep 2026 16:28:51 +0000 Subject: [PATCH 4/6] Fix AutoEP ZeRO-1/2 universal optimizer states Consolidate ZeRO FP32 master expert weights and optimizer states after per-expert model files, and verify native and Universal loss-resume parity for ZeRO stages 1 and 2. Signed-off-by: Jin, Youzhi Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- deepspeed/checkpoint/ds_to_universal.py | 8 ++- tests/unit/v1/moe/test_autoep_checkpoint.py | 58 +++++++++++++++++++-- 2 files changed, 60 insertions(+), 6 deletions(-) diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py index 825705d17101..4500af4e82ab 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/deepspeed/checkpoint/ds_to_universal.py @@ -1189,7 +1189,8 @@ def main(args): slice_shapes = _group_per_tp_shapes(slice_shapes_by_tp, ds_checkpoint.pp_degree, ds_checkpoint.tp_degree) temp_dir = os.path.join(args.output_folder, 'tmp') - from deepspeed.checkpoint.autoep_universal import get_autoep_zero12_expert_param_info + from deepspeed.checkpoint.autoep_universal import (consolidate_autoep_zero12_expert_states, + get_autoep_zero12_expert_param_info) expert_files = glob.glob(os.path.join(args.input_folder, 'layer_*_expert_*_model_states.pt')) autoep_expert_file_type = _classify_autoep_expert_file_consolidation(autoep_metadata, expert_files) @@ -1225,6 +1226,11 @@ def main(args): from deepspeed.checkpoint.autoep_universal import consolidate_autoep_expert_files consolidate_autoep_expert_files(args.input_folder, args.output_folder, autoep_metadata) print(f' Consolidated {len(autoep_metadata)} AutoEP layer(s)') + print('*** 2.6. Consolidating AutoEP ZeRO-1/2 expert states') + autoep_slice_shapes = {name: per_tp_shapes[0] for name, per_tp_shapes in slice_shapes.items()} + consolidate_autoep_zero12_expert_states(temp_dir, args.output_folder, autoep_expert_param_info, + autoep_slice_shapes, ds_checkpoint.dp_degree, + ds_checkpoint.tp_degree, use_data_before_expert_parallel) elif autoep_expert_file_type == 'native_moe': print(f' Found {len(expert_files)} expert checkpoint file(s) but no AutoEP metadata; ' 'assuming native DeepSpeed MoE and skipping AutoEP consolidation') diff --git a/tests/unit/v1/moe/test_autoep_checkpoint.py b/tests/unit/v1/moe/test_autoep_checkpoint.py index 8b0c2967af38..90dc56941b49 100644 --- a/tests/unit/v1/moe/test_autoep_checkpoint.py +++ b/tests/unit/v1/moe/test_autoep_checkpoint.py @@ -60,8 +60,14 @@ def _load_universal_file(universal_dir, param_name, key): def _load_universal_dense_state(universal_dir, param_name, key): + from deepspeed.checkpoint.constants import PARAM + state = _load_universal_file(universal_dir, param_name, key) - assert torch.is_tensor(state), f"expected raw tensor state for dense ZeRO-3 parameter {param_name}/{key}" + if isinstance(state, dict): + tensor = state.get(PARAM) + assert torch.is_tensor(tensor), f"expected tensor payload for dense parameter {param_name}/{key}" + return tensor + assert torch.is_tensor(state), f"expected tensor state for dense parameter {param_name}/{key}" return state @@ -212,6 +218,7 @@ def _assert_router_params_match_universal(engine, universal_dir): for param_name, param in _router_params(engine): restored = _gather_zero_param(param).cpu() expected = _load_universal_dense_state(universal_dir, param_name, "fp32").view_as(restored) + expected = expected.to(dtype=restored.dtype) torch.testing.assert_close(restored, expected, rtol=0, atol=0) @@ -219,6 +226,7 @@ def _assert_shared_params_match_universal(engine, universal_dir): for param_name, param in _shared_params(engine): restored = _gather_zero_param(param).cpu() expected = _load_universal_dense_state(universal_dir, param_name, "fp32").view_as(restored) + expected = expected.to(dtype=restored.dtype) torch.testing.assert_close(restored, expected, rtol=0, atol=0) @@ -228,6 +236,7 @@ def _assert_expert_params_match_universal(engine, universal_dir): restored = _collect_by_ep_rank(local_experts, module.ep_rank, module.ep_size, engine.device) if dist.get_rank() == 0: expected = _load_universal_expert_state(universal_dir, param_name, "fp32") + expected = expected.to(dtype=restored.dtype) torch.testing.assert_close(restored, expected, rtol=0, atol=0) @@ -305,6 +314,22 @@ def _run_training_steps_with_engine_input_dtype(engine, num_steps=2, seq_len=8, return losses +def _fixed_training_inputs(num_steps, seq_len=8, hidden_dim=64): + generator = torch.Generator().manual_seed(29041) + return [torch.randn(1, seq_len, hidden_dim, generator=generator) for _ in range(num_steps)] + + +def _run_training_steps_with_inputs(engine, inputs): + losses = [] + for x in inputs: + x = x.to(device=engine.device, dtype=engine_input_dtype(engine)) + loss = engine(x).mean() + engine.backward(loss) + engine.step() + losses.append(loss.item()) + return losses + + @pytest.mark.parametrize("data_first,dp_order", [(False, [0, 2, 1, 3]), (True, [0, 1, 2, 3])], ids=["E+D", "D+E"]) def test_zero12_expert_state_consolidation_exact(tmpdir, data_first, dp_order): temp_dir = os.path.join(str(tmpdir), "tmp") @@ -462,25 +487,48 @@ class TestAutoEPZero12UniversalCheckpoint(DistributedTest): world_size = 4 @pytest.mark.parametrize("zero_stage", [1, 2]) - def test_round_trip_weights_and_training_step(self, tmpdir, zero_stage): + def test_round_trip_loss_parity(self, tmpdir, zero_stage): skip_unless_h100_tests_enabled("AutoEP ZeRO-1/2 Universal Checkpoint coverage requires H100") seed_everything(8100 + zero_stage) + inputs = _fixed_training_inputs(num_steps=4) config = make_autoep_config(zero_stage=zero_stage, ep_size=2, mixed_precision=True) engine, _, _, _ = deepspeed.initialize(model=MockMoETransformer(), config=config) - _run_training_steps_with_engine_input_dtype(engine, num_steps=2) save_dir, tag = str(tmpdir), f"autoep-zero{zero_stage}" + baseline_losses = _run_training_steps_with_inputs(engine, inputs[:2]) engine.save_checkpoint(save_dir, tag=tag) - universal_dir = _convert_checkpoint_to_universal(save_dir, tag) + baseline_losses.extend(_run_training_steps_with_inputs(engine, inputs[2:])) engine.destroy() + universal_dir = _convert_checkpoint_to_universal(save_dir, tag) + if dist.get_rank() == 0: + _assert_universal_expert_metadata(universal_dir, num_experts=4) + + seed_everything(8100 + zero_stage) + native, _, _, _ = deepspeed.initialize(model=MockMoETransformer(), config=config) + native.load_checkpoint(save_dir, tag=tag) + native_losses = _run_training_steps_with_inputs(native, inputs[2:]) + native.destroy() + + for actual, expected in zip(native_losses, baseline_losses[2:]): + torch.testing.assert_close(torch.tensor(actual), torch.tensor(expected), rtol=1e-4, atol=1e-5) + + seed_everything(8100 + zero_stage) load_config = make_autoep_config(zero_stage=zero_stage, ep_size=2, mixed_precision=True) load_config["checkpoint"] = {"load_universal": True} restored, _, _, _ = deepspeed.initialize(model=MockMoETransformer(), config=load_config) restored.load_checkpoint(save_dir, tag=f"{tag}_universal") _assert_module_params_match_universal(restored, universal_dir) - assert torch.isfinite(torch.tensor(_run_training_steps_with_engine_input_dtype(restored, num_steps=1)[0])) + universal_losses = _run_training_steps_with_inputs(restored, inputs[2:]) restored.destroy() + for path_name, losses in (("native", native_losses), ("universal", universal_losses)): + for step, (actual, expected) in enumerate(zip(losses, baseline_losses[2:]), start=3): + torch.testing.assert_close(torch.tensor(actual), + torch.tensor(expected), + rtol=1e-4, + atol=1e-5, + msg=f"{path_name} resume loss diverged at step {step}.") + class TestAutoEPZero3UniversalCheckpoint(DistributedTest): world_size = 2 From b7df13e44cfc4ca83a1adf71fdb94e98addef836 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Thu, 24 Sep 2026 03:21:23 +0000 Subject: [PATCH 5/6] Avoid duplicate AutoEP FP32 checkpoint writes Remove the unreachable expp_rank optimizer consolidation helper; ZeRO-1/2 optimizer state is handled by the active shard merger. Keep expert-file FP32 fallbacks in the temporary directory and let the merger write each final tensor once. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- deepspeed/checkpoint/autoep_universal.py | 283 +++++------------- deepspeed/checkpoint/ds_to_universal.py | 23 +- .../test_autoep_affine_universal.py | 69 ++++- 3 files changed, 169 insertions(+), 206 deletions(-) diff --git a/deepspeed/checkpoint/autoep_universal.py b/deepspeed/checkpoint/autoep_universal.py index fd386bcf908b..5f300b19a8da 100644 --- a/deepspeed/checkpoint/autoep_universal.py +++ b/deepspeed/checkpoint/autoep_universal.py @@ -11,8 +11,6 @@ import os import glob import torch - -from deepspeed.utils import logger from .affine import AFFINE_MAP_FORMAT_VERSION from .constants import ( @@ -127,70 +125,6 @@ def validate_folding_metadata(metadata, return folding -def _state_entry(state, param_id): - """Get optimizer state entry by param id, handling int/str key variants.""" - if param_id in state: - return state[param_id] - - pid_str = str(param_id) - if pid_str in state: - return state[pid_str] - - if isinstance(param_id, str): - try: - pid_int = int(param_id) - except ValueError: - return None - return state.get(pid_int) - - return None - - -def _ordered_param_ids(optim_sd): - """Return optimizer param ids in param_groups order, deduplicated.""" - ordered = [] - seen = set() - for group in optim_sd.get('param_groups', []): - for param_id in group.get('params', []): - key = str(param_id) - if key in seen: - continue - seen.add(key) - ordered.append(param_id) - - if ordered: - return ordered - - # Fallback for unexpected optimizer formats. - state = optim_sd.get('state', {}) - return list(state.keys()) - - -def _param_name_to_id(optim_sd): - """Build optional mapping from parameter name to optimizer param id.""" - mapping = {} - for group in optim_sd.get('param_groups', []): - params = group.get('params', []) - param_names = group.get('param_names', None) - if not isinstance(param_names, list): - continue - if len(param_names) != len(params): - continue - for param_id, param_name in zip(params, param_names): - mapping[param_name] = param_id - return mapping - - -def _is_expert_optimizer_state(param_state, num_local): - for state_key in ('exp_avg', 'exp_avg_sq'): - tensor = param_state.get(state_key) - if tensor is None: - continue - if tensor.dim() == 3 and tensor.shape[0] == num_local: - return True - return False - - def resolve_expert_ckpt_path(checkpoint_dir, moe_layer_id, global_expert_id): """Find the expert checkpoint file for a given (layer, expert) pair. @@ -317,14 +251,48 @@ def _autoep_zero12_dp_ranks(ep_rank, dp_degree, ep_size, use_data_before_expert_ return range(ep_rank, dp_degree, ep_size) +def get_autoep_zero12_fp32_fallback_param_names(temp_dir, expert_param_info, slice_shapes, dp_degree, + use_data_before_expert_parallel): + """Find parameters that need per-expert weights as a fallback for missing master fragments.""" + fallback_param_names = set() + for param_name, metadata in expert_param_info.items(): + if param_name not in slice_shapes: + raise RuntimeError(f"AutoEP expert parameter {param_name} is missing from checkpoint parameter shapes.") + + ep_size = metadata['ep_size'] + local_shape = tuple(slice_shapes[param_name]) + affine_map = None + if metadata.get('expert_placement') is not None: + logical_shape = (metadata['num_experts'], ) + local_shape[1:] + affine_map = autoep_metadata_to_affine_map(metadata, logical_shape) + + for ep_rank in range(ep_size): + expected_shape = affine_map.shard_shapes[ep_rank] if affine_map is not None else local_shape + if torch.Size(expected_shape).numel() == 0: + continue + dp_ranks = _autoep_zero12_dp_ranks(ep_rank, dp_degree, ep_size, use_data_before_expert_parallel) + if any(not os.path.isfile(_zero12_fragment_path(temp_dir, param_name, "fp32", dp_rank)) + for dp_rank in dp_ranks): + fallback_param_names.add(param_name) + break + + return fallback_param_names + + def _zero12_values_equal(left, right): if torch.is_tensor(left) and torch.is_tensor(right): return torch.equal(left, right) return left == right -def consolidate_autoep_zero12_expert_states(temp_dir, output_dir, expert_param_info, slice_shapes, dp_degree, - tp_degree, use_data_before_expert_parallel): +def consolidate_autoep_zero12_expert_states(temp_dir, + output_dir, + expert_param_info, + slice_shapes, + dp_degree, + tp_degree, + use_data_before_expert_parallel, + fp32_fallback_dir=None): """Consolidate AutoEP expert FP32 and Adam states from ZeRO-1/2 fragments.""" if tp_degree != 1: raise NotImplementedError("ZeRO-1/2 Universal Checkpoint conversion for AutoEP with tensor parallelism " @@ -355,6 +323,7 @@ def consolidate_autoep_zero12_expert_states(temp_dir, output_dir, expert_param_i for state_name in ('fp32', 'exp_avg', 'exp_avg_sq'): ep_tensors = {} + fp32_fallback_tensor = None 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 = [] @@ -377,10 +346,39 @@ def consolidate_autoep_zero12_expert_states(temp_dir, output_dir, expert_param_i ep_tensors[ep_rank] = local_tensor continue if not fragments: + fallback_path = (os.path.join(fp32_fallback_dir, param_name, "fp32.pt") + if fp32_fallback_dir is not None else os.path.join(param_dir, "fp32.pt")) + if state_name == "fp32" and os.path.isfile(fallback_path): + if fp32_fallback_tensor is None: + fp32_fallback_state = torch.load(fallback_path, map_location='cpu', weights_only=False) + fp32_fallback_tensor = fp32_fallback_state[PARAM] + if affine_map is not None: + local_tensor = affine_map.extract(fp32_fallback_tensor, ep_rank) + else: + start = ep_rank * metadata['num_local_experts'] + end = start + metadata['num_local_experts'] + local_tensor = fp32_fallback_tensor[start:end] + ep_tensors[ep_rank] = local_tensor.reshape(expected_shape) + continue raise RuntimeError(f"Missing AutoEP {state_name} fragments for {param_name}, EP rank {ep_rank}.") local_tensor = torch.cat(fragments, dim=0) if local_tensor.numel() != expected_numel: + fallback_path = (os.path.join(fp32_fallback_dir, param_name, "fp32.pt") + if fp32_fallback_dir is not None else os.path.join(param_dir, "fp32.pt")) + if (state_name == "fp32" and local_tensor.numel() < expected_numel + and os.path.isfile(fallback_path)): + if fp32_fallback_tensor is None: + fp32_fallback_state = torch.load(fallback_path, map_location='cpu', weights_only=False) + fp32_fallback_tensor = fp32_fallback_state[PARAM] + if affine_map is not None: + local_tensor = affine_map.extract(fp32_fallback_tensor, ep_rank) + else: + start = ep_rank * metadata['num_local_experts'] + end = start + metadata['num_local_experts'] + local_tensor = fp32_fallback_tensor[start:end] + ep_tensors[ep_rank] = local_tensor.reshape(expected_shape) + continue 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[ep_rank] = local_tensor.reshape(expected_shape) @@ -412,7 +410,11 @@ def consolidate_autoep_zero12_expert_states(temp_dir, output_dir, expert_param_i torch.save(step_values[0], os.path.join(param_dir, "step.pt")) -def consolidate_autoep_expert_files(checkpoint_dir, output_dir, autoep_layers_metadata): +def consolidate_autoep_expert_files(checkpoint_dir, + output_dir, + autoep_layers_metadata, + fp32_fallback_param_names=None, + fp32_fallback_dir=None): """Consolidate per-expert checkpoint files into full-expert universal format. For each AutoEP layer, loads all per-expert files, stacks into @@ -422,6 +424,8 @@ def consolidate_autoep_expert_files(checkpoint_dir, output_dir, autoep_layers_me checkpoint_dir: Path to DeepSpeed checkpoint directory. output_dir: Path to universal checkpoint output directory. autoep_layers_metadata: AutoEP metadata list from main checkpoint. + fp32_fallback_param_names: Parameters whose ZeRO FP32 fragments are incomplete. + fp32_fallback_dir: Temporary directory for fallback tensors consumed by ZeRO consolidation. Raises: FileNotFoundError: If expected expert files are missing. @@ -446,7 +450,9 @@ def consolidate_autoep_expert_files(checkpoint_dir, output_dir, autoep_layers_me raise RuntimeError(f"AutoEP expert placement disagrees with num_experts for {prefix}.") for wname in ('w1', 'w2', 'w3'): - expert_tensors = [] + param_name = f"{prefix}.{wname}" + save_fp32 = fp32_fallback_param_names is None or param_name in fp32_fallback_param_names + expert_tensors = [] if save_fp32 else None folding_metadata = None for global_eid in range(num_experts): ckpt_path = resolve_expert_ckpt_path(checkpoint_dir, moe_layer_id, global_eid) @@ -456,14 +462,17 @@ def consolidate_autoep_expert_files(checkpoint_dir, output_dir, autoep_layers_me key = f"{prefix}.{wname}.{global_eid}" if key not in sd: raise RuntimeError(f"Expected key '{key}' not found in {ckpt_path}") - expert_tensors.append(sd[key]) + if expert_tensors is not None: + expert_tensors.append(sd[key]) # Stack to full fused tensor [E_total, H, D] + if expert_tensors is None: + continue full_tensor = torch.stack(expert_tensors, dim=0) - # Save in universal format - param_name = f"{prefix}.{wname}" - param_dir = os.path.join(output_dir, "zero", param_name) + # Store fallbacks separately so the ZeRO merger writes each final fp32.pt once. + param_dir = (os.path.join(fp32_fallback_dir, param_name) if fp32_fallback_dir is not None + and fp32_fallback_param_names is not None else os.path.join(output_dir, "zero", param_name)) os.makedirs(param_dir, exist_ok=True) universal_state = { PARAM: full_tensor, @@ -474,129 +483,3 @@ def consolidate_autoep_expert_files(checkpoint_dir, output_dir, autoep_layers_me if folding_metadata is not None: universal_state[FOLDING_METADATA_KEY] = folding_metadata torch.save(universal_state, os.path.join(param_dir, "fp32.pt")) - - -def consolidate_autoep_optimizer_states(checkpoint_dir, output_dir, autoep_layers_metadata, ep_size): - """Consolidate expert optimizer states from expp_rank files into universal format. - - Loads optimizer states from all expp_rank_*_optim_states.pt files, - extracts per-expert-parameter states (exp_avg, exp_avg_sq, etc.), - concatenates along the expert dimension (dim 0) to form full - [E_total, H, D] optimizer states, and saves alongside the model - parameter in universal format. - - Args: - checkpoint_dir: Path to DeepSpeed checkpoint directory. - output_dir: Path to universal checkpoint output directory. - autoep_layers_metadata: AutoEP metadata list from main checkpoint. - ep_size: Expert parallel world size (number of expp_rank files to load). - - Raises: - FileNotFoundError: If expected optimizer state files are missing. - RuntimeError: If expert parameter states cannot be extracted. - """ - if autoep_layers_metadata is None: - raise RuntimeError("AutoEP metadata is missing. Cannot consolidate optimizer states.") - - # Load all expp_rank optimizer states - optim_states = [] - for rank in range(ep_size): - pattern = os.path.join(checkpoint_dir, f'expp_rank_{rank}_mp_rank_*_optim_states.pt') - matches = glob.glob(pattern) - if not matches: - # No optimizer state files (e.g., ZeRO handles optimizer differently) - return - optim_path = matches[0] - sd = torch.load(optim_path, map_location='cpu', weights_only=False) - optim_states.append(sd) - - if not optim_states: - return - - # Extract optimizer state dict - optim_sd = optim_states[0].get('optimizer') - if optim_sd is None: - logger.warning("AutoEP per-expert optimizer checkpoint has no optimizer payload; ZeRO-1/2 checkpoints " - "must be consolidated from their ZeRO optimizer shards instead.") - return - - state = optim_sd.get('state', {}) - - if not state: - return - - ordered_param_ids = _ordered_param_ids(optim_sd) - name_to_param_id = _param_name_to_id(optim_sd) - consumed_param_ids = set() - - # For each AutoEP layer, extract and consolidate optimizer states - for layer_info in autoep_layers_metadata: - prefix = layer_info['expert_key_prefix'] - num_experts = layer_info['num_experts'] - num_local = layer_info['num_local_experts'] - layer_param_ids = {} - - # If optimizer state carries param names, map weights by exact identity. - for wname in ('w1', 'w2', 'w3'): - param_name = f"{prefix}.{wname}" - param_id = name_to_param_id.get(param_name) - if param_id is None: - continue - layer_param_ids[wname] = param_id - consumed_param_ids.add(str(param_id)) - - # Fallback: consume expert-like params in optimizer param_groups order. - missing_wnames = [w for w in ('w1', 'w2', 'w3') if w not in layer_param_ids] - if missing_wnames: - candidates = [] - for param_id in ordered_param_ids: - if str(param_id) in consumed_param_ids: - continue - param_state = _state_entry(state, param_id) - if param_state is None: - continue - if not _is_expert_optimizer_state(param_state, num_local): - continue - candidates.append(param_id) - - for wname, param_id in zip(missing_wnames, candidates): - layer_param_ids[wname] = param_id - consumed_param_ids.add(str(param_id)) - - for wname in ('w1', 'w2', 'w3'): - param_name = f"{prefix}.{wname}" - param_dir = os.path.join(output_dir, "zero", param_name) - os.makedirs(param_dir, exist_ok=True) - param_id = layer_param_ids.get(wname) - if param_id is None: - continue - - # Consolidate optimizer states for this specific expert parameter id. - for state_key in ('exp_avg', 'exp_avg_sq'): - rank_tensors = [] - - for rank in range(ep_size): - rank_optim_sd = optim_states[rank].get('optimizer', {}) - rank_state = rank_optim_sd.get('state', {}) - param_state = _state_entry(rank_state, param_id) - if param_state is None: - rank_tensors = [] - break - tensor = param_state.get(state_key) - if tensor is None: - rank_tensors = [] - break - if tensor.dim() != 3 or tensor.shape[0] != num_local: - rank_tensors = [] - break - rank_tensors.append(tensor) - - if len(rank_tensors) == ep_size: - full_tensor = torch.cat(rank_tensors, dim=0) - torch.save( - { - PARAM: full_tensor, - CAT_DIM: 0, - EP_IS_EXPERT_PARAM: True, - EP_NUM_EXPERTS: num_experts, - }, os.path.join(param_dir, f"{state_key}.pt")) diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py index 4500af4e82ab..3d7a1668a2c6 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/deepspeed/checkpoint/ds_to_universal.py @@ -1190,6 +1190,7 @@ def main(args): temp_dir = os.path.join(args.output_folder, 'tmp') from deepspeed.checkpoint.autoep_universal import (consolidate_autoep_zero12_expert_states, + get_autoep_zero12_fp32_fallback_param_names, get_autoep_zero12_expert_param_info) expert_files = glob.glob(os.path.join(args.input_folder, 'layer_*_expert_*_model_states.pt')) @@ -1224,13 +1225,25 @@ def main(args): if autoep_expert_file_type == 'autoep': print('*** 2.5. Consolidating AutoEP expert model states') from deepspeed.checkpoint.autoep_universal import consolidate_autoep_expert_files - consolidate_autoep_expert_files(args.input_folder, args.output_folder, autoep_metadata) + autoep_slice_shapes = {name: per_tp_shapes[0] for name, per_tp_shapes in slice_shapes.items()} + fp32_fallback_param_names = get_autoep_zero12_fp32_fallback_param_names( + temp_dir, autoep_expert_param_info, autoep_slice_shapes, ds_checkpoint.dp_degree, + use_data_before_expert_parallel) + consolidate_autoep_expert_files(args.input_folder, + args.output_folder, + autoep_metadata, + fp32_fallback_param_names=fp32_fallback_param_names, + fp32_fallback_dir=temp_dir) print(f' Consolidated {len(autoep_metadata)} AutoEP layer(s)') print('*** 2.6. Consolidating AutoEP ZeRO-1/2 expert states') - autoep_slice_shapes = {name: per_tp_shapes[0] for name, per_tp_shapes in slice_shapes.items()} - consolidate_autoep_zero12_expert_states(temp_dir, args.output_folder, autoep_expert_param_info, - autoep_slice_shapes, ds_checkpoint.dp_degree, - ds_checkpoint.tp_degree, use_data_before_expert_parallel) + consolidate_autoep_zero12_expert_states(temp_dir, + args.output_folder, + autoep_expert_param_info, + autoep_slice_shapes, + ds_checkpoint.dp_degree, + ds_checkpoint.tp_degree, + use_data_before_expert_parallel, + fp32_fallback_dir=temp_dir) elif autoep_expert_file_type == 'native_moe': print(f' Found {len(expert_files)} expert checkpoint file(s) but no AutoEP metadata; ' 'assuming native DeepSpeed MoE and skipping AutoEP consolidation') diff --git a/tests/unit/checkpoint/test_autoep_affine_universal.py b/tests/unit/checkpoint/test_autoep_affine_universal.py index 885998a9a030..10e69e458214 100644 --- a/tests/unit/checkpoint/test_autoep_affine_universal.py +++ b/tests/unit/checkpoint/test_autoep_affine_universal.py @@ -8,7 +8,9 @@ import torch from deepspeed.checkpoint.autoep_affine import make_autoep_placement_descriptor -from deepspeed.checkpoint.autoep_universal import (consolidate_autoep_zero12_expert_states, +from deepspeed.checkpoint.autoep_universal import (consolidate_autoep_expert_files, + consolidate_autoep_zero12_expert_states, _zero12_fragment_path, + get_autoep_zero12_fp32_fallback_param_names, get_autoep_zero12_expert_param_info) from deepspeed.checkpoint.constants import ( AUTOEP_EXPERT_PLACEMENT, AUTOEP_LAYERS_KEY, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY, @@ -106,6 +108,71 @@ def test_zero12_uniform_affine_consolidation_matches_legacy_order(tmp_path): assert torch.equal(_load_consolidated(str(tmp_path / "output"), param_name)[PARAM], full) +def test_zero12_consolidation_uses_expert_file_fallback_for_missing_master_rank(tmp_path): + param_name = "layer.experts.w1" + placement = make_autoep_placement_descriptor(4, [[2, 0, 3], [1]]) + model_tensor = torch.arange(8, dtype=torch.float32).view(4, 2) + master_tensor = model_tensor + 10 + fragments_dir = str(tmp_path / "fragments") + output_dir = str(tmp_path / "output") + fallback_dir = str(tmp_path / "fallback") + _write_zero12_fragments(fragments_dir, param_name, {0: master_tensor[[2, 0, 3]], 1: model_tensor[1:2]}) + os.remove(_zero12_fragment_path(fragments_dir, param_name, "fp32", 1)) + + os.makedirs(os.path.join(fallback_dir, param_name)) + torch.save({PARAM: model_tensor}, os.path.join(fallback_dir, param_name, "fp32.pt")) + info = { + param_name: { + 'num_experts': 4, + 'num_local_experts': 2, + 'ep_size': 2, + 'expert_placement': placement, + } + } + slice_shapes = {param_name: (2, 2)} + + fallback_names = get_autoep_zero12_fp32_fallback_param_names(fragments_dir, info, slice_shapes, 2, False) + assert fallback_names == {param_name} + + consolidate_autoep_zero12_expert_states(fragments_dir, + output_dir, + info, + slice_shapes, + 2, + 1, + False, + fp32_fallback_dir=fallback_dir) + + expected = model_tensor.clone() + expected[[2, 0, 3]] = master_tensor[[2, 0, 3]] + assert torch.equal(_load_consolidated(output_dir, param_name)[PARAM], expected) + + +def test_expert_file_consolidation_can_skip_fp32_output(tmp_path): + checkpoint_dir = tmp_path / "checkpoint" + output_dir = tmp_path / "output" + checkpoint_dir.mkdir() + for expert_id in range(2): + state = { + f"layer.experts.{weight}.{expert_id}": torch.full((2, 2), expert_id, dtype=torch.float32) + for weight in ("w1", "w2", "w3") + } + torch.save(state, checkpoint_dir / f"layer_0_expert_{expert_id}_mp_rank_00_model_states.pt") + + consolidate_autoep_expert_files( + str(checkpoint_dir), + str(output_dir), + [{ + "moe_layer_id": 0, + "num_experts": 2, + "expert_key_prefix": "layer.experts" + }], + fp32_fallback_param_names=set(), + ) + + assert not (output_dir / "zero" / "layer.experts.w1" / "fp32.pt").exists() + + def test_zero12_affine_consolidation_validates_replicated_experts(tmp_path): param_name = "layer.experts.w2" placement = make_autoep_placement_descriptor(3, [[0, 1], [1, 2]]) From 87a00c9df5cc605175d6068d216b06aeb4890650 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Thu, 24 Sep 2026 05:01:11 +0000 Subject: [PATCH 6/6] fix null module state in universal conversion Signed-off-by: Jin, Youzhi --- deepspeed/checkpoint/ds_to_universal.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py index 3d7a1668a2c6..99ddd1206641 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/deepspeed/checkpoint/ds_to_universal.py @@ -1214,7 +1214,11 @@ def main(args): zero_output_folder = os.path.join(args.output_folder, "zero") for model_file in ds_checkpoint.mp_rank_files: model_state = torch.load(model_file, map_location="cpu", weights_only=False) - for name, tensor in model_state.get("module", {}).items(): + module_state = model_state.get("module") + if module_state is None: + # Pipeline checkpoints store model weights in separate layer files. + continue + for name, tensor in module_state.items(): param_dir = os.path.join(zero_output_folder, name) fp32_path = os.path.join(param_dir, "fp32.pt") if os.path.isfile(fp32_path):