Mixture of Experts MoE（三）：稀疏大模型架构实战

# MoE 层实现（简化版）
import torch
import torch.nn as nn
import torch.nn.functional as F

class Expert(nn.Module):
    """单个专家——标准 FFN"""

    def __init__(self, dim: int, hidden_dim: int):
        super().__init__()
        self.w1 = nn.Linear(dim, hidden_dim)
        self.w2 = nn.Linear(hidden_dim, dim)
        self.w3 = nn.Linear(dim, hidden_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU 激活
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class MoELayer(nn.Module):
    """Mixture of Experts 层"""

    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        num_experts: int = 8,
        top_k: int = 2,
        aux_loss_weight: float = 0.01,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.aux_loss_weight = aux_loss_weight

        # 门控网络
        self.gate = nn.Linear(dim, num_experts, bias=False)

        # 专家网络
        self.experts = nn.ModuleList([
            Expert(dim, hidden_dim) for _ in range(num_experts)
        ])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (batch, seq_len, dim)
        """
        batch_size, seq_len, dim = x.shape

        # 门控计算
        gate_logits = self.gate(x)  # (batch, seq_len, num_experts)

        # Top-K 路由
        top_k_weights, top_k_indices = torch.topk(
            gate_logits, self.top_k, dim=-1
        )  # (batch, seq_len, top_k)

        # 归一化权重
        top_k_weights = F.softmax(top_k_weights, dim=-1)

        # 计算辅助损失（负载均衡）
        aux_loss = self._compute_aux_loss(gate_logits)

        # 专家计算
        output = torch.zeros_like(x)
        for expert_idx in range(self.num_experts):
            # 找到被路由到当前专家的所有 token
            mask = (top_k_indices == expert_idx)  # (batch, seq_len, top_k)
            token_weights = top_k_weights * mask.float()
            token_weights = token_weights.sum(dim=-1)  # (batch, seq_len)

            if token_weights.sum() > 0:
                # 只计算选中的 token
                selected = x[mask.any(dim=-1)]
                expert_out = self.experts[expert_idx](selected)
                output[mask.any(dim=-1)] += expert_out * token_weights[mask.any(dim=-1)].unsqueeze(-1)

        self.aux_loss = aux_loss
        return output

    def _compute_aux_loss(self, gate_logits: torch.Tensor) -> torch.Tensor:
        """Switch Transformer 负载均衡损失"""
        # 门控概率: P_i = mean(gate_probs_i)
        gate_probs = F.softmax(gate_logits, dim=-1)
        # token 分配比例: f_i = count(tokens assigned to i) / total_tokens
        _, top_k_indices = torch.topk(gate_logits, self.top_k, dim=-1)
        one_hot = F.one_hot(top_k_indices, num_classes=self.num_experts)
        token_counts = one_hot.sum(dim=[0, 1, 2]).float()
        f_i = token_counts / token_counts.sum()
        P_i = gate_probs.mean(dim=[0, 1])
        aux_loss = self.num_experts * (f_i * P_i).sum()
        return aux_loss

# 共享专家 + 路由专家的 MoE 实现
class SharedPlusRoutedMoE(nn.Module):
    """共享专家与路由专家结合的 MoE 层"""

    def __init__(
        self, dim, hidden_dim,
        num_routed_experts=8, top_k=2,
        num_shared_experts=1
    ):
        super().__init__()
        self.top_k = top_k
        self.num_routed_experts = num_routed_experts

        # 路由专家：按 token 动态选择
        self.gate = nn.Linear(dim, num_routed_experts, bias=False)
        self.routed_experts = nn.ModuleList([
            Expert(dim, hidden_dim)
            for _ in range(num_routed_experts)
        ])

        # 共享专家：始终激活，处理通用知识
        self.shared_experts = nn.ModuleList([
            Expert(dim, hidden_dim)
            for _ in range(num_shared_experts)
        ])

    def forward(self, x):
        # 共享专家输出（所有 token 都经过）
        shared_out = sum(e(x) for e in self.shared_experts)

        # 路由专家输出
        gate_logits = self.gate(x)
        top_k_w, top_k_idx = torch.topk(gate_logits, self.top_k, dim=-1)
        top_k_w = F.softmax(top_k_w, dim=-1)

        routed_out = torch.zeros_like(x)
        for i in range(self.num_routed_experts):
            mask = (top_k_idx == i)
            w = (top_k_w * mask.float()).sum(dim=-1)
            if w.sum() > 0:
                selected = x[mask.any(dim=-1)]
                routed_out[mask.any(dim=-1)] += (
                    self.routed_experts[i](selected) *
                    w[mask.any(dim=-1)].unsqueeze(-1)
                )

        return shared_out + routed_out

# 专家利用率监控——训练诊断的关键工具
import torch
import numpy as np

class MoEMonitor:
    """监控 MoE 训练中的专家使用情况"""

    def __init__(self, num_experts, window_size=100):
        self.num_experts = num_experts
        self.window_size = window_size
        self.expert_counts = np.zeros(window_size, dtype=int)
        self.step = 0

    def record(self, expert_assignments):
        """记录当前批次的专家分配情况
        expert_assignments: (batch * seq_len, top_k) 的专家索引
        """
        # 统计每个专家被使用的次数
        counts = np.bincount(
            expert_assignments.flatten(),
            minlength=self.num_experts
        )
        self.expert_counts[self.step % self.window_size] = counts.sum()
        self.step += 1

    def check_collapse(self, threshold=0.05):
        """检查是否有专家崩溃（使用率低于阈值）
        返回: (是否有崩溃, 各专家使用率, 熵)
        """
        recent = self.expert_counts[self.step % self.window_size:]
        if recent.sum() == 0:
            return True, None, 0.0

        rates = recent / recent.sum()
        entropy = -np.sum(rates * np.log(rates + 1e-10))
        max_entropy = np.log(self.num_experts)

        collapsed = np.where(rates < threshold / self.num_experts)[0]
        is_collapsed = len(collapsed) > 0

        if is_collapsed:
            print(f"警告: 专家 {collapsed} 使用率低于阈值!")
            for idx in collapsed:
                print(f"  专家 {idx}: 使用率 {rates[idx]:.4f}")

        print(f"路由熵: {entropy:.2f} / {max_entropy:.2f} "
              f"({entropy/max_entropy:.1%} 理想)")

        return is_collapsed, rates, entropy


# 使用示例
monitor = MoEMonitor(num_experts=8)
# 在每个训练步调用:
# monitor.record(top_k_indices.cpu().numpy())
# monitor.check_collapse()

# MoE 推理引擎核心逻辑（简化）
import torch
from collections import OrderedDict

class MoEInferenceEngine:
    """MoE 推理引擎——支持专家卸载和缓存"""

    def __init__(self, model, gpu_memory_limit_gb=80):
        self.model = model
        self.device = torch.device("cuda")
        self.cpu_device = torch.device("cpu")

        # 专家 LRU 缓存
        self.expert_cache = OrderedDict()
        self.cache_capacity = 4  # 最多缓存几个专家在 GPU 上

        # 预计算每个专家需要的显存
        self.expert_sizes = {}
        for name, module in model.named_modules():
            if "expert" in name.lower() and isinstance(module, torch.nn.Module):
                size = sum(p.numel() * p.element_size()
                          for p in module.parameters())
                self.expert_sizes[name] = size / (1024**3)  # GB

    def load_expert_to_gpu(self, expert_name, expert_module):
        """将专家加载到 GPU（LRU 缓存管理）"""
        if expert_name in self.expert_cache:
            # 已经在缓存中，移到最前面
            self.expert_cache.move_to_end(expert_name)
            return

        # 缓存满了，卸载最久未使用的专家
        while len(self.expert_cache) >= self.cache_capacity:
            oldest = next(iter(self.expert_cache))
            self._offload_expert(oldest)

        # 加载到 GPU
        expert_module.to(self.device)
        self.expert_cache[expert_name] = expert_module

    def _offload_expert(self, expert_name):
        """卸载专家到 CPU"""
        if expert_name in self.expert_cache:
            expert = self.expert_cache.pop(expert_name)
            expert.to(self.cpu_device)
            torch.cuda.empty_cache()

    def generate(self, input_ids, max_new_tokens=100):
        """MoE 模型文本生成"""
        with torch.no_grad():
            for step in range(max_new_tokens):
                outputs = self.model(input_ids)
                next_token = outputs.logits[:, -1, :].argmax(dim=-1)
                input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=-1)

                # 可选：检查生成是否结束
                if next_token.item() == self.model.config.eos_token_id:
                    break

        return input_ids

# MoE 模型选择决策树
def choose_model_architecture(requirements):
    """根据需求选择模型架构"""
    if requirements.get("inference_latency_ms", float("inf")) < 50:
        return {
            "architecture": "小参数稠密模型",
            "reason": "低延迟场景不适合 MoE 的路由开销",
            "suggestion": "使用 7B 以下稠密模型 + 量化",
        }

    if requirements.get("max_gpu_memory_gb", float("inf")) < 20:
        return {
            "architecture": "稠密模型 + INT4 量化",
            "reason": "MoE 的总参数量大，显存需求高",
            "suggestion": "使用 13B 稠密模型 + AWQ 量化到 ~8GB",
        }

    if requirements.get("quality_priority", False):
        return {
            "architecture": "大规模 MoE",
            "reason": "MoE 在同等计算预算下质量最优",
            "suggestion": "Mixtral 8x22B 或 DeepSeek-V3",
        }

    if requirements.get("cost_sensitive", False):
        return {
            "architecture": "中等规模 MoE",
            "reason": "MoE 的推理成本与激活参数相关",
            "suggestion": "Mixtral 8x7B (46B总/13B激活)",
        }

    return {
        "architecture": "稠密模型",
        "reason": "无特殊需求时，稠密模型最简单",
        "suggestion": "Llama 3 8B 或 Qwen 7B",
    }


# MoE 能力 Scaling Law 经验公式
def estimate_moe_capability(
    total_params_b: float,  # 总参数（B）
    active_params_b: float,  # 激活参数（B）
    num_experts: int,
    training_tokens_t: float,  # 训练数据（T）
) -> dict:
    """估算 MoE 模型的能力（经验公式）"""
    # 关键发现：MoE 的能力由总参数和激活参数共同决定
    # 总参数决定知识容量，激活参数决定推理质量

    # 基于 Chinchilla 定律的修正
    optimal_tokens = 20 * active_params_b  # B tokens
    data_efficiency = min(1.0, training_tokens_t / optimal_tokens)

    # 专家多样性收益（额外专家数带来的增益递减）
    expert_diversity = 1.0 + 0.1 * (num_experts ** 0.5)

    # 综合评分（0-100，仅供参考）
    capability_score = (
        20 * (active_params_b ** 0.35) *
        (total_params_b ** 0.1) *
        (data_efficiency ** 0.5) *
        expert_diversity
    )

    return {
        "capability_score": min(capability_score, 100),
        "data_efficiency": data_efficiency,
        "expert_diversity": expert_diversity,
        "estimated_mmlu": min(
            40 + 15 * (active_params_b ** 0.25) *
            (total_params_b ** 0.08),
            90
        ),
        "recommendation": (
            "数据不足" if data_efficiency < 0.5
            else "可以训练" if data_efficiency < 1.0
            else "数据充足"
        ),
    }