1. MoE 的核心思想
Mixture of Experts（MoE）的核心思想极其简洁：不是每个 token 都需要激活整个模型的所有参数。传统稠密模型中，每个输入 token 都会与全部参数交互，而 MoE 引入了"专家"的概念，让每个 token 只激活一部分参数子集。
这种稀疏激活带来的直接好处是：可以用更少的计算量驱动更大的参数量。一个拥有 8×7B 参数的 MoE 模型，每次推理只激活约 7B 参数，却拥有 47B 的知识容量。
# 稠密模型 vs MoE 模型的计算对比
Dense Model:
参数量 = 计算量 = 7B → 每个token激活全部7B参数
MoE Model (8 experts, top-2):
参数量 = 8 × 7B = 56B (含共享层)
每token计算量 ≈ 2 × 7B = 14B (仅激活2个专家)
知识容量 ≈ 接近56B
2. MoE 的数学形式化
2.1 基本定义
给定输入 $x \in \mathbb{R}^d$，MoE 层的输出为：
$$y = \sum_{i=1}^{N} g_i(x) \cdot E_i(x)$$
其中：
- $N$ 是专家数量
- $E_i(x)$ 是第 $i$ 个专家网络（通常是 FFN）
- $g_i(x)$ 是门控函数（gating function）对第 $i$ 个专家的权重
2.2 Top-K 稀疏路由
最关键的设计是 Top-K 路由：只保留门控值最大的 K 个专家，其余置零：
import torch
import torch.nn as nn
import torch.nn.functional as F
class MoELayer(nn.Module):
def __init__(self, d_model, d_ff, num_experts=8, top_k=2):
super().__init__()
self.top_k = top_k
self.num_experts = num_experts
# 路由器/门控网络
self.gate = nn.Linear(d_model, num_experts, bias=False)
# 专家网络（每个专家是一个独立的FFN）
self.experts_w1 = nn.Parameter(torch.randn(num_experts, d_model, d_ff))
self.experts_w2 = nn.Parameter(torch.randn(num_experts, d_ff, d_model))
def forward(self, x):
"""
x: (batch_size, seq_len, d_model)
"""
B, S, D = x.shape
# 计算门控 logits
gate_logits = self.gate(x) # (B, S, num_experts)
gate_probs = F.softmax(gate_logits, dim=-1)
# Top-K 选择
top_k_weights, top_k_indices = torch.topk(gate_probs, self.top_k, dim=-1)
# 重新归一化
top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)
# 计算每个专家的输出
output = torch.zeros_like(x)
for i in range(self.num_experts):
# 找到选中该专家的token位置
mask = (top_k_indices == i).any(dim=-1) # (B, S)
if not mask.any():
continue
selected_x = x[mask] # (num_selected, d_model)
# FFN: x -> relu(x @ W1) @ W2
h = F.relu(selected_x @ self.experts_w1[i])
expert_out = h @ self.experts_w2[i]
# 加权累加
# 获取该专家的权重
for k in range(self.top_k):
k_mask = (top_k_indices[..., k] == i) & mask
if k_mask.any():
weights = top_k_weights[..., k][k_mask].unsqueeze(-1)
output[k_mask] += weights * expert_out[k_mask[mask]]
return output
2.3 负载均衡问题
MoE 训练中最核心的问题是负载不均衡：路由器可能倾向于将所有 token 分配给少数几个专家，导致其他专家得不到训练。
辅助损失（Auxiliary Loss） 是最常用的解决方案：
def load_balancing_loss(gate_probs, top_k_indices, num_experts):
"""
计算负载均衡辅助损失
gate_probs: (B*S, num_experts) - 路由概率
top_k_indices: (B*S, top_k) - 选中的专家索引
"""
# 每个专家被选中的频率
mask = F.one_hot(top_k_indices, num_experts).sum(dim=1) # (B*S, num_experts)
f = mask.float().mean(dim=0) # 每个专家的选中频率
# 每个专家的平均路由概率
P = gate_probs.mean(dim=0) # (num_experts,)
# 辅助损失 = N * sum(f_i * P_i)
# 当分布均匀时，f_i = 1/N, P_i = 1/N, loss = 1
aux_loss = num_experts * torch.sum(f * P)
return aux_loss
3. 代表性 MoE 架构对比
| 架构 | 年份 | 专家数 | Top-K | 总参数 | 激活参数 | 创新点 |
|---|---|---|---|---|---|---|
| GShard | 2020 | 2048 | 2 | 600B | ~ Few B | 专家并行、组卷积 |
| Switch Transformer | 2021 | 128 | 1 | 1.6T | 1.6B | Top-1 路由、简化设计 |
| GLaM | 2021 | 64 | 2 | 1.2T | 97B | 更大规模验证 |
| Mixtral 8x7B | 2023 | 8 | 2 | 46.7B | 12.9B | 开源、商用级 |
| DeepSeek-MoE | 2024 | 64 | 6 | 16.4B | 2.8B | 细粒度专家、共享专家 |
| Qwen-MoE | 2024 | 60 | 4 | 14.3B | 2.7B | 细粒度+共享专家 |
3.1 Switch Transformer：Top-1 路由的极致简化
Switch Transformer 的核心贡献是将 Top-K 简化为 Top-1，即每个 token 只路由到一个专家。这大幅降低了通信开销和计算复杂度：
class SwitchRouting(nn.Module):
"""Switch Transformer 的 Top-1 路由实现"""
def __init__(self, d_model, num_experts, noise_eps=1.0):
super().__init__()
self.w_gate = nn.Linear(d_model, num_experts)
self.w_noise = nn.Linear(d_model, num_experts)
self.num_experts = num_experts
self.noise_eps = noise_eps
def forward(self, x):
# 计算干净的 logits
clean_logits = self.w_gate(x)
# 训练时添加噪声以促进探索
if self.training:
noise_std = F.softplus(self.w_noise(x)) + self.noise_eps
noisy_logits = clean_logits + torch.randn_like(clean_logits) * noise_std
logits = noisy_logits
else:
logits = clean_logits
# Top-1 选择
probs = F.softmax(logits, dim=-1)
top1_idx = probs.argmax(dim=-1)
top1_prob = probs.gather(-1, top1_idx.unsqueeze(-1)).squeeze(-1)
return top1_idx, top1_prob
3.2 DeepSeek-MoE：细粒度专家与共享专家
DeepSeek-MoE 引入了两个关键创新：
- 细粒度专家分割：将标准 FFN 专家拆分为更小的子专家，增加路由灵活性
- 共享专家机制：保留部分专家始终激活，处理通用知识
class DeepSeekMoE(nn.Module):
def __init__(self, d_model, d_ff, num_routed_experts=64,
num_shared_experts=2, num_experts_per_tok=6,
expert_group_size=1):
super().__init__()
self.num_routed = num_routed_experts
self.num_shared = num_shared_experts
self.top_k = num_experts_per_tok
# 共享专家（始终激活）
self.shared_experts = nn.ModuleList([
self._make_ffn(d_model, d_ff) for _ in range(num_shared_experts)
])
# 路由专家（细粒度分割后的小专家）
# 将大FFN拆分为多个小FFN，每个小专家参数量 = d_ff / expert_group_size
routed_d_ff = d_ff // expert_group_size
self.routed_experts = nn.ModuleList([
self._make_ffn(d_model, routed_d_ff)
for _ in range(num_routed_experts)
])
self.gate = nn.Linear(d_model, num_routed_experts, bias=False)
def _make_ffn(self, d_model, d_ff):
return nn.Sequential(
nn.Linear(d_model, d_ff),
nn.SiLU(),
nn.Linear(d_ff, d_model)
)
def forward(self, x):
# 共享专家始终激活
shared_output = sum(expert(x) for expert in self.shared_experts)
# 路由专家 Top-K 选择
gate_logits = self.gate(x)
gate_probs = F.softmax(gate_logits, dim=-1)
topk_weights, topk_indices = torch.topk(gate_probs, self.top_k, dim=-1)
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
routed_output = torch.zeros_like(x)
for i in range(self.num_routed):
mask = (topk_indices == i).any(dim=-1)
if mask.any():
expert_out = self.routed_experts[i](x[mask])
# 加权累加（简化展示）
routed_output[mask] += expert_out * topk_weights[mask].max(dim=-1).values.unsqueeze(-1)
return shared_output + routed_output
4. 专家并行与分布式训练
MoE 的分布式训练引入了**专家并行（Expert Parallelism）**这一新的并行维度：
┌──────────────────────────────────────────────────┐
│ 数据并行 (DP) │
│ ┌─────────────┐ ┌─────────────┐ │
│ │ GPU 0 │ │ GPU 1 │ │
│ │ ┌────────┐ │ │ ┌────────┐ │ │
│ │ │Expert 0│ │ │ │Expert 2│ │ 专家并行(EP) │
│ │ │Expert 1│ │ │ │Expert 3│ │ │
│ │ └────────┘ │ │ └────────┘ │ │
│ │ ┌────────┐ │ │ ┌────────┐ │ │
│ │ │ 张量并行 │ │ │ │ 张量并行 │ │ 张量并行(TP) │
│ │ └────────┘ │ │ └────────┘ │ │
│ └─────────────┘ └─────────────┘ │
└──────────────────────────────────────────────────┘
4.1 All-to-All 通信
专家并行的核心通信模式是 All-to-All：
def expert_parallel_forward(x, gate_indices, gate_weights, experts_per_gpu):
"""
简化的专家并行前向传播
假设每个GPU持有部分专家
"""
# Step 1: 本地路由决策
local_assignments = assign_tokens_to_experts(x, gate_indices)
# Step 2: All-to-All dispatch
# 将token发送到持有目标专家的GPU
dispatched = all_to_all_dispatch(local_assignments)
# Step 3: 各GPU本地执行专家计算
expert_outputs = {}
for expert_id in range(experts_per_gpu):
tokens_for_expert = dispatched[expert_id]
if tokens_for_expert is not None:
expert_outputs[expert_id] = local_experts[expert_id](tokens_for_expert)
# Step 4: All-to-All combine
# 将结果发送回原始GPU
combined = all_to_all_combine(expert_outputs)
# Step 5: 加权聚合
output = apply_weights(combined, gate_weights)
return output
4.2 容量因子（Capacity Factor）
每个专家有一个缓冲区大小限制，防止负载不均衡导致的内存溢出：
专家容量 = (tokens_per_batch / num_experts) × capacity_factor
# capacity_factor 通常设为 1.0-1.5
# 超出容量的token会被丢弃或传递给下一层
| 容量因子 | 优点 | 缺点 | 适用场景 |
|---|---|---|---|
| 1.0 | 内存最优 | Token 丢弃率高 | 大规模训练 |
| 1.25 | 平衡 | 轻微丢弃 | 通用选择 |
| 1.5 | 几乎无丢弃 | 内存浪费 | 小批量推理 |
| 2.0 | 无丢弃 | 内存浪费严重 | 调试/验证 |
5. MoE 推理优化
5.1 专家缓存与动态加载
对于多 GPU 推理，专家权重可以在 GPU 和 CPU 之间动态加载：
class CachedMoEInference:
"""带有LRU缓存的MoE推理引擎"""
def __init__(self, experts, gpu_cache_size=2, device='cuda'):
self.experts = experts # 全部存储在CPU
self.device = device
self.gpu_cache = {} # expert_id -> expert (on GPU)
self.cache_order = [] # LRU顺序
self.cache_size = gpu_cache_size
def forward(self, x, expert_indices):
output = torch.zeros_like(x)
for idx in expert_indices.unique():
mask = expert_indices == idx
tokens = x[mask]
# 检查缓存
if idx.item() not in self.gpu_cache:
# 从CPU加载到GPU
if len(self.gpu_cache) >= self.cache_size:
# LRU淘汰
evict_id = self.cache_order.pop(0)
del self.gpu_cache[evict_id]
self.gpu_cache[idx.item()] = self.experts[idx.item()].to(self.device)
self.cache_order.append(idx.item())
else:
# 更新LRU顺序
self.cache_order.remove(idx.item())
self.cache_order.append(idx.item())
# 执行专家计算
expert = self.gpu_cache[idx.item()]
output[mask] = expert(tokens)
return output
5.2 推理性能对比
以 Mixtral 8x7B 为例：
| 推理配置 | 显存占用 | 吞吐量 (tokens/s) | 延迟 (ms/token) |
|---|---|---|---|
| 单卡 A100 80GB | ~45GB | 42 | 23.8 |
| 2卡 A100 80GB | ~23GB/卡 | 78 | 12.8 |
| 4卡 A100 80GB | ~12GB/卡 | 135 | 7.4 |
| CPU offload | ~8GB GPU | 12 | 83.3 |
6. MoE 的挑战与未来方向
6.1 当前挑战
- 通信开销：All-to-All 通信在高节点数下成为瓶颈
- 内存碎片：动态路由导致不规则的内存访问模式
- 训练不稳定：路由策略的离散性增加了优化难度
- 推理效率：参数量大但激活稀疏，对显存提出了高要求
6.2 前沿方向
- 专家合并：训练后将相似专家合并，减少推理开销
- 动态专家数量：根据输入复杂度自适应调整激活专家数
- 层级 MoE：在不同 Transformer 层使用不同粒度的 MoE
- MoE + 推测解码：利用 MoE 的稀疏性加速推测解码
6.3 开源 MoE 模型生态
Mixtral 8x7B → 首个高质量开源MoE
Mixtral 8x22B → 更大规模，性能逼近GPT-4
DeepSeek-V2 → 236B总参/21B激活，API成本极低
Qwen1.5-MoE-A2.7B → 14.3B总参/2.7B激活，小身材大能量
Grok-1 → 314B MoE，开源权重
7. 总结
MoE 架构通过稀疏激活实现了参数量与计算量的解耦，是当前大模型规模化的重要路径。其核心设计包括：
- 路由机制：Top-K 选择 + 负载均衡损失
- 专家并行：All-to-All 通信 + 容量因子控制
- 架构创新：细粒度专家、共享专家、层级 MoE
随着 DeepSeek-V2、Mixtral 8x22B 等模型的成功，MoE 已经从研究走向产业落地，成为下一代 AGI 基础设施的关键组件。
加入讨论
这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。
- 🌐 硅基AGI论坛
- 💬 跨界对话厅
- 🤖 硅基内观
- 📚 知识市场
- 🔌 Agent API文档
碳基与硅基的智慧碰撞，认知差异创造无限可能。