mirror of https://github.com/razor-ai/soup.git
127 lines
3.8 KiB
Python
127 lines
3.8 KiB
Python
"""MoE (Mixture of Experts) model detection and LoRA target module helpers.
|
|
|
|
Supports: Qwen3 MoE (30B-A3B), Mixtral (8x7B, 8x22B), DeepSeek V3, DBRX,
|
|
OLMoE, JetMoE, and other models using MoE / sparse expert architectures.
|
|
"""
|
|
|
|
from typing import Optional
|
|
|
|
# Known MoE architecture config keys that indicate expert layers
|
|
MOE_CONFIG_KEYS = (
|
|
"num_experts",
|
|
"num_local_experts",
|
|
"num_experts_per_tok",
|
|
"num_experts_per_token",
|
|
"n_routed_experts",
|
|
"moe_num_experts",
|
|
)
|
|
|
|
# Common expert FFN module name patterns across MoE architectures
|
|
MOE_EXPERT_PATTERNS = [
|
|
"experts", # Mixtral, Qwen3, DeepSeek
|
|
"gate_proj", # Expert gate projection
|
|
"up_proj", # Expert up projection
|
|
"down_proj", # Expert down projection
|
|
"w1", # DeepSeek V3 expert naming
|
|
"w2",
|
|
"w3",
|
|
]
|
|
|
|
# Standard attention + MLP target modules (non-expert layers)
|
|
STANDARD_TARGET_MODULES = [
|
|
"q_proj",
|
|
"k_proj",
|
|
"v_proj",
|
|
"o_proj",
|
|
]
|
|
|
|
|
|
def detect_moe_model(model) -> bool:
|
|
"""Detect whether a model uses Mixture of Experts architecture.
|
|
|
|
Checks model config for known MoE indicators.
|
|
"""
|
|
if not hasattr(model, "config"):
|
|
return False
|
|
|
|
config = model.config
|
|
|
|
# Check for known MoE config keys
|
|
for key in MOE_CONFIG_KEYS:
|
|
value = getattr(config, key, None)
|
|
if value is not None and isinstance(value, (int, float)) and value > 1:
|
|
return True
|
|
|
|
# Check model_type for known MoE architectures
|
|
model_type = getattr(config, "model_type", "")
|
|
moe_types = {"mixtral", "qwen3_moe", "qwen2_moe", "dbrx", "deepseek_v2",
|
|
"deepseek_v3", "olmoe", "jetmoe", "arctic", "grok"}
|
|
if model_type.lower() in moe_types:
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def get_moe_target_modules(model) -> Optional[list[str]]:
|
|
"""Get LoRA target modules for MoE models (ScatterMoE LoRA).
|
|
|
|
Returns a list of module name patterns that includes both attention
|
|
layers and expert FFN layers for comprehensive LoRA coverage.
|
|
Returns None if the model is not an MoE model.
|
|
"""
|
|
if not detect_moe_model(model):
|
|
return None
|
|
|
|
# Scan model for expert layer names
|
|
expert_modules = set()
|
|
for name, _module in model.named_modules():
|
|
name_lower = name.lower()
|
|
# Look for expert-specific patterns
|
|
if "expert" in name_lower or "moe" in name_lower:
|
|
parts = name.split(".")
|
|
for part in parts:
|
|
if part in ("gate_proj", "up_proj", "down_proj", "w1", "w2", "w3"):
|
|
expert_modules.add(part)
|
|
|
|
# Combine attention targets with discovered expert targets
|
|
targets = list(STANDARD_TARGET_MODULES)
|
|
if expert_modules:
|
|
targets.extend(sorted(expert_modules))
|
|
else:
|
|
# Fallback: add common expert FFN patterns
|
|
targets.extend(["gate_proj", "up_proj", "down_proj"])
|
|
|
|
return targets
|
|
|
|
|
|
def get_moe_info(model) -> dict:
|
|
"""Extract MoE architecture details from a model config.
|
|
|
|
Returns a dict with num_experts, num_active_experts, and model_type.
|
|
Returns empty dict if not an MoE model.
|
|
"""
|
|
if not hasattr(model, "config"):
|
|
return {}
|
|
|
|
config = model.config
|
|
info = {}
|
|
|
|
# Number of total experts
|
|
for key in ("num_local_experts", "num_experts", "n_routed_experts", "moe_num_experts"):
|
|
value = getattr(config, key, None)
|
|
if value is not None:
|
|
info["num_experts"] = value
|
|
break
|
|
|
|
# Number of active experts per token
|
|
for key in ("num_experts_per_tok", "num_experts_per_token", "num_selected_experts"):
|
|
value = getattr(config, key, None)
|
|
if value is not None:
|
|
info["num_active_experts"] = value
|
|
break
|
|
|
|
if info:
|
|
info["model_type"] = getattr(config, "model_type", "unknown")
|
|
|
|
return info
|