wehub-resource-sync
388 lines
15 KiB
Python
388 lines
15 KiB
Python
# Copyright (c) DeepSpeed Team.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""AutoEP universal checkpoint conversion utilities.
|
|
|
|
Consolidates per-expert checkpoint files (and their optimizer states) into
|
|
topology-agnostic universal format for EP resharding support.
|
|
"""
|
|
|
|
import os
|
|
import glob
|
|
import torch
|
|
|
|
from .constants import (
|
|
PARAM,
|
|
CAT_DIM,
|
|
EP_IS_EXPERT_PARAM,
|
|
EP_NUM_EXPERTS,
|
|
FOLDING_METADATA_KEY,
|
|
FOLDING_METADATA_VERSION,
|
|
FOLDING_TP_SIZE,
|
|
FOLDING_TP_RANK,
|
|
FOLDING_EP_SIZE,
|
|
FOLDING_EP_RANK,
|
|
FOLDING_ETP_SIZE,
|
|
FOLDING_ETP_RANK,
|
|
FOLDING_ZERO_PARTITION_GROUP,
|
|
FOLDING_ZERO_PARTITION_RANK,
|
|
FOLDING_ZERO_PARTITION_COUNT,
|
|
FOLDING_DISPATCH_STRATEGY,
|
|
FOLDING_SHARED_EXPERT_PLACEMENT,
|
|
FOLDING_FAMILY,
|
|
FOLDING_PARAM_FAMILIES,
|
|
)
|
|
|
|
|
|
def make_folding_metadata(*,
|
|
tp_size,
|
|
tp_rank,
|
|
ep_size,
|
|
ep_rank,
|
|
zero_partition_group,
|
|
zero_partition_rank,
|
|
zero_partition_count,
|
|
family,
|
|
param_families=None):
|
|
metadata = {
|
|
"version": FOLDING_METADATA_VERSION,
|
|
FOLDING_TP_SIZE: tp_size,
|
|
FOLDING_TP_RANK: tp_rank,
|
|
FOLDING_EP_SIZE: ep_size,
|
|
FOLDING_EP_RANK: ep_rank,
|
|
FOLDING_ETP_SIZE: 1,
|
|
FOLDING_ETP_RANK: 0,
|
|
FOLDING_ZERO_PARTITION_GROUP: zero_partition_group,
|
|
FOLDING_ZERO_PARTITION_RANK: zero_partition_rank,
|
|
FOLDING_ZERO_PARTITION_COUNT: zero_partition_count,
|
|
FOLDING_DISPATCH_STRATEGY: "route_full_partition_dispatch",
|
|
FOLDING_SHARED_EXPERT_PLACEMENT: "tp_sharded",
|
|
FOLDING_FAMILY: family,
|
|
}
|
|
if param_families is not None:
|
|
metadata[FOLDING_PARAM_FAMILIES] = dict(param_families)
|
|
return metadata
|
|
|
|
|
|
def validate_folding_metadata(metadata,
|
|
*,
|
|
tp_size,
|
|
ep_size,
|
|
etp_size=1,
|
|
tp_rank=None,
|
|
ep_rank=None,
|
|
etp_rank=None,
|
|
zero_partition_group=None,
|
|
zero_partition_rank=None,
|
|
zero_partition_count=None,
|
|
family=None,
|
|
param_families=None,
|
|
shared_expert_placement=None,
|
|
dispatch_strategy=None):
|
|
if not isinstance(metadata, dict) or FOLDING_METADATA_KEY not in metadata:
|
|
raise RuntimeError("Missing AutoEP+AutoTP folding metadata in folded checkpoint.")
|
|
folding = metadata[FOLDING_METADATA_KEY]
|
|
if folding.get("version") != FOLDING_METADATA_VERSION:
|
|
raise RuntimeError(f"Unsupported folding metadata version: {folding.get('version')}")
|
|
expected = {
|
|
FOLDING_TP_SIZE: tp_size,
|
|
FOLDING_EP_SIZE: ep_size,
|
|
FOLDING_ETP_SIZE: etp_size,
|
|
}
|
|
optional_expected = {
|
|
FOLDING_TP_RANK: tp_rank,
|
|
FOLDING_EP_RANK: ep_rank,
|
|
FOLDING_ETP_RANK: etp_rank,
|
|
FOLDING_ZERO_PARTITION_GROUP: zero_partition_group,
|
|
FOLDING_ZERO_PARTITION_RANK: zero_partition_rank,
|
|
FOLDING_ZERO_PARTITION_COUNT: zero_partition_count,
|
|
FOLDING_FAMILY: family,
|
|
FOLDING_PARAM_FAMILIES: param_families,
|
|
FOLDING_SHARED_EXPERT_PLACEMENT: shared_expert_placement,
|
|
FOLDING_DISPATCH_STRATEGY: dispatch_strategy,
|
|
}
|
|
expected.update({key: value for key, value in optional_expected.items() if value is not None})
|
|
for key, value in expected.items():
|
|
if folding.get(key) != value:
|
|
raise RuntimeError(f"Folding metadata mismatch for {key}: saved={folding.get(key)} runtime={value}")
|
|
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.
|
|
|
|
Resolves using glob pattern without assuming mp_rank=0.
|
|
|
|
Returns:
|
|
Path to the single matching expert checkpoint file.
|
|
|
|
Raises:
|
|
FileNotFoundError: No matching file found.
|
|
NotImplementedError: Multiple matching files found (multi-mp_rank).
|
|
"""
|
|
pattern = os.path.join(checkpoint_dir, f'layer_{moe_layer_id}_expert_{global_expert_id}_mp_rank_*_model_states.pt')
|
|
matches = glob.glob(pattern)
|
|
if len(matches) == 0:
|
|
raise FileNotFoundError(f"Expert checkpoint file not found: layer_{moe_layer_id} "
|
|
f"expert_{global_expert_id} in {checkpoint_dir}")
|
|
if len(matches) > 1:
|
|
for match in matches:
|
|
state = torch.load(match, map_location='cpu', weights_only=False)
|
|
if FOLDING_METADATA_KEY in state:
|
|
raise NotImplementedError("Universal checkpoint conversion for folded AutoEP+AutoTP expert shards "
|
|
"is not supported yet. Load this checkpoint with a matching folded "
|
|
"runtime, or consolidate the tensor-parallel expert shards before "
|
|
"running ds_to_universal.")
|
|
raise NotImplementedError(f"Multiple expert checkpoint files found for layer_{moe_layer_id} "
|
|
f"expert_{global_expert_id}: {matches}. Multi-mp_rank expert files "
|
|
f"are not yet supported.")
|
|
return matches[0]
|
|
|
|
|
|
def consolidate_autoep_expert_files(checkpoint_dir, output_dir, autoep_layers_metadata):
|
|
"""Consolidate per-expert checkpoint files into full-expert universal format.
|
|
|
|
For each AutoEP layer, loads all per-expert files, stacks into
|
|
[E_total, H, D] tensors, and saves in universal checkpoint 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.
|
|
|
|
Raises:
|
|
FileNotFoundError: If expected expert files are missing.
|
|
NotImplementedError: If multiple mp_rank files match one (layer, expert).
|
|
RuntimeError: If metadata is missing or malformed.
|
|
"""
|
|
if autoep_layers_metadata is None:
|
|
raise RuntimeError("AutoEP metadata is missing from checkpoint. Cannot consolidate "
|
|
"expert files without ds_autoep_layers metadata.")
|
|
if not isinstance(autoep_layers_metadata, list):
|
|
raise RuntimeError(f"AutoEP metadata is malformed: expected list, got "
|
|
f"{type(autoep_layers_metadata).__name__}")
|
|
|
|
for layer_info in autoep_layers_metadata:
|
|
moe_layer_id = layer_info['moe_layer_id']
|
|
num_experts = layer_info['num_experts']
|
|
prefix = layer_info['expert_key_prefix']
|
|
|
|
for wname in ('w1', 'w2', 'w3'):
|
|
expert_tensors = []
|
|
folding_metadata = None
|
|
for global_eid in range(num_experts):
|
|
ckpt_path = resolve_expert_ckpt_path(checkpoint_dir, moe_layer_id, global_eid)
|
|
sd = torch.load(ckpt_path, map_location='cpu', weights_only=False)
|
|
if folding_metadata is None:
|
|
folding_metadata = sd.get(FOLDING_METADATA_KEY)
|
|
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])
|
|
|
|
# Stack to full fused tensor [E_total, H, D]
|
|
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)
|
|
os.makedirs(param_dir, exist_ok=True)
|
|
universal_state = {
|
|
PARAM: full_tensor,
|
|
CAT_DIM: 0,
|
|
EP_IS_EXPERT_PARAM: True,
|
|
EP_NUM_EXPERTS: num_experts,
|
|
}
|
|
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:
|
|
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"))
|