系统 第 4 讲 / 共 17 讲

第 4 讲：Mixture of Experts (MoE)

1. MoE 的动机

Scaling Laws 告诉我们：更多参数 = 更好的性能。 但传统 dense model 的问题是：参数量增加意味着每个 token 的计算量也等比增加。 训练一个 1T 参数的 dense model 需要天文数字的计算资源。

Mixture of Experts (MoE) 提出了一个巧妙的解决方案： 增加参数量，但不等比增加每个 token 的计算量。 关键思想是：对于每个输入 token，只激活模型的一个子集（一部分 expert）， 而不是使用全部参数。

核心权衡

MoE 模型的总参数量可能是 dense model 的 8-16 倍，但每个 token 的计算量（active parameters）只增加少量。例如 Mixtral 8x7B 有约 47B 总参数，但每个 token 只使用约 13B active parameters，计算成本接近一个 13B 的 dense model。

指标Dense 7BMoE 8×7B（如 Mixtral）
总参数量7B~47B
Active parameters per token7B~13B
FLOPs per token~1x~2x
显存需求~1x~6-7x
性能基准接近 dense 30B+

2. Expert 结构

在 MoE Transformer 中，只有 FFN 层被替换为 MoE 层， Attention 层保持不变（所有 token 共享）。每个 MoE 层包含多个 expert， 每个 expert 本质上就是一个独立的 FFN： 

class Expert(nn.Module):
    """Each expert is simply a standard FFN."""
    def __init__(self, d_model: int, d_ff: int):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)

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

这些 expert 的结构完全相同，但参数不共享。 每个 expert 在训练过程中会自然地"专门化"，学习处理不同类型的输入。

3. Router / Gating 机制

MoE 的核心问题是：对于每个输入 token，如何决定将它路由到哪些 expert？ 这由 Router（也叫 Gating Network）决定。

3.1 Top-K Routing

最常用的路由策略是 Top-K Routing： Router 为每个 token 计算一个分数向量（每个 expert 一个分数）， 然后选择分数最高的 K 个 expert。 

class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, num_experts, top_k):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # Router: a simple linear layer
        self.router = nn.Linear(d_model, num_experts, bias=False)

        # Create expert pool
        self.experts = nn.ModuleList([
            Expert(d_model, d_ff) for _ in range(num_experts)
        ])

    def forward(self, x):
        """
        x: (B, T, d_model)
        """
        B, T, d = x.shape

        # Step 1: Compute router logits
        router_logits = self.router(x)  # (B, T, num_experts)

        # Step 2: Top-K selection
        top_k_logits, top_k_indices = router_logits.topk(
            self.top_k, dim=-1
        )  # each (B, T, top_k)

        # Step 3: Softmax over selected experts only
        top_k_weights = F.softmax(top_k_logits, dim=-1)  # (B, T, top_k)

        # Step 4: Compute weighted sum of expert outputs
        output = torch.zeros_like(x)
        for i in range(self.top_k):
            expert_idx = top_k_indices[:, :, i]  # (B, T)
            weight = top_k_weights[:, :, i:i+1]  # (B, T, 1)

            for e in range(self.num_experts):
                mask = (expert_idx == e)
                if mask.any():
                    expert_input = x[mask]
                    expert_output = self.experts[e](expert_input)
                    output[mask] += weight[mask] * expert_output

        return output

Softmax Gating 的直觉

Router 的 softmax 输出可以理解为"置信度"——模型对每个 expert 的信任程度。Top-K 选择确保每个 token 只使用 K 个 expert，保持计算量可控。K 的典型值为 1 或 2。Mixtral 使用 top_k=2，即每个 token 选择 2 个 expert。

3.2 Router 的数学形式

完整的 routing 过程可以用以下公式描述：

g(x) = softmax(TopK(x · Wrouter))
y = Σi∈TopK gi(x) · Experti(x)

其中 TopK 操作将未被选中的 expert 的分数设为 -∞（在 softmax 之前）， 从而只在选中的 expert 之间分配权重。

4. Load Balancing 问题

MoE 训练中最核心的挑战是 Load Balancing（负载均衡）。 如果不加干预，训练过程中很容易出现专家坍缩（expert collapse）： 所有 token 都被路由到少数几个 expert，而其他 expert 完全不被使用。

4.1 为什么会发生专家坍缩？

这是一个正反馈循环：

  1. 某个 expert 偶然在初期表现略好
  2. 更多 token 被路由到这个 expert
  3. 这个 expert 获得更多梯度更新，进一步提升
  4. 其他 expert 因为缺少训练信号而停滞
  5. 最终只有少数 expert 活跃——失去了 MoE 的意义

4.2 Auxiliary Loss（辅助损失）

解决 load balancing 的标准方法是添加一个辅助损失函数， 惩罚不均衡的路由：

Lbalance = α · num_experts · Σi fi · Pi

其中：

def load_balance_loss(router_logits, top_k_indices, num_experts, alpha=0.01):
    """
    router_logits: (B*T, num_experts) - raw router outputs
    top_k_indices: (B*T, top_k) - selected expert indices
    """
    num_tokens = router_logits.shape[0]

    # f_i: fraction of tokens assigned to each expert
    expert_counts = torch.zeros(num_experts, device=router_logits.device)
    for k in range(top_k_indices.shape[1]):
        expert_counts.scatter_add_(
            0, top_k_indices[:, k],
            torch.ones(num_tokens, device=router_logits.device)
        )
    f = expert_counts / (num_tokens * top_k_indices.shape[1])

    # P_i: average router probability for each expert
    router_probs = F.softmax(router_logits, dim=-1)
    P = router_probs.mean(dim=0)

    # Auxiliary loss: minimize when f and P are uniform
    loss = alpha * num_experts * (f * P).sum()

    return loss

直觉理解 Auxiliary Loss

当所有 expert 被均匀使用时（fi = 1/N, Pi = 1/N），辅助损失取最小值。这个损失会"推动"router 更均匀地分配 token，同时不会过强地干预 router 的自然学习。α 的选择很关键——太大会损害模型质量，太小则无法防止坍缩。

5. Token Dropping 策略

在分布式训练中，各 expert 分布在不同的 GPU 上。如果某个 expert 收到过多的 token， 会导致该 GPU 成为瓶颈（其他 GPU 必须等待它完成）。 Token Dropping 是一种强制负载均衡的策略： 

# Token dropping example
capacity = int(num_tokens * top_k / num_experts * capacity_factor)

for expert_idx in range(num_experts):
    assigned_tokens = get_tokens_for_expert(expert_idx)
    if len(assigned_tokens) > capacity:
        # Drop excess tokens (lowest router score first)
        scores = get_router_scores(assigned_tokens, expert_idx)
        keep_indices = scores.topk(capacity).indices
        dropped_indices = set(range(len(assigned_tokens))) - set(keep_indices.tolist())
        # Dropped tokens use residual (skip this expert)

Token Dropping 的问题

Token dropping 在训练和推理之间引入了不一致性：训练时可能丢弃 token，但推理时通常不丢弃。这种 train-test mismatch 可能影响性能。一些新方法（如 Sinkhorn routing）试图在不丢弃 token 的情况下实现负载均衡。

6. 代表模型

6.1 Switch Transformer (Google, 2021)

6.2 GShard (Google, 2020)

6.3 Mixtral 8x7B (Mistral, 2024)

模型总参数Active 参数Expert 数Top-K
Switch-Base3.8B0.2B1281
GShard600B—20482
Mixtral 8x7B47B13B82
DeepSeek-V2236B21B1606

7. MoE 训练的挑战

7.1 训练不稳定性

MoE 模型比同规模的 dense model 更容易出现训练不稳定（loss spikes, divergence）。 主要原因包括：

常用的缓解策略：

7.2 通信开销

在分布式训练中，不同 expert 通常分布在不同 GPU 上（Expert Parallelism）。 每个 MoE 层需要两次 All-to-All 通信：

  1. Dispatch：将 token 发送到各自被分配的 expert 所在的 GPU
  2. Combine：将 expert 输出发回原始 GPU

All-to-All 的通信量与 batch_size × d_model 成正比。在大规模训练中， 通信开销可能占到总时间的 20-40%。优化通信（如通信计算重叠）是 MoE 系统的核心工程挑战。 

# Simplified MoE distributed training flow
# (using Expert Parallelism across 8 GPUs, 8 experts)

# Forward pass:
# 1. Each GPU processes its local tokens through shared layers (Attention)
# 2. Router decides which expert each token goes to
# 3. All-to-All: dispatch tokens to expert GPUs
#    GPU 0 sends tokens for Expert 3 to GPU 3, etc.
# 4. Each GPU processes tokens through its local expert(s)
# 5. All-to-All: combine results back to original GPUs
# 6. Continue with next layer

7.3 显存开销

虽然 MoE 的计算量接近 dense model，但显存需求远大于 dense model。 因为所有 expert 的参数都需要存储在内存中（即使每次只有一部分被使用）。 对于 Mixtral 8x7B，需要存储 47B 参数 + 对应的优化器状态， 这在单个 GPU 上是不可能的——必须使用模型并行策略。

8. 本讲总结