Mixture of Experts 架构：稀疏激活如何重塑大语言模型的效率与规模

┌──────────────────────────────────────────────────────────────┐
│                    Dense Model vs MoE Model                   │
│                                                               │
│  Dense:                                                       │
│  Token ──→ [FFN: 所有参数激活] ──→ Output                     │
│                                                               │
│  MoE:                                                         │
│  Token ──→ [Router] ──→ 选中 Expert_2, Expert_5              │
│           ├──→ Expert_1 (静默)                                │
│           ├──→ Expert_2 (激活) ──→ Weighted Sum ──→ Output    │
│           ├──→ Expert_3 (静默)                                │
│           ├──→ Expert_4 (静默)                                │
│           ├──→ Expert_5 (激活) ──↗                            │
│           ├──→ Expert_6 (静默)                                │
│           ├──→ Expert_7 (静默)                                │
│           └──→ Expert_8 (静默)                                │
└──────────────────────────────────────────────────────────────┘

标准 Transformer Block:                MoE Transformer Block:

┌────────────────────┐               ┌────────────────────┐
│   Layer Normalization│               │   Layer Normalization│
└────────┬───────────┘               └────────┬───────────┘
         ▼                                     ▼
┌────────────────────┐               ┌────────────────────┐
│  Masked Self-Attention│              │  Masked Self-Attention│
└────────┬───────────┘               └────────┬───────────┘
         ▼                                     ▼
┌────────────────────┐               ┌────────────────────┐
│   Residual + Norm  │               │   Residual + Norm  │
└────────┬───────────┘               └────────┬───────────┘
         ▼                                     ▼
┌────────────────────┐               ┌────────────────────┐
│        FFN         │        →      │    MoE Layer       │
│   (单个 FFN 网络)   │               │  (Router + N 个 FFN)│
└────────┬───────────┘               └────────┬───────────┘
         ▼                                     ▼
┌────────────────────┐               ┌────────────────────┐
│   Residual + Norm  │               │   Residual + Norm  │
└────────────────────┘               └────────────────────┘

MoE(x) = Σᵢ g(x)ᵢ · Eᵢ(x)

import torch
import torch.nn.functional as F

def moe_forward(x, experts, router_weight, top_k=2):
    # x: [batch_size, seq_len, d_model]
    # router_weight: [d_model, num_experts]

    # 1. 计算每个专家的路由分数
    logits = F.linear(x, router_weight)  # [batch, seq, num_experts]

    # 2. 选择 top-k 个专家
    top_k_logits, top_k_indices = torch.topk(logits, top_k, dim=-1)

    # 3. 计算门控权重（仅对选中的专家做 softmax）
    top_k_gates = F.softmax(top_k_logits, dim=-1)

    # 4. 逐专家计算并加权求和
    output = torch.zeros_like(x)
    for k in range(top_k):
        expert_idx = top_k_indices[:, :, k]   # [batch, seq]
        gate = top_k_gates[:, :, k]           # [batch, seq]
        expert_out = experts[expert_idx](x)   # 选中专家的输出
        output += gate.unsqueeze(-1) * expert_out

    return output

理想分布:        实际坍缩分布:
E1: ████ 12.5%   E1: ████████████████████████ 80%
E2: ████ 12.5%   E2: ███ 10%
E3: ████ 12.5%   E3: █ 3%
E4: ████ 12.5%   E4: █ 2%
E5: ████ 12.5%   E5: █ 2%
E6: ████ 12.5%   E6: █ 1.5%
E7: ████ 12.5%   E7: █ 1%
E8: ████ 12.5%   E8: █ 0.5%

def load_balancing_loss(gate_logits, num_experts, top_k):
    # gate_logits: [batch * seq, num_experts]
    routing_weights = F.softmax(gate_logits, dim=-1)

    # 每个专家被选中的比例 fᵢ
    _, top_k_indices = torch.topk(routing_weights, top_k, dim=-1)
    mask = torch.zeros_like(routing_weights).scatter_(1, top_k_indices, 1)
    tokens_per_expert = mask.float().mean(dim=0)

    # 每个专家的平均路由权重 pᵢ
    router_prob_per_expert = routing_weights.mean(dim=0)

    # 辅助损失 = N * Σᵢ (fᵢ * pᵢ)
    # 当所有专家均匀负载时，此损失最小
    loss = num_experts * (tokens_per_expert * router_prob_per_expert).sum()
    return loss

def loss_free_balancing_routing(x, router_weight, bias, top_k, gamma=0.001):
    # 1. 计算原始路由分数（使用 Sigmoid）
    scores = torch.sigmoid(F.linear(x, router_weight))

    # 2. 在路由决策前，加上偏置项
    biased_scores = scores + bias  # bias: [num_experts]

    # 3. 基于偏置后的分数做 top-K 选择
    _, top_k_indices = torch.topk(biased_scores, top_k, dim=-1)

    # 4. 门控权重使用原始分数（不受偏置影响）
    gates = scores.gather(-1, top_k_indices)
    gates = gates / gates.sum(dim=-1, keepdim=True)

    return top_k_indices, gates

def update_bias(expert_load, target_load, bias, gamma=0.001):
    # 根据实际负载动态调整偏置
    for i in range(len(bias)):
        if expert_load[i] > target_load:
            bias[i] -= gamma   # 过载专家 → 降低偏置 → 减少被选概率
        elif expert_load[i] < target_load:
            bias[i] += gamma   # 冷门专家 → 提高偏置 → 增加被选概率

def router_z_loss(logits):
    # logits: [batch * seq, num_experts]
    # 对 logits 的 log-sum-exp 进行惩罚
    z = torch.logsumexp(logits, dim=-1).square().mean()
    return z

传统 MoE:  Token → Router → Top-2 of 8 Experts → Output

DeepSeekMoE:
  Token ──→ Router → Top-8 of 256 Routed Experts ──→ Weighted Sum ──→ Output
        └──→ Shared Expert (始终激活) ────────────────↗

┌────────────────────────────────────────────────────┐
│              Expert Parallelism 示意                 │
│                                                     │
│  GPU 0: Expert 0, 1    ┐                            │
│  GPU 1: Expert 2, 3    ├─ All-to-All 通信            │
│  GPU 2: Expert 4, 5    │  (token 按路由目标重分布)     │
│  GPU 3: Expert 6, 7    ┘                            │
│                                                     │
│  Token Batch:                                       │
│    [T1→E0,E3] [T2→E1,E7] [T3→E2,E5] [T4→E0,E3]   │
│                                                     │
│  All-to-All 后:                                     │
│    GPU 0: [T1,T4→E0] [T1,T4→E3]                    │
│    GPU 1: [T2→E1] [T2→E7]                           │
│    GPU 2: [T3→E2] [T3→E5]                           │
└────────────────────────────────────────────────────┘