wehub-resource-sync
212 lines
8.3 KiB
Python
212 lines
8.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# DeepSpeed Team
|
|
"""AutoEP preset registry and config override helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
from dataclasses import replace
|
|
|
|
from deepspeed.module_inject.auto_ep_presets.base import (
|
|
_UNSET,
|
|
AutoEPConfig,
|
|
AutoEPPresetAdapter,
|
|
MoEModelPreset,
|
|
)
|
|
from deepspeed.module_inject.auto_ep_presets import deepseek_v2, deepseek_v3, mixtral, qwen3_5_moe, qwen3_moe
|
|
from deepspeed.utils import logger
|
|
|
|
_PRESET_MODULES = (
|
|
mixtral,
|
|
qwen3_moe,
|
|
qwen3_5_moe,
|
|
deepseek_v2,
|
|
deepseek_v3,
|
|
)
|
|
|
|
PRESET_MODELS: dict[str, MoEModelPreset] = {module.PRESET_NAME: module.PRESET for module in _PRESET_MODULES}
|
|
|
|
_PRESET_ADAPTERS: dict[str, AutoEPPresetAdapter] = {
|
|
"default": AutoEPPresetAdapter(),
|
|
}
|
|
for _module in _PRESET_MODULES:
|
|
_PRESET_ADAPTERS.update(getattr(_module, "PRESET_ADAPTERS", {}))
|
|
|
|
|
|
def _validate_registered_preset_adapters(
|
|
preset_models: dict[str, MoEModelPreset] | None = None,
|
|
preset_adapters: dict[str, AutoEPPresetAdapter] | None = None,
|
|
) -> None:
|
|
"""Fail fast if a registered preset references an adapter that is not registered."""
|
|
preset_models = PRESET_MODELS if preset_models is None else preset_models
|
|
preset_adapters = _PRESET_ADAPTERS if preset_adapters is None else preset_adapters
|
|
|
|
missing_presets = []
|
|
for preset_name, preset in preset_models.items():
|
|
if preset.preset_adapter not in preset_adapters:
|
|
missing_presets.append((preset_name, preset.preset_adapter))
|
|
|
|
if not missing_presets:
|
|
return
|
|
|
|
details = ", ".join(f"{preset_name}:{adapter_name}" for preset_name, adapter_name in missing_presets)
|
|
raise RuntimeError(f"AutoEP preset registry is inconsistent; missing preset_adapter registration(s): {details}")
|
|
|
|
|
|
_validate_registered_preset_adapters()
|
|
|
|
_PRESET_DEFAULT_EXPLICIT_FLAGS = {
|
|
"load_balance_coeff": "_load_balance_coeff_explicit",
|
|
}
|
|
|
|
|
|
def available_preset_names() -> tuple[str, ...]:
|
|
"""Return built-in AutoEP preset names in registry order."""
|
|
return tuple(PRESET_MODELS.keys())
|
|
|
|
|
|
def get_preset(preset_name: str) -> MoEModelPreset:
|
|
"""Return a registered AutoEP preset by name."""
|
|
preset = PRESET_MODELS.get(preset_name)
|
|
if preset is None:
|
|
raise ValueError(f"Unknown preset_model '{preset_name}'. Available presets: {list(available_preset_names())}")
|
|
return preset
|
|
|
|
|
|
def get_preset_adapter(adapter_name: str) -> AutoEPPresetAdapter:
|
|
"""Return a registered AutoEP preset adapter by name."""
|
|
adapter = _PRESET_ADAPTERS.get(adapter_name)
|
|
if adapter is None:
|
|
raise ValueError(f"Unknown AutoEP preset adapter '{adapter_name}'")
|
|
return adapter
|
|
|
|
|
|
def preset_name_for_hf_model_type(model_type: str) -> str | None:
|
|
"""Return the AutoEP preset name for a supported HF model_type."""
|
|
for preset_name, preset in PRESET_MODELS.items():
|
|
if model_type in preset.hf_model_types:
|
|
return preset_name
|
|
return None
|
|
|
|
|
|
def unsupported_preset_for_hf_model_type(model_type: str) -> tuple[str, MoEModelPreset] | None:
|
|
"""Return a preset carrying an actionable diagnostic for an unsupported HF model_type."""
|
|
for preset_name, preset in PRESET_MODELS.items():
|
|
if model_type in preset.unsupported_hf_model_type_notes:
|
|
return preset_name, preset
|
|
return None
|
|
|
|
|
|
def resolve_autoep_config_defaults(config: AutoEPConfig, preset_name: str | None) -> AutoEPConfig:
|
|
"""Return config with preset-level AutoEP defaults applied where the user did not override."""
|
|
if preset_name is None or preset_name not in PRESET_MODELS:
|
|
return config
|
|
|
|
preset_defaults = PRESET_MODELS[preset_name].autoep_config_defaults
|
|
if not preset_defaults:
|
|
return config
|
|
|
|
resolved = copy.copy(config)
|
|
for field_name, default_value in preset_defaults.items():
|
|
explicit_flag = _PRESET_DEFAULT_EXPLICIT_FLAGS.get(field_name)
|
|
if explicit_flag is None:
|
|
continue
|
|
if not getattr(config, explicit_flag, False):
|
|
setattr(resolved, field_name, default_value)
|
|
return resolved
|
|
|
|
|
|
def apply_config_overrides(config: AutoEPConfig, preset: MoEModelPreset) -> MoEModelPreset:
|
|
"""Apply explicit AutoEP config overrides to a preset.
|
|
|
|
Return the original preset object when there are no overrides. When overrides
|
|
are present, return a dataclass copy so the registered preset remains unchanged.
|
|
"""
|
|
overrides = {}
|
|
if config.moe_layer_pattern is not None:
|
|
overrides["moe_layer_pattern"] = config.moe_layer_pattern
|
|
if config.router_pattern is not None:
|
|
overrides["router_pattern"] = config.router_pattern
|
|
if config.expert_pattern is not None:
|
|
overrides["experts_pattern"] = config.expert_pattern
|
|
if config.expert_w1 is not None:
|
|
overrides["expert_w1"] = config.expert_w1
|
|
if config.expert_w2 is not None:
|
|
overrides["expert_w2"] = config.expert_w2
|
|
if config.expert_w3 is not _UNSET:
|
|
overrides["expert_w3"] = config.expert_w3
|
|
if config.num_experts_attr is not None:
|
|
overrides["num_experts_attr"] = config.num_experts_attr
|
|
if config.top_k_attr is not None:
|
|
overrides["top_k_attr"] = config.top_k_attr
|
|
if config.has_shared_experts is not None:
|
|
overrides["has_shared_experts"] = config.has_shared_experts
|
|
if config.shared_experts_pattern is not None:
|
|
overrides["shared_experts_pattern"] = config.shared_experts_pattern
|
|
if config.shared_experts_gate_pattern is not None:
|
|
overrides["shared_experts_gate_pattern"] = config.shared_experts_gate_pattern
|
|
if not overrides:
|
|
return preset
|
|
return replace(preset, **overrides)
|
|
|
|
|
|
def resolve_preset_candidates(
|
|
config: AutoEPConfig,
|
|
model_config,
|
|
) -> list[tuple[str, MoEModelPreset]]:
|
|
"""Resolve ordered preset candidates for AutoEP detection."""
|
|
if config.preset_model is not None:
|
|
preset = apply_config_overrides(config, get_preset(config.preset_model))
|
|
_validate_preset_compatibility(config.preset_model, preset, model_config)
|
|
return [(config.preset_model, preset)]
|
|
|
|
if model_config is not None:
|
|
model_type = getattr(model_config, 'model_type', None)
|
|
if model_type:
|
|
preset_name = preset_name_for_hf_model_type(model_type)
|
|
if preset_name is not None:
|
|
logger.info(f"AutoEP: auto-detected model_type='{model_type}', using preset '{preset_name}'")
|
|
preset = apply_config_overrides(config, get_preset(preset_name))
|
|
_validate_preset_compatibility(preset_name, preset, model_config)
|
|
return [(preset_name, preset)]
|
|
|
|
unsupported_preset = unsupported_preset_for_hf_model_type(model_type)
|
|
if unsupported_preset is not None:
|
|
preset_name, preset = unsupported_preset
|
|
_validate_preset_compatibility(preset_name, preset, model_config)
|
|
|
|
if config.moe_layer_pattern:
|
|
return [("custom", _build_custom_preset(config))]
|
|
|
|
return [(name, apply_config_overrides(config, preset)) for name, preset in PRESET_MODELS.items()]
|
|
|
|
|
|
def _validate_preset_compatibility(
|
|
preset_name: str,
|
|
preset: MoEModelPreset,
|
|
model_config,
|
|
) -> None:
|
|
adapter = get_preset_adapter(preset.preset_adapter)
|
|
adapter.validate_compatibility(preset_name, preset, model_config)
|
|
|
|
|
|
def _build_custom_preset(config: AutoEPConfig) -> MoEModelPreset:
|
|
return MoEModelPreset(
|
|
moe_layer_pattern=config.moe_layer_pattern,
|
|
router_pattern=config.router_pattern or "gate",
|
|
experts_pattern=config.expert_pattern or "experts",
|
|
expert_storage="fused_3d",
|
|
expert_w1=config.expert_w1 or "gate_up_proj",
|
|
expert_w2=config.expert_w2 or "down_proj",
|
|
expert_w3=(None if config.expert_w3 is _UNSET else config.expert_w3),
|
|
num_experts_attr=config.num_experts_attr or "num_local_experts",
|
|
top_k_attr=config.top_k_attr or "num_experts_per_tok",
|
|
score_func=(config.score_func if config.score_func != "auto" else "softmax"),
|
|
score_apply=(config.score_apply if config.score_apply != "auto" else "post"),
|
|
route_norm=(config.route_norm if config.route_norm is not None else True),
|
|
gate_bias=False,
|
|
has_shared_experts=(config.has_shared_experts if config.has_shared_experts is not None else False),
|
|
shared_experts_pattern=config.shared_experts_pattern or "",
|
|
shared_experts_gate_pattern=config.shared_experts_gate_pattern or "",
|
|
)
|