wehub-resource-sync
600 lines
27 KiB
Python
600 lines
27 KiB
Python
# Copyright (c) DeepSpeed Team.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""AutoEP: Automatic Expert Parallelism for MoE models.
|
|
|
|
Phase 3: MoE layer detection and structural validation.
|
|
Phase 5: Layer replacement (replace_moe_layer filled in).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import re
|
|
from collections import OrderedDict
|
|
from typing import Literal
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from deepspeed.module_inject.auto_ep_config import (
|
|
fill_autoep_config_from_hf,
|
|
AutoEPConfig,
|
|
MoELayerSpec,
|
|
MoEModelPreset,
|
|
)
|
|
from deepspeed.module_inject.auto_ep_presets.base import ForwardContract
|
|
from deepspeed.module_inject.auto_ep_presets.registry import (
|
|
apply_config_overrides,
|
|
get_preset_adapter,
|
|
preset_name_for_hf_model_type,
|
|
resolve_preset_candidates,
|
|
unsupported_preset_for_hf_model_type,
|
|
)
|
|
from deepspeed.moe.fused_expert_layout import classify_fused_gate_up_layout
|
|
from deepspeed.runtime.zero.utils import is_zero_param
|
|
from deepspeed.utils import logger
|
|
|
|
|
|
def _remove_transformers_output_capture_hooks(model: nn.Module) -> int:
|
|
"""Remove HF output-capturing hooks so they can be reinstalled after AutoEP conversion."""
|
|
removed = 0
|
|
for module in model.modules():
|
|
hooks = getattr(module, "_forward_hooks", None)
|
|
if not hooks:
|
|
continue
|
|
|
|
for hook_id, hook in list(hooks.items()):
|
|
if getattr(hook, "__name__", "") != "output_capturing_hook":
|
|
continue
|
|
del hooks[hook_id]
|
|
removed += 1
|
|
hooks_with_kwargs = getattr(module, "_forward_hooks_with_kwargs", None)
|
|
if hooks_with_kwargs is not None:
|
|
hooks_with_kwargs.pop(hook_id, None)
|
|
hooks_always_called = getattr(module, "_forward_hooks_always_called", None)
|
|
if hooks_always_called is not None:
|
|
hooks_always_called.pop(hook_id, None)
|
|
return removed
|
|
|
|
|
|
def _is_known_hf_model_type(model_type: str | None) -> bool:
|
|
if model_type is None:
|
|
return False
|
|
return (preset_name_for_hf_model_type(model_type) is not None
|
|
or unsupported_preset_for_hf_model_type(model_type) is not None)
|
|
|
|
|
|
def _raise_if_duplicate_moe_specs(specs: list[MoELayerSpec]) -> None:
|
|
by_module: dict[str, list[MoELayerSpec]] = {}
|
|
for spec in specs:
|
|
by_module.setdefault(spec.moe_module_name, []).append(spec)
|
|
|
|
duplicates = {name: matches for name, matches in by_module.items() if len(matches) > 1}
|
|
if not duplicates:
|
|
return
|
|
|
|
details = "; ".join(f"{name}: {', '.join(spec.model_family for spec in matches)}"
|
|
for name, matches in sorted(duplicates.items()))
|
|
raise ValueError("AutoEP detection is ambiguous and produced multiple replacement specs for the same "
|
|
f"MoE module(s): {details}. Set expert_parallel.preset_model or provide custom "
|
|
"AutoEP patterns so each MoE module matches exactly one preset.")
|
|
|
|
|
|
def _source_param_shape(param: torch.Tensor | nn.Parameter) -> torch.Size:
|
|
if is_zero_param(param):
|
|
return torch.Size(param.ds_shape)
|
|
return torch.Size(param.shape)
|
|
|
|
|
|
def _source_param_ndim(param: torch.Tensor | nn.Parameter) -> int:
|
|
return len(_source_param_shape(param))
|
|
|
|
|
|
def _has_3d_expert_params(module: nn.Module, preset: MoEModelPreset) -> bool:
|
|
"""Check if module stores expert weights as 3D parameter tensors (transformers 5.0.0+).
|
|
|
|
Returns True if the module has a parameter named preset.expert_w1 (e.g., "gate_up_proj")
|
|
with 3 dimensions (num_experts, ..., ...).
|
|
"""
|
|
w1_name = preset.expert_w1
|
|
param = getattr(module, w1_name, None)
|
|
if param is None:
|
|
return False
|
|
if isinstance(param, nn.Parameter) or isinstance(param, torch.Tensor):
|
|
return _source_param_ndim(param) == 3
|
|
return False
|
|
|
|
|
|
def _get_num_experts_from_config(model_config, preset: MoEModelPreset) -> int | None:
|
|
"""Extract num_experts from model.config using the preset's attribute name."""
|
|
return getattr(model_config, preset.num_experts_attr, None)
|
|
|
|
|
|
def _get_top_k_from_config(model_config, preset: MoEModelPreset) -> int | None:
|
|
"""Extract top_k from model.config using the preset's attribute name."""
|
|
return getattr(model_config, preset.top_k_attr, None)
|
|
|
|
|
|
def _as_finite_float(value, field_name: str) -> float:
|
|
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
|
raise ValueError(f"{field_name} must be a finite number")
|
|
|
|
value = float(value)
|
|
if not math.isfinite(value):
|
|
raise ValueError(f"{field_name} must be a finite number")
|
|
return value
|
|
|
|
|
|
def _resolve_route_scale(config: AutoEPConfig, model_config) -> float:
|
|
"""Resolve the single scale applied by TokenChoiceTopKRouter."""
|
|
routed_scaling_factor = config.routed_scaling_factor
|
|
|
|
if routed_scaling_factor != "auto":
|
|
route_scale = _as_finite_float(routed_scaling_factor, "routed_scaling_factor")
|
|
if config.route_scale != 1.0:
|
|
logger.warning("AutoEP: routed_scaling_factor=%s overrides route_scale=%s.", routed_scaling_factor,
|
|
config.route_scale)
|
|
return route_scale
|
|
|
|
cfg_routed_scaling_factor = getattr(model_config, 'routed_scaling_factor', None)
|
|
if cfg_routed_scaling_factor is not None:
|
|
route_scale = _as_finite_float(cfg_routed_scaling_factor, "model.config.routed_scaling_factor")
|
|
if config.route_scale != 1.0:
|
|
logger.warning("AutoEP: model.config.routed_scaling_factor=%s overrides route_scale=%s.",
|
|
cfg_routed_scaling_factor, config.route_scale)
|
|
return route_scale
|
|
|
|
return _as_finite_float(config.route_scale, "route_scale")
|
|
|
|
|
|
def _detect_expert_storage(experts_module: nn.Module, preset: MoEModelPreset) -> Literal["fused_3d", "module_list"]:
|
|
"""Determine whether experts are stored as fused 3D tensors or nn.ModuleList."""
|
|
if _has_3d_expert_params(experts_module, preset):
|
|
return "fused_3d"
|
|
if isinstance(experts_module, nn.ModuleList):
|
|
return "module_list"
|
|
# Check children for 3D params as fallback
|
|
for name, param in experts_module.named_parameters(recurse=False):
|
|
if _source_param_ndim(param) == 3:
|
|
return "fused_3d"
|
|
return "module_list"
|
|
|
|
|
|
def _infer_hidden_and_ffn_size(
|
|
experts_module: nn.Module,
|
|
preset: MoEModelPreset,
|
|
storage: Literal["fused_3d", "module_list"],
|
|
num_experts: int,
|
|
) -> tuple[int, int]:
|
|
"""Infer hidden_size and ffn_hidden_size from expert weight shapes."""
|
|
if storage == "fused_3d":
|
|
w1_param = getattr(experts_module, preset.expert_w1, None)
|
|
w2_param = getattr(experts_module, preset.expert_w2, None)
|
|
if w1_param is not None and w2_param is not None:
|
|
w1_shape = _source_param_shape(w1_param)
|
|
w2_shape = _source_param_shape(w2_param)
|
|
if preset.expert_w3 is None:
|
|
layout = classify_fused_gate_up_layout(tuple(w1_shape), tuple(w2_shape))
|
|
if layout is None:
|
|
raise ValueError("expert_w3=None expects fused gate+up weights with either "
|
|
f"[E, 2*ffn, hidden]/[E, hidden, ffn] or [E, hidden, 2*ffn]/[E, ffn, hidden], "
|
|
f"but got {preset.expert_w1}={tuple(w1_shape)} and "
|
|
f"{preset.expert_w2}={tuple(w2_shape)}.")
|
|
hidden_size = layout.hidden_size
|
|
ffn_hidden_size = layout.ffn_hidden_size
|
|
else:
|
|
# Separate gate and up: w1 shape is [E, ffn, hidden]
|
|
w3_param = getattr(experts_module, preset.expert_w3, None)
|
|
if w3_param is None:
|
|
raise ValueError(f"expert_w3='{preset.expert_w3}' is set but no such weight "
|
|
f"exists on experts module.")
|
|
hidden_size = w1_shape[2]
|
|
ffn_hidden_size = w1_shape[1]
|
|
return hidden_size, ffn_hidden_size
|
|
elif storage == "module_list":
|
|
# Legacy: individual expert modules
|
|
if isinstance(experts_module, nn.ModuleList) and len(experts_module) > 0:
|
|
expert0 = experts_module[0]
|
|
w1 = getattr(expert0, preset.expert_w1, None)
|
|
if w1 is None:
|
|
# Try weight attribute for nn.Linear
|
|
for name, child in expert0.named_children():
|
|
if preset.expert_w1 in name:
|
|
w1 = child.weight if hasattr(child, 'weight') else None
|
|
break
|
|
if w1 is not None:
|
|
if isinstance(w1, nn.Linear):
|
|
return w1.in_features, w1.out_features
|
|
elif isinstance(w1, (nn.Parameter, torch.Tensor)):
|
|
w1_shape = _source_param_shape(w1)
|
|
if len(w1_shape) == 2:
|
|
return w1_shape[1], w1_shape[0]
|
|
|
|
raise ValueError(f"Could not infer hidden_size/ffn_hidden_size from experts module "
|
|
f"with storage={storage}, preset.expert_w1={preset.expert_w1}")
|
|
|
|
|
|
def _detect_forward_contract(
|
|
moe_module: nn.Module,
|
|
router_module: nn.Module,
|
|
) -> ForwardContract:
|
|
"""Detect the forward contract for router logits capture.
|
|
|
|
Returns:
|
|
ForwardContract with router-logit return and capture metadata.
|
|
"""
|
|
# Check for OutputRecorder on the model (transformers 5.0.0 pattern)
|
|
# Look for _can_record_outputs attribute on parent modules
|
|
capture_target: Literal["moe_block", "router", "none"] = "none"
|
|
capture_index: int | None = None
|
|
capture_layer_name: str | None = None
|
|
return_router_logits = False
|
|
|
|
# Check for OutputRecorder pattern on router class
|
|
router_class = type(router_module)
|
|
if hasattr(router_class, '_can_record_outputs'):
|
|
capture_target = "router"
|
|
record_config = router_class._can_record_outputs
|
|
if isinstance(record_config, dict):
|
|
for key, val in record_config.items():
|
|
if isinstance(val, dict):
|
|
capture_index = val.get('index', 0)
|
|
capture_layer_name = val.get('layer_name', None)
|
|
else:
|
|
capture_index = 0
|
|
elif isinstance(record_config, (list, tuple)):
|
|
capture_index = 0
|
|
logger.debug(f"Detected OutputRecorder on router class {router_class.__name__}: "
|
|
f"index={capture_index}, layer_name={capture_layer_name}")
|
|
|
|
# Check if MoE block has tuple return contract (legacy transformers)
|
|
if hasattr(moe_module, '_can_record_outputs'):
|
|
record_config = moe_module._can_record_outputs
|
|
if record_config:
|
|
capture_target = "moe_block"
|
|
return_router_logits = True
|
|
if isinstance(record_config, dict):
|
|
for key, val in record_config.items():
|
|
if isinstance(val, dict):
|
|
capture_index = val.get('index', None)
|
|
elif isinstance(val, int):
|
|
capture_index = val
|
|
|
|
return ForwardContract(
|
|
return_router_logits=return_router_logits,
|
|
capture_target=capture_target,
|
|
capture_index=capture_index,
|
|
capture_layer_name=capture_layer_name,
|
|
)
|
|
|
|
|
|
class AutoEP:
|
|
"""Automatic Expert Parallelism: detect and replace MoE layers."""
|
|
|
|
def __init__(self, model: nn.Module, config: AutoEPConfig) -> None:
|
|
self.model = model
|
|
self.config = config
|
|
self.model_config = getattr(model, 'config', None)
|
|
self._retargeted_transformers_output_recorders: set[str] = set()
|
|
fill_autoep_config_from_hf(self.config, self.model_config)
|
|
|
|
def ep_parser(self) -> list[MoELayerSpec]:
|
|
"""Traverse model and detect MoE layers. Returns list of MoELayerSpec."""
|
|
specs = []
|
|
|
|
# Determine which preset(s) to use
|
|
presets_to_try = self._resolve_presets()
|
|
|
|
for preset_name, preset in presets_to_try:
|
|
adapter = get_preset_adapter(preset.preset_adapter)
|
|
pattern = re.compile(preset.moe_layer_pattern)
|
|
|
|
for module_name, module in self.model.named_modules():
|
|
if not pattern.fullmatch(module_name):
|
|
continue
|
|
|
|
# Structural validation: check for experts child
|
|
experts_child = getattr(module, preset.experts_pattern, None)
|
|
if experts_child is None:
|
|
logger.debug(
|
|
"Skipping %s: pattern matched but no '%s' child (likely dense FFN)",
|
|
module_name,
|
|
preset.experts_pattern,
|
|
)
|
|
continue
|
|
|
|
expert_layout = adapter.resolve_expert_layout(experts_child, preset)
|
|
|
|
# Accept both: nn.ModuleList (legacy) and Experts class (transformers 5.0.0+)
|
|
has_expert_params = (isinstance(experts_child, nn.ModuleList)
|
|
or _has_3d_expert_params(experts_child, expert_layout))
|
|
if not has_expert_params:
|
|
logger.debug(
|
|
"Skipping %s: '%s' child exists but has no expert parameters",
|
|
module_name,
|
|
preset.experts_pattern,
|
|
)
|
|
continue
|
|
|
|
# Check for router
|
|
router_child = getattr(module, preset.router_pattern, None)
|
|
if router_child is None:
|
|
logger.debug(
|
|
"Skipping %s: no router child '%s'",
|
|
module_name,
|
|
preset.router_pattern,
|
|
)
|
|
continue
|
|
|
|
# Detect storage format
|
|
storage = _detect_expert_storage(experts_child, expert_layout)
|
|
|
|
# Get num_experts and top_k from config or weights
|
|
num_experts = None
|
|
top_k = None
|
|
|
|
if self.model_config is not None:
|
|
num_experts = _get_num_experts_from_config(self.model_config, preset)
|
|
top_k = _get_top_k_from_config(self.model_config, preset)
|
|
|
|
# Validate/derive from router weight shape
|
|
router_weight = getattr(router_child, 'weight', None)
|
|
router_weight_shape = _source_param_shape(router_weight) if router_weight is not None else None
|
|
if router_weight_shape is not None and len(router_weight_shape) == 2:
|
|
num_experts_from_weight = router_weight_shape[0]
|
|
hidden_from_weight = router_weight_shape[1]
|
|
if num_experts is not None and num_experts != num_experts_from_weight:
|
|
raise ValueError(f"Config num_experts={num_experts} mismatches router weight "
|
|
f"shape {router_weight_shape} (expected {num_experts_from_weight}) "
|
|
f"in layer '{module_name}'")
|
|
num_experts = num_experts_from_weight
|
|
|
|
if num_experts is None:
|
|
raise ValueError(f"Could not determine num_experts for layer '{module_name}'. "
|
|
f"Set model.config.{preset.num_experts_attr} or use a preset.")
|
|
|
|
# Override top_k from config if user specified
|
|
if isinstance(self.config.top_k, int):
|
|
top_k = self.config.top_k
|
|
elif top_k is None:
|
|
raise ValueError(f"Could not determine top_k for layer '{module_name}'. "
|
|
f"Set model.config.{preset.top_k_attr} or config top_k.")
|
|
|
|
# Infer hidden sizes
|
|
try:
|
|
hidden_size, ffn_hidden_size = _infer_hidden_and_ffn_size(experts_child, expert_layout, storage,
|
|
num_experts)
|
|
except ValueError as e:
|
|
if self._requires_selected_preset_detection():
|
|
raise ValueError(f"AutoEP: preset '{preset_name}' matched layer '{module_name}' "
|
|
f"with router and experts, but shape inference failed: {e}") from e
|
|
logger.warning(f"Skipping {module_name}: {e}")
|
|
continue
|
|
|
|
# Cross-validate hidden_size with router
|
|
if router_weight_shape is not None and len(router_weight_shape) == 2:
|
|
if hidden_size != router_weight_shape[1]:
|
|
raise ValueError(f"hidden_size={hidden_size} from expert weights mismatches "
|
|
f"router weight dim={router_weight_shape[1]} in '{module_name}'")
|
|
|
|
# Validate top_k <= num_experts
|
|
if top_k > num_experts:
|
|
raise ValueError(f"top_k={top_k} exceeds num_experts={num_experts} "
|
|
f"in layer '{module_name}'")
|
|
|
|
# Resolve score_func
|
|
if self.config.score_func != "auto":
|
|
score_func = self.config.score_func
|
|
else:
|
|
# Check model config for scoring_func attribute
|
|
cfg_score = getattr(self.model_config, 'scoring_func', None)
|
|
if cfg_score in ("softmax", "sigmoid"):
|
|
score_func = cfg_score
|
|
else:
|
|
score_func = preset.score_func
|
|
|
|
# Resolve score_apply
|
|
if self.config.score_apply != "auto":
|
|
score_apply = self.config.score_apply
|
|
else:
|
|
score_apply = preset.score_apply
|
|
|
|
route_norm = adapter.resolve_route_norm(self.config, preset, self.model_config)
|
|
|
|
route_scale = _resolve_route_scale(self.config, self.model_config)
|
|
|
|
group_routing = adapter.resolve_group_routing(self.config, self.model_config)
|
|
|
|
# Check gate bias
|
|
gate_bias = preset.gate_bias
|
|
if router_weight is not None:
|
|
gate_bias = getattr(router_child, 'bias', None) is not None
|
|
|
|
forward_contract = adapter.adjust_forward_contract(_detect_forward_contract(module, router_child))
|
|
|
|
# Check shared experts
|
|
has_shared = False
|
|
shared_name = ""
|
|
shared_gate_name = ""
|
|
if preset.has_shared_experts and preset.shared_experts_pattern:
|
|
shared = getattr(module, preset.shared_experts_pattern, None)
|
|
if shared is not None:
|
|
has_shared = True
|
|
shared_name = preset.shared_experts_pattern
|
|
if preset.shared_experts_gate_pattern:
|
|
shared_gate = getattr(module, preset.shared_experts_gate_pattern, None)
|
|
if shared_gate is not None:
|
|
shared_gate_name = preset.shared_experts_gate_pattern
|
|
|
|
# Warn about router stochasticity/precision settings
|
|
if self.model_config is not None:
|
|
jitter = getattr(self.model_config, 'router_jitter_noise', 0.0)
|
|
if jitter and jitter > 0:
|
|
logger.warning(f"Layer {module_name}: model has router_jitter_noise={jitter}, "
|
|
f"AutoEP router does not implement jitter.")
|
|
z_loss = getattr(self.model_config, 'router_z_loss_coef', 0.0)
|
|
if z_loss and z_loss > 0:
|
|
logger.warning(f"Layer {module_name}: model has router_z_loss_coef={z_loss}, "
|
|
f"AutoEP router does not implement z-loss.")
|
|
|
|
spec = MoELayerSpec(
|
|
moe_module_name=module_name,
|
|
model_family=preset_name,
|
|
router_name=preset.router_pattern,
|
|
experts_name=preset.experts_pattern,
|
|
expert_storage=storage,
|
|
expert_w1_name=expert_layout.expert_w1,
|
|
expert_w2_name=expert_layout.expert_w2,
|
|
expert_w3_name=expert_layout.expert_w3,
|
|
num_experts=num_experts,
|
|
top_k=top_k,
|
|
hidden_size=hidden_size,
|
|
ffn_hidden_size=ffn_hidden_size,
|
|
score_func=score_func,
|
|
score_apply=score_apply,
|
|
route_norm=route_norm,
|
|
gate_bias=gate_bias,
|
|
return_router_logits=forward_contract.return_router_logits,
|
|
router_logits_capture_target=forward_contract.capture_target,
|
|
router_logits_capture_index=forward_contract.capture_index,
|
|
router_logits_capture_layer_name=forward_contract.capture_layer_name,
|
|
has_shared_experts=has_shared,
|
|
shared_experts_name=shared_name,
|
|
shared_experts_gate_name=shared_gate_name,
|
|
route_scale=route_scale,
|
|
num_expert_groups=group_routing.num_expert_groups,
|
|
num_limited_groups=group_routing.num_limited_groups,
|
|
group_score_func=group_routing.group_score_func,
|
|
supports_expert_bias=preset.supports_expert_bias,
|
|
unsupported_router_bias_names=preset.unsupported_router_bias_names,
|
|
preset_adapter=preset.preset_adapter,
|
|
router_logits_capture_mode=forward_contract.router_logits_capture_mode,
|
|
moe_output_shape=forward_contract.moe_output_shape,
|
|
)
|
|
specs.append(spec)
|
|
logger.debug(f"Detected MoE layer: {module_name} (family={preset_name}, "
|
|
f"experts={num_experts}, top_k={top_k}, storage={storage})")
|
|
|
|
if not specs:
|
|
if self._requires_selected_preset_detection():
|
|
self._raise_no_moe_layers_detected(presets_to_try)
|
|
logger.warning("AutoEP: no MoE layers detected in model.")
|
|
else:
|
|
_raise_if_duplicate_moe_specs(specs)
|
|
|
|
return specs
|
|
|
|
def _replace_moe_layer_without_retarget(
|
|
self,
|
|
spec: MoELayerSpec,
|
|
ep_size: int,
|
|
ep_rank: int,
|
|
) -> nn.Module:
|
|
from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer
|
|
|
|
# Navigate to the parent module and get the child name
|
|
parts = spec.moe_module_name.split(".")
|
|
parent = self.model
|
|
for part in parts[:-1]:
|
|
parent = getattr(parent, part)
|
|
child_name = parts[-1]
|
|
source_module = getattr(parent, child_name)
|
|
|
|
# Create replacement layer
|
|
replacement = AutoEPMoELayer(
|
|
spec=spec,
|
|
source_module=source_module,
|
|
ep_size=ep_size,
|
|
ep_rank=ep_rank,
|
|
config=self.config,
|
|
)
|
|
|
|
# Replace in-place on parent
|
|
setattr(parent, child_name, replacement)
|
|
return replacement
|
|
|
|
def _retarget_transformers_output_recorders(self, spec: MoELayerSpec, replacement: nn.Module) -> None:
|
|
adapter = get_preset_adapter(spec.preset_adapter)
|
|
adapter.retarget_transformers_output_recorders(
|
|
self.model,
|
|
spec,
|
|
replacement,
|
|
self._retargeted_transformers_output_recorders,
|
|
_remove_transformers_output_capture_hooks,
|
|
)
|
|
|
|
def replace_moe_layer(
|
|
self,
|
|
spec: MoELayerSpec,
|
|
ep_size: int,
|
|
ep_rank: int,
|
|
) -> None:
|
|
"""Replace a single MoE module with AutoEPMoELayer in-place on the model."""
|
|
replacement = self._replace_moe_layer_without_retarget(spec, ep_size, ep_rank)
|
|
self._retarget_transformers_output_recorders(spec, replacement)
|
|
|
|
logger.info(f"AutoEP: replaced '{spec.moe_module_name}' with AutoEPMoELayer "
|
|
f"(ep_size={ep_size}, ep_rank={ep_rank}, "
|
|
f"local_experts={replacement.num_local_experts})")
|
|
|
|
def replace_moe_layers(
|
|
self,
|
|
specs: list[MoELayerSpec],
|
|
ep_size: int,
|
|
ep_rank: int,
|
|
) -> None:
|
|
"""Replace multiple MoE modules and batch post-replacement recorder retargeting."""
|
|
replacements: list[tuple[MoELayerSpec, nn.Module]] = []
|
|
for spec in specs:
|
|
replacement = self._replace_moe_layer_without_retarget(spec, ep_size, ep_rank)
|
|
replacements.append((spec, replacement))
|
|
logger.info(f"AutoEP: replaced '{spec.moe_module_name}' with AutoEPMoELayer "
|
|
f"(ep_size={ep_size}, ep_rank={ep_rank}, "
|
|
f"local_experts={replacement.num_local_experts})")
|
|
|
|
retarget_groups: OrderedDict[tuple[str, str, type], tuple[MoELayerSpec, nn.Module]] = OrderedDict()
|
|
for spec, replacement in replacements:
|
|
retarget_key = (spec.preset_adapter, spec.model_family, replacement.__class__)
|
|
retarget_groups.setdefault(retarget_key, (spec, replacement))
|
|
|
|
for spec, replacement in retarget_groups.values():
|
|
self._retarget_transformers_output_recorders(spec, replacement)
|
|
|
|
def _apply_config_overrides(self, preset: MoEModelPreset) -> MoEModelPreset:
|
|
return apply_config_overrides(self.config, preset)
|
|
|
|
def _requires_selected_preset_detection(self) -> bool:
|
|
"""Return whether empty detection should fail for the selected preset."""
|
|
if self.config.preset_model is not None:
|
|
return True
|
|
if self.config.moe_layer_pattern is not None:
|
|
return True
|
|
if self.model_config is None:
|
|
return False
|
|
model_type = getattr(self.model_config, 'model_type', None)
|
|
return _is_known_hf_model_type(model_type)
|
|
|
|
def _raise_no_moe_layers_detected(self, presets_to_try: list[tuple[str, MoEModelPreset]]) -> None:
|
|
model_type = getattr(self.model_config, 'model_type', None)
|
|
if self.config.preset_model is not None:
|
|
source = f"preset_model='{self.config.preset_model}'"
|
|
elif self.config.moe_layer_pattern is not None:
|
|
source = f"moe_layer_pattern='{self.config.moe_layer_pattern}'"
|
|
else:
|
|
source = f"model_type='{model_type}'"
|
|
|
|
expected = "; ".join(f"{preset_name}: moe_layer_pattern='{preset.moe_layer_pattern}', "
|
|
f"router='{preset.router_pattern}', experts='{preset.experts_pattern}'"
|
|
for preset_name, preset in presets_to_try)
|
|
raise ValueError(f"AutoEP: no MoE layers detected for {source}. "
|
|
f"Expected MoE structure for selected preset(s): {expected}. "
|
|
"This usually means the selected preset does not match the model implementation, "
|
|
"or the installed Transformers version exposes a different structure. Choose a matching "
|
|
"preset, upgrade Transformers, or provide custom AutoEP patterns.")
|
|
|
|
def _resolve_presets(self) -> list[tuple[str, MoEModelPreset]]:
|
|
"""Determine which preset(s) to use for detection."""
|
|
return resolve_preset_candidates(self.config, self.model_config)
|