deepspeedai--deepspeed/deepspeed/module_inject/auto_ep_config.py

# Copyright (c) DeepSpeed Team.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team
"""AutoEP configuration: config parsing, model presets, and validation."""

from __future__ import annotations

from deepspeed.module_inject.auto_ep_presets.base import (
    _UNSET,
    _raise_unsupported_load_balance_coeff,
    AutoEPConfig,
    MoELayerSpec,
    MoEModelPreset,
)
from deepspeed.module_inject.auto_ep_presets.registry import (
    PRESET_MODELS,
    available_preset_names,
    resolve_autoep_config_defaults,
)
from deepspeed.module_inject.auto_ep_folding import build_folding_spec, validate_folding_global
from deepspeed.utils import logger

__all__ = [
    "_UNSET",
    "AutoEPConfig",
    "MoELayerSpec",
    "MoEModelPreset",
    "PRESET_MODELS",
    "parse_autoep_config",
    "resolve_autoep_config_defaults",
    "validate_autoep_config",
    "validate_autoep_post_detection",
]

# ---------------------------------------------------------------------------
# Config parsing
# ---------------------------------------------------------------------------


def parse_autoep_config(param_dict: dict) -> AutoEPConfig:
    """Parse the 'expert_parallel' section from DS config JSON."""
    if not param_dict:
        return AutoEPConfig()

    config = AutoEPConfig()
    config.enabled = param_dict.get("enabled", False)
    config.autoep_size = param_dict.get("autoep_size", 1)
    config.expert_tensor_parallel_size = param_dict.get("expert_tensor_parallel_size", 1)
    config.validate_folding_routing = param_dict.get("validate_folding_routing", False)
    config.preset_model = param_dict.get("preset_model", None)
    config.moe_layer_pattern = param_dict.get("moe_layer_pattern", None)
    config.expert_pattern = param_dict.get("expert_pattern", None)
    config.router_pattern = param_dict.get("router_pattern", None)
    config.use_grouped_mm = param_dict.get("use_grouped_mm", True)
    config.route_norm = param_dict.get("route_norm", None)
    config.route_scale = param_dict.get("route_scale", 1.0)
    config.score_apply = param_dict.get("score_apply", "auto")
    config.combine_impl = param_dict.get("combine_impl", "auto")
    config.num_expert_groups = param_dict.get("num_expert_groups", None)
    config.num_limited_groups = param_dict.get("num_limited_groups", None)
    config.score_func = param_dict.get("score_func", "auto")
    config.top_k = param_dict.get("top_k", "auto")
    if "load_balance_coeff" in param_dict:
        value = param_dict["load_balance_coeff"]
        if value is not None:
            _raise_unsupported_load_balance_coeff(value)
        config.load_balance_coeff = None
        config._load_balance_coeff_explicit = True
    else:
        config.load_balance_coeff = None
        config._load_balance_coeff_explicit = False
    config.routed_scaling_factor = param_dict.get("routed_scaling_factor", "auto")
    config.expert_w1 = param_dict.get("expert_w1", None)
    config.expert_w2 = param_dict.get("expert_w2", None)
    # expert_w3: key absent → _UNSET (preset default); key present with null → None (fused); key present with string → custom name
    if "expert_w3" in param_dict:
        config.expert_w3 = param_dict["expert_w3"]  # None or string
    else:
        config.expert_w3 = _UNSET
    config.num_experts_attr = param_dict.get("num_experts_attr", None)
    config.top_k_attr = param_dict.get("top_k_attr", None)
    config.has_shared_experts = param_dict.get("has_shared_experts", None)
    config.shared_experts_pattern = param_dict.get("shared_experts_pattern", None)
    config.shared_experts_gate_pattern = param_dict.get("shared_experts_gate_pattern", None)

    return config


# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------


def validate_autoep_config(
    config: AutoEPConfig,
    world_size: int,
    pp_size: int,
    tp_size: int,
    sp_size: int,
    *,
    zero_stage: int = 0,
    deepcompile_enabled: bool = False,
    tp_preset_model: str | None = None,
    use_data_before_expert_parallel: bool = False,
    mpu=None,
    zero_offload_optimizer: bool = False,
    zero_offload_param: bool = False,
) -> None:
    """Validate config constraints. Raises ValueError on invalid config."""
    if config.load_balance_coeff is not None:
        _raise_unsupported_load_balance_coeff(config.load_balance_coeff)

    if not isinstance(config.validate_folding_routing, bool):
        raise ValueError("expert_parallel.validate_folding_routing must be a boolean")

    if not config.enabled:
        return

    folding_spec = build_folding_spec(
        world_size=world_size,
        pp_size=pp_size,
        tp_size=max(tp_size, 1),
        ep_size=config.autoep_size,
        etp_size=config.expert_tensor_parallel_size,
        mp_mode="tp" if tp_size > 1 else "sp",
    )
    validate_folding_global(
        folding_spec,
        zero_stage=zero_stage,
        sp_size=sp_size,
        deepcompile_enabled=deepcompile_enabled,
        use_data_before_expert_parallel=use_data_before_expert_parallel,
        mpu=mpu,
        autoep_enabled=config.enabled,
        tp_preset=tp_preset_model,
        ep_preset=config.preset_model,
        zero_offload_optimizer=zero_offload_optimizer,
        zero_offload_param=zero_offload_param,
    )

    # Validate preset_model if specified
    if config.preset_model is not None and config.preset_model not in PRESET_MODELS:
        raise ValueError(f"Unknown preset_model '{config.preset_model}'. "
                         f"Available presets: {list(available_preset_names())}")

    # Validate score_apply
    valid_score_apply = ("auto", "pre", "post")
    if config.score_apply not in valid_score_apply:
        raise ValueError(f"score_apply must be one of {valid_score_apply}, "
                         f"got '{config.score_apply}'")

    # Validate combine_impl
    valid_combine_impl = ("auto", "weighted_sum", "legacy_bmm")
    if config.combine_impl not in valid_combine_impl:
        raise ValueError(f"combine_impl must be one of {valid_combine_impl}, "
                         f"got '{config.combine_impl}'")

    # Validate score_func
    valid_score_func = ("auto", "softmax", "sigmoid")
    if config.score_func not in valid_score_func:
        raise ValueError(f"score_func must be one of {valid_score_func}, "
                         f"got '{config.score_func}'")

    # Validate group-limited routing constraints
    if config.num_limited_groups is not None:
        if config.num_limited_groups < 1:
            raise ValueError(f"num_limited_groups must be >= 1, got {config.num_limited_groups}")

    if config.num_expert_groups is not None:
        if config.num_expert_groups < 1:
            raise ValueError(f"num_expert_groups must be >= 1, got {config.num_expert_groups}")
        if config.num_limited_groups is not None and config.num_limited_groups > config.num_expert_groups:
            raise ValueError(f"num_limited_groups ({config.num_limited_groups}) must be <= "
                             f"num_expert_groups ({config.num_expert_groups})")
        logger.warning("num_expert_groups is set; interaction with EP topology "
                       "is not yet optimized.")

    # Warn if autoep_size == 1 (no EP needed)
    if config.autoep_size == 1:
        logger.warning("autoep_size=1 means every rank owns all experts with no AllToAll. "
                       "AutoEP replacement remains enabled, but expert-parallel communication "
                       "is bypassed because every rank owns all experts.")

    # Helper validators (local to validate_autoep_config)
    def _validate_attr_name(field_name: str, value, *, allow_dot: bool = False) -> None:
        if value is None:
            return
        if not isinstance(value, str) or value == "":
            raise ValueError(f"{field_name} must be a non-empty string")
        if not allow_dot and "." in value:
            raise ValueError(f"{field_name} must be a direct attribute name (no dots)")

    # Validate expert weight names
    _validate_attr_name("expert_w1", config.expert_w1)
    _validate_attr_name("expert_w2", config.expert_w2)
    if config.expert_w3 is not _UNSET and config.expert_w3 is not None:
        _validate_attr_name("expert_w3", config.expert_w3)

    # Validate model.config attribute names
    _validate_attr_name("num_experts_attr", config.num_experts_attr)
    _validate_attr_name("top_k_attr", config.top_k_attr)

    # Validate child-name fields (direct attribute names, not regex/path)
    _validate_attr_name("router_pattern", config.router_pattern)
    _validate_attr_name("expert_pattern", config.expert_pattern)
    _validate_attr_name("shared_experts_pattern", config.shared_experts_pattern)
    _validate_attr_name("shared_experts_gate_pattern", config.shared_experts_gate_pattern)

    # Validate has_shared_experts type
    if config.has_shared_experts is not None and not isinstance(config.has_shared_experts, bool):
        raise ValueError("has_shared_experts must be a boolean when set")

    # Warn if explicit top_k overrides top_k_attr
    if isinstance(config.top_k, int) and config.top_k_attr is not None:
        logger.warning("top_k is explicitly set; top_k_attr will be ignored.")

    if config.routed_scaling_factor != "auto" and not isinstance(config.routed_scaling_factor, (int, float)):
        raise ValueError("routed_scaling_factor must be a number or 'auto'")

    # Validate shared expert field pairing
    if config.has_shared_experts is True and not config.shared_experts_pattern:
        logger.warning("has_shared_experts=True but shared_experts_pattern is not set. "
                       "Shared expert detection requires both fields.")
    if config.shared_experts_pattern and config.has_shared_experts is not True:
        logger.warning(f"shared_experts_pattern='{config.shared_experts_pattern}' is set "
                       f"but has_shared_experts is not True. Pattern will be ignored.")
    if config.shared_experts_gate_pattern and config.has_shared_experts is not True:
        logger.warning(f"shared_experts_gate_pattern='{config.shared_experts_gate_pattern}' is set "
                       f"but has_shared_experts is not True. Pattern will be ignored.")

    # Warn if custom override fields are set alongside preset_model or auto-detect
    custom_fields_set = []
    if config.moe_layer_pattern is not None:
        custom_fields_set.append("moe_layer_pattern")
    if config.router_pattern is not None:
        custom_fields_set.append("router_pattern")
    if config.expert_pattern is not None:
        custom_fields_set.append("expert_pattern")
    if config.expert_w1 is not None:
        custom_fields_set.append("expert_w1")
    if config.expert_w2 is not None:
        custom_fields_set.append("expert_w2")
    if config.expert_w3 is not _UNSET:
        custom_fields_set.append("expert_w3")
    if config.num_experts_attr is not None:
        custom_fields_set.append("num_experts_attr")
    if config.top_k_attr is not None:
        custom_fields_set.append("top_k_attr")
    if config.has_shared_experts is not None:
        custom_fields_set.append("has_shared_experts")
    if config.shared_experts_pattern is not None:
        custom_fields_set.append("shared_experts_pattern")
    if config.shared_experts_gate_pattern is not None:
        custom_fields_set.append("shared_experts_gate_pattern")
    if custom_fields_set and config.preset_model is not None:
        logger.warning(f"Custom preset fields {custom_fields_set} are set alongside "
                       f"preset_model='{config.preset_model}'. Custom fields will override "
                       f"preset defaults during detection.")
    if custom_fields_set and config.preset_model is None and config.moe_layer_pattern is None:
        logger.warning(f"Custom preset fields {custom_fields_set} are set without preset_model or "
                       f"moe_layer_pattern. Overrides will apply to auto-detected presets or try-all.")


def validate_autoep_post_detection(
    config: AutoEPConfig,
    specs: list[MoELayerSpec],
) -> None:
    """Post-detection validation: ep_size vs num_experts constraints."""
    if not config.enabled or not specs:
        return

    for spec in specs:
        # ep_size must not exceed num_experts
        if config.autoep_size > spec.num_experts:
            valid_divisors = _divisors(spec.num_experts)
            raise ValueError(f"autoep_size={config.autoep_size} exceeds num_experts="
                             f"{spec.num_experts} in layer '{spec.moe_module_name}'. "
                             f"Each rank must own at least one expert. "
                             f"Valid autoep_size values (divisors of {spec.num_experts}): "
                             f"{valid_divisors}")

        # num_experts must be divisible by ep_size
        if spec.num_experts % config.autoep_size != 0:
            valid_sizes = [d for d in _divisors(spec.num_experts) if d <= spec.num_experts]
            raise ValueError(f"num_experts={spec.num_experts} in layer "
                             f"'{spec.moe_module_name}' is not divisible by "
                             f"autoep_size={config.autoep_size}. "
                             f"Suggested autoep_size values: {valid_sizes}")

        num_expert_groups = spec.num_expert_groups if spec.num_expert_groups is not None else config.num_expert_groups
        num_limited_groups = spec.num_limited_groups if spec.num_limited_groups is not None else config.num_limited_groups

        # Validate group-limited routing constraints after layer-specific defaults.
        if num_limited_groups is not None and num_expert_groups is None:
            raise ValueError(f"num_limited_groups requires num_expert_groups to be set "
                             f"in layer '{spec.moe_module_name}'")

        if num_expert_groups is not None:
            if num_expert_groups < 1:
                raise ValueError(f"num_expert_groups must be >= 1 in layer '{spec.moe_module_name}', "
                                 f"got {num_expert_groups}")
            if spec.num_experts % num_expert_groups != 0:
                raise ValueError(f"num_expert_groups ({num_expert_groups}) must divide "
                                 f"num_experts ({spec.num_experts}) in layer "
                                 f"'{spec.moe_module_name}'")
            if num_limited_groups is None:
                raise ValueError(f"num_limited_groups must be set when num_expert_groups is set "
                                 f"in layer '{spec.moe_module_name}'")
            if num_limited_groups < 1:
                raise ValueError(f"num_limited_groups must be >= 1 in layer '{spec.moe_module_name}', "
                                 f"got {num_limited_groups}")
            if num_limited_groups > num_expert_groups:
                raise ValueError(f"num_limited_groups ({num_limited_groups}) must be <= "
                                 f"num_expert_groups ({num_expert_groups}) in layer "
                                 f"'{spec.moe_module_name}'")


def _divisors(n: int) -> list[int]:
    """Return sorted list of positive divisors of n."""
    divs = []
    for i in range(1, int(n**0.5) + 1):
        if n % i == 0:
            divs.append(i)
            if i != n // i:
                divs.append(n // i)
    return sorted(divs)


def fill_autoep_config_from_hf(config: AutoEPConfig, model_config) -> None:
    """Back-fill AutoEPConfig fields from HF model config when user hasn't set them.

    HF field names (e.g. n_group, topk_group, routed_scaling_factor) differ from
    AutoEP's internal names, so we map them explicitly rather than relying on the
    user to duplicate these values in the DS config JSON.
    """
    if model_config is None:
        return
    # n_group / topk_group: DeepSeek-style node-limited routing groups
    if config.num_expert_groups is None:
        config.num_expert_groups = getattr(model_config, 'n_group', None)
    if config.num_limited_groups is None:
        config.num_limited_groups = getattr(model_config, 'topk_group', None)
    # routed_scaling_factor: sigmoid score scaling (DeepSeek-V3 / Moonlight)
    if config.routed_scaling_factor == "auto":
        hf_scale = getattr(model_config, 'routed_scaling_factor', None)
        if hf_scale is not None:
            config.route_scale = float(hf_scale)