77 lines
2.3 KiB
Python
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(),
|
|
}
|