Mixture of Experts（MoE）架构深度解析：从 Switch Transformer 到 Mixtral 的演进与高效扩展策略

import torch
import torch.nn as nn
import torch.nn.functional as F

class GatingNetwork(nn.Module):
    """MoE 基础门控网络"""
    def __init__(self, d_model, num_experts, top_k=2):
        super().__init__()
        self.top_k = top_k
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        logits = self.gate(x)  # (batch_size, seq_len, num_experts)

        # 选择 Top-k 专家
        top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)
        top_k_gates = F.softmax(top_k_logits, dim=-1)

        return top_k_gates, top_k_indices

class MoELayer(nn.Module):
    """基础 Sparse MoE 层"""
    def __init__(self, d_model, d_ff, num_experts, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gate = GatingNetwork(d_model, num_experts, top_k)

        # 每个专家都是独立的 FFN
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.GELU(),
                nn.Linear(d_ff, d_model)
            )
            for _ in range(num_experts)
        ])

    def forward(self, x):
        B, T, D = x.shape
        gates, indices = self.gate(x)  # gates: (B,T,k), indices: (B,T,k)

        # 初始化输出
        output = torch.zeros_like(x)

        # 按专家逐一进行批处理
        for i in range(self.num_experts):
            # 分配给第 i 个专家的 token 掩码
            expert_mask = (indices == i).any(dim=-1)  # (B, T)
            if not expert_mask.any():
                continue

            # 提取对应 token 并执行专家运算
            expert_input = x[expert_mask]
            expert_output = self.experts[i](expert_input)

            # 应用门控权重
            gate_values = gates[indices == i]
            output[expert_mask] += gate_values.unsqueeze(-1) * expert_output

        return output

class SwitchRouter(nn.Module):
    """Switch Transformer 路由器（Top-1）"""
    def __init__(self, d_model, num_experts, capacity_factor=1.25):
        super().__init__()
        self.num_experts = num_experts
        self.capacity_factor = capacity_factor
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):
        B, T, D = x.shape
        # 门控 logits
        logits = self.gate(x)  # (B, T, num_experts)
        probs = F.softmax(logits, dim=-1)

        # 选择 Top-1
        gate_values, expert_indices = probs.max(dim=-1)  # (B, T)

        # 计算专家容量
        capacity = int(self.capacity_factor * T / self.num_experts)

        # 处理超出容量的 token 丢弃
        dispatch_mask = torch.zeros(B, T, self.num_experts, dtype=torch.bool)
        for i in range(self.num_experts):
            expert_mask = (expert_indices == i)
            # 超出容量时丢弃超出部分
            positions = expert_mask.nonzero(as_tuple=True)
            if len(positions[1]) > capacity:
                drop_indices = positions[1][capacity:]
                expert_mask[positions[0][capacity:], drop_indices] = False
            dispatch_mask[:, :, i] = expert_mask

        return gate_values, expert_indices, dispatch_mask

def load_balancing_loss(gates, expert_indices, num_experts, alpha=0.01):
    """Switch Transformer 负载均衡辅助损失"""
    B, T = expert_indices.shape

    # f_i：分配给各专家的 token 比例
    f = torch.zeros(num_experts, device=gates.device)
    for i in range(num_experts):
        f[i] = (expert_indices == i).float().sum() / (B * T)

    # P_i：各专家的平均路由概率
    probs = F.softmax(gates, dim=-1)  # 针对全部专家的概率
    P = probs.mean(dim=[0, 1])  # (num_experts,)

    # 均衡损失
    loss = alpha * num_experts * (f * P).sum()
    return loss

class MixtralMoELayer(nn.Module):
    """Mixtral 8x7B 风格 MoE 层"""
    def __init__(self, d_model=4096, d_ff=14336, num_experts=8, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # 路由器：线性门控
        self.gate = nn.Linear(d_model, num_experts, bias=False)

        # 8 个独立的 SwiGLU FFN 专家
        self.experts = nn.ModuleList([
            SwiGLUExpert(d_model, d_ff) for _ in range(num_experts)
        ])

    def forward(self, x):
        B, T, D = x.shape
        x_flat = x.view(-1, D)  # (B*T, D)

        # 计算路由概率
        logits = self.gate(x_flat)  # (B*T, 8)
        weights, indices = torch.topk(logits, self.top_k, dim=-1)
        weights = F.softmax(weights, dim=-1)  # 在 Top-k 内做归一化

        # 对专家输出做加权求和
        output = torch.zeros_like(x_flat)
        for i in range(self.top_k):
            expert_idx = indices[:, i]  # (B*T,)
            gate_weight = weights[:, i]  # (B*T,)

            for j in range(self.num_experts):
                mask = (expert_idx == j)
                if mask.any():
                    expert_out = self.experts[j](x_flat[mask])
                    output[mask] += gate_weight[mask].unsqueeze(-1) * expert_out

        return output.view(B, T, D)


class SwiGLUExpert(nn.Module):
    """Mixtral 使用的基于 SwiGLU 的 FFN 专家"""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.w3 = nn.Linear(d_model, d_ff, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

class DeepSeekMoELayer(nn.Module):
    """DeepSeek-MoE 风格层：细粒度专家 + 共享专家"""
    def __init__(
        self,
        d_model=2048,
        d_ff=10944,
        num_routed_experts=64,
        num_shared_experts=2,
        top_k=6,
    ):
        super().__init__()
        self.top_k = top_k

        # 细粒度路由专家（较小的 FFN）
        expert_d_ff = d_ff // 4  # 相比原来缩小到 1/4
        self.routed_experts = nn.ModuleList([
            SwiGLUExpert(d_model, expert_d_ff)
            for _ in range(num_routed_experts)
        ])

        # 共享专家（始终激活）
        self.shared_experts = nn.ModuleList([
            SwiGLUExpert(d_model, d_ff)
            for _ in range(num_shared_experts)
        ])

        # 路由器
        self.gate = nn.Linear(d_model, num_routed_experts, bias=False)

    def forward(self, x):
        B, T, D = x.shape
        x_flat = x.view(-1, D)

        # 共享专家输出（始终激活）
        shared_out = sum(expert(x_flat) for expert in self.shared_experts)

        # 路由专家输出
        logits = self.gate(x_flat)
        weights, indices = torch.topk(logits, self.top_k, dim=-1)
        weights = F.softmax(weights, dim=-1)

        routed_out = torch.zeros_like(x_flat)
        for i in range(self.top_k):
            expert_idx = indices[:, i]
            gate_weight = weights[:, i]
            for j in range(len(self.routed_experts)):
                mask = (expert_idx == j)
                if mask.any():
                    out = self.routed_experts[j](x_flat[mask])
                    routed_out[mask] += gate_weight[mask].unsqueeze(-1) * out

        return (shared_out + routed_out).view(B, T, D)

class ExpertChoiceRouter(nn.Module):
    """Expert Choice 路由：由专家选择 token"""
    def __init__(self, d_model, num_experts, capacity_factor=1.0):
        super().__init__()
        self.num_experts = num_experts
        self.capacity_factor = capacity_factor
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):
        B, T, D = x.shape
        x_flat = x.view(-1, D)  # (N, D) where N = B*T
        N = x_flat.shape[0]

        logits = self.gate(x_flat)  # (N, num_experts)
        scores = F.softmax(logits, dim=0)  # 在 token 维度上做 softmax

        # 每个专家要处理的 token 数
        k = int(self.capacity_factor * N / self.num_experts)

        # 各专家选择 Top-k 个 token
        top_k_scores, top_k_indices = torch.topk(
            scores.t(), k, dim=-1
        )  # (num_experts, k)

        return top_k_scores, top_k_indices

# DeepSpeed MoE 配置示例
deepspeed_config = {
    "train_batch_size": 256,
    "fp16": {"enabled": True},
    "zero_optimization": {"stage": 2},
    "moe": {
        "enabled": True,
        "ep_size": 8,            # Expert Parallelism：将专家分散到 8 个 GPU 上
        "num_experts": 64,
        "top_k": 2,
        "capacity_factor": 1.25,
        "min_capacity": 4,
        "use_residual": True,    # 残差 MoE
        "moe_param_group": True  # 专家参数单独分组
    }
}