Files
2026-07-13 13:18:33 +08:00

77 lines
2.3 KiB
Python

# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
"""DeepSeek-V2 AutoEP preset and parser adapter."""
from __future__ import annotations
from deepspeed.module_inject.auto_ep_presets.base import (
AutoEPConfig,
AutoEPPresetAdapter,
GroupRoutingConfig,
MoEModelPreset,
)
PRESET_NAME = "deepseek_v2"
PRESET = MoEModelPreset(
moe_layer_pattern=r"model\.layers\.\d+\.mlp",
router_pattern="gate",
experts_pattern="experts",
expert_storage="fused_3d",
expert_w1="gate_up_proj",
expert_w2="down_proj",
expert_w3=None,
num_experts_attr="n_routed_experts",
top_k_attr="num_experts_per_tok",
score_func="softmax",
score_apply="post",
route_norm=False,
gate_bias=False,
has_shared_experts=True,
shared_experts_pattern="shared_experts",
autoep_config_defaults={"load_balance_coeff": None},
supports_expert_bias=False,
preset_adapter="deepseek_v2",
hf_model_types=("deepseek_v2", ),
min_transformers_version="5.0.0",
docs_support_notes=("load_balance_coeff / expert-bias auxiliary-loss-free load balancing "
"is not currently supported; non-null values are rejected."),
)
class DeepSeekV2PresetAdapter(AutoEPPresetAdapter):
"""DeepSeek-V2 keeps native top-k normalization and optional group-limited routing."""
def _requires_transformers_version_validation(self) -> bool:
return True
def resolve_route_norm(
self,
config: AutoEPConfig,
preset: MoEModelPreset,
model_config,
) -> bool:
if config.route_norm is not None:
return config.route_norm
return preset.route_norm
def resolve_group_routing(
self,
config: AutoEPConfig,
model_config,
) -> GroupRoutingConfig:
group_routing = super().resolve_group_routing(config, model_config)
if getattr(model_config, 'topk_method', None) != "group_limited_greedy":
return group_routing
return GroupRoutingConfig(
num_expert_groups=group_routing.num_expert_groups or getattr(model_config, 'n_group', None),
num_limited_groups=group_routing.num_limited_groups or getattr(model_config, 'topk_group', None),
group_score_func="max",
)
PRESET_ADAPTERS = {
"deepseek_v2": DeepSeekV2PresetAdapter(),
}