Mixture of Experts (MoE) 架构论文深度解析：从 GShard 到 DeepSeek-MoE

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

class BasicMoEGating(nn.Module):
    """Sparsely-Gated MoE 的基础门控网络"""
    def __init__(self, d_model: int, num_experts: int, top_k: int = 2,
                 noise_std: float = 1.0):
        super().__init__()
        self.top_k = top_k
        self.noise_std = noise_std
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        self.noise_gate = nn.Linear(d_model, num_experts, bias=False)

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

        # 训练时加入噪声（促进探索）
        if self.training:
            noise = F.softplus(self.noise_gate(x))
            logits = logits + torch.randn_like(logits) * noise * self.noise_std

        # 选择 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 层（含 Capacity Factor）"""
    def __init__(self, d_model: int, d_ff: int, num_experts: int,
                 top_k: int = 2, capacity_factor: float = 1.25):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.capacity_factor = capacity_factor
        self.gate = BasicMoEGating(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: torch.Tensor):
        B, T, D = x.shape
        gates, indices = self.gate(x)
        expert_capacity = int(self.capacity_factor * T / self.num_experts)

        output = torch.zeros_like(x)
        tokens_dropped = 0

        for i in range(self.num_experts):
            # 分配给第 i 个专家的 token 掩码
            expert_mask = (indices == i).any(dim=-1)  # (B, T)
            # 丢弃超出容量的 token
            token_count = expert_mask.sum(dim=-1)
            for b in range(B):
                if token_count[b] > expert_capacity:
                    overflow = token_count[b] - expert_capacity
                    tokens_dropped += overflow.item()

            if not expert_mask.any():
                continue

            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, tokens_dropped

class GShardRouter(nn.Module):
    """GShard Top-2 路由器（含随机路由）"""
    def __init__(self, d_model: int, num_experts: int,
                 capacity_factor: float = 2.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: torch.Tensor):
        # x: (batch_size, seq_len, d_model)
        B, T, D = x.shape
        logits = self.gate(x)
        probs = F.softmax(logits, dim=-1)

        # Top-1：确定性选择
        top1_probs, top1_indices = probs.max(dim=-1)

        # Top-2：概率性选择（从其余专家中采样）
        mask = torch.zeros_like(probs)
        mask.scatter_(-1, top1_indices.unsqueeze(-1), 1.0)
        remaining_probs = probs * (1 - mask)
        remaining_probs = remaining_probs / (remaining_probs.sum(dim=-1, keepdim=True) + 1e-8)

        # 按概率采样第二个专家
        top2_indices = torch.multinomial(
            remaining_probs.view(-1, self.num_experts), 1
        ).view(B, T)
        top2_probs = probs.gather(-1, top2_indices.unsqueeze(-1)).squeeze(-1)

        # 重新归一化门控值
        total = top1_probs + top2_probs
        top1_gates = top1_probs / total
        top2_gates = top2_probs / total

        return (top1_gates, top1_indices, top2_gates, top2_indices)

class SwitchRouter(nn.Module):
    """Switch Transformer 路由器（Top-1 单一专家）"""
    def __init__(self, d_model: int, num_experts: int,
                 capacity_factor: float = 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: torch.Tensor):
        B, T, D = x.shape
        logits = self.gate(x)
        probs = F.softmax(logits, dim=-1)

        # Top-1 选择（单一专家）
        gate_values, expert_indices = probs.max(dim=-1)  # (B, T)

        # 基于 Capacity 的 token 丢弃
        capacity = int(self.capacity_factor * T / self.num_experts)

        # 追踪各专家分配到的 token 数
        expert_counts = torch.zeros(B, self.num_experts, device=x.device)
        drop_mask = torch.zeros(B, T, dtype=torch.bool, device=x.device)

        for t in range(T):
            for b in range(B):
                eid = expert_indices[b, t].item()
                if expert_counts[b, eid] >= capacity:
                    drop_mask[b, t] = True  # 丢弃超出容量的 token
                else:
                    expert_counts[b, eid] += 1

        # 把被丢弃 token 的门控值设为 0
        gate_values = gate_values * (~drop_mask).float()

        return gate_values, expert_indices, drop_mask

def switch_load_balancing_loss(router_probs: torch.Tensor,
                                expert_indices: torch.Tensor,
                                num_experts: int,
                                alpha: float = 0.01) -> torch.Tensor:
    """Switch Transformer 的 Simplified Load Balancing Loss

    Args:
        router_probs: 路由器概率 (B, T, num_experts)
        expert_indices: 所选专家索引 (B, T)
        num_experts: 专家数
        alpha: 均衡系数
    """
    # f_i：分配给各专家的 token 比例
    one_hot = F.one_hot(expert_indices, num_experts).float()
    f = one_hot.mean(dim=[0, 1])  # (num_experts,)

    # P_i：各专家的平均路由器概率
    P = router_probs.mean(dim=[0, 1])  # (num_experts,)

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

class DeepSeekMoELayer(nn.Module):
    """DeepSeek-MoE 层（Fine-grained Segmentation + Shared Expert）"""
    def __init__(self, d_model: int, d_ff: int, num_routed_experts: int,
                 num_shared_experts: int, top_k: int, segmentation_factor: int = 4):
        super().__init__()
        self.num_routed = num_routed_experts
        self.num_shared = num_shared_experts
        self.top_k = top_k

        # Fine-grained：把 FFN 中间维度除以 segmentation_factor
        fine_d_ff = d_ff // segmentation_factor

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

        # 路由专家（细分后的尺寸）
        self.routed_experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, fine_d_ff),
                nn.SiLU(),
                nn.Linear(fine_d_ff, d_model)
            ) for _ in range(num_routed_experts)
        ])

        # 门控网络（用于路由专家）
        self.gate = nn.Linear(d_model, num_routed_experts, bias=False)

    def forward(self, x: torch.Tensor):
        B, T, D = x.shape

        # 1. 共享专家输出（应用于所有 token）
        shared_output = sum(expert(x) for expert in self.shared_experts)

        # 2. 选择路由专家并计算输出
        logits = self.gate(x)
        probs = F.softmax(logits, dim=-1)
        top_k_probs, top_k_indices = torch.topk(probs, self.top_k, dim=-1)
        top_k_gates = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        routed_output = torch.zeros_like(x)
        for k in range(self.top_k):
            expert_idx = top_k_indices[:, :, k]  # (B, T)
            gate_val = top_k_gates[:, :, k]       # (B, T)
            for i in range(self.num_routed):
                mask = (expert_idx == i)
                if mask.any():
                    inp = x[mask]
                    out = self.routed_experts[i](inp)
                    routed_output[mask] += gate_val[mask].unsqueeze(-1) * out

        return shared_output + routed_output

def compute_expert_utilization(expert_indices: torch.Tensor,
                                num_experts: int) -> dict:
    """专家利用率监控工具

    Args:
        expert_indices: 所选专家索引 (B, T) 或 (B, T, K)
        num_experts: 专家总数

    Returns:
        利用率指标字典
    """
    flat = expert_indices.flatten()
    counts = torch.bincount(flat, minlength=num_experts).float()
    total = flat.numel()

    # 各专家的 token 比例
    fractions = counts / total

    # 利用率指标
    max_fraction = fractions.max().item()
    min_fraction = fractions.min().item()
    ideal_fraction = 1.0 / num_experts

    # 不均衡比例（越接近 1 越均匀）
    balance_ratio = min_fraction / (max_fraction + 1e-8)

    # 非激活专家数
    dead_experts = (counts == 0).sum().item()

    return {
        "balance_ratio": balance_ratio,
        "max_load": max_fraction,
        "min_load": min_fraction,
        "ideal_load": ideal_fraction,
        "dead_experts": dead_experts,
        "load_std": fractions.std().item(),
    }

def router_z_loss(router_logits: torch.Tensor) -> torch.Tensor:
    """Router z-loss：通过抑制路由器 logit 的规模来提升稳定性

    Args:
        router_logits: 路由器的原始 logit (B, T, num_experts)
    """
    # log-sum-exp 的平方
    log_z = torch.logsumexp(router_logits, dim=-1)  # (B, T)
    z_loss = (log_z ** 2).mean()
    return z_loss


class StableMoETrainer:
    """用于稳定 MoE 训练的 Trainer 示例"""
    def __init__(self, model, optimizer, alpha_balance=0.01,
                 alpha_z=0.001):
        self.model = model
        self.optimizer = optimizer
        self.alpha_balance = alpha_balance
        self.alpha_z = alpha_z

    def train_step(self, input_ids, labels):
        self.optimizer.zero_grad()

        # Forward pass
        logits, router_logits, expert_indices = self.model(input_ids)

        # Task loss
        task_loss = F.cross_entropy(
            logits.view(-1, logits.size(-1)), labels.view(-1)
        )

        # Load balancing loss
        router_probs = F.softmax(router_logits, dim=-1)
        balance_loss = switch_load_balancing_loss(
            router_probs, expert_indices, self.model.num_experts
        )

        # Router z-loss
        z_loss = router_z_loss(router_logits)

        # Total loss
        total_loss = task_loss + self.alpha_balance * balance_loss + self.alpha_z * z_loss

        total_loss.backward()

        # 梯度裁剪（在 MoE 训练中尤为重要）
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)

        self.optimizer.step()

        return {
            "task_loss": task_loss.item(),
            "balance_loss": balance_loss.item(),
            "z_loss": z_loss.item(),
            "total_loss": total_loss.item(),
        }