在大语言模型参数量从百亿迈向千亿、万亿的时代,一个核心矛盾愈发凸显:模型容量越大效果越好,但推理和训练的算力开销呈线性甚至二次增长。Mixture of Experts(MoE,混合专家)架构通过稀疏激活机制破解了这一困局——模型总参数量可达数千亿,但每次前向传播只激活其中一小部分”专家”模块,从而在保持推理速度不变的前提下大幅扩展模型容量。本文将从 MoE 的数学原理出发,深入剖析路由机制、负载均衡、训练工程和推理部署的完整技术栈。
一、MoE 的核心思想:用稀疏性换取容量
传统稠密(Dense)Transformer 中,每个 token 都会经过所有参数的计算。而 MoE 架构将前馈网络(FFN)层替换为多个并行的”专家”网络,由一个门控路由器(Gating Router)决定每个 token 应该被发送给哪个专家。假设模型有 N 个专家,每次只激活 Top-k 个,那么单次前向传播的计算量仅为稠密模型的 k/N。
以 Mixtral 8x7B 为例:模型总参数量约 46.7B,但每个 token 只激活约 12.9B 参数(2个专家),推理速度与 14B 稠密模型相当,效果却接近 70B 级别。这种”以存储换计算”的策略是 MoE 的本质。
1.1 门控路由的数学表达
给定输入 token 的隐状态
1 | h |
,门控路由器计算每个专家的得分:
1
2
3
4
5
6
7
8 # 标准 Top-K 路由
def gating(h, W_gate, num_experts=8, top_k=2):
# h: [batch * seq_len, dim]
# W_gate: [dim, num_experts]
logits = h @ W_gate # [B, num_experts]
weights, indices = torch.topk(logits, top_k, dim=-1)
weights = torch.softmax(weights, dim=-1)
return weights, indices
输出
1 | weights |
是归一化后的路由权重,
1 | indices |
是被选中的专家编号。随后 token 只会被发送到对应的专家 FFN 进行计算,最终输出按权重加权求和。
1.2 专家容量与桶式分配
为了避免某些专家被分配过多 token 导致负载不均,MoE 引入了专家容量(Expert Capacity)的概念:
1 capacity = (batch_size * seq_len / num_experts) * capacity_factor
其中
1 | capacity_factor |
通常设为 1.0-1.25。超出容量的 token 会被”溢出”(overflow),直接通过残差连接跳过该 MoE 层。这保证了每个专家的计算量有上界,但代价是部分信息丢失。在训练中,溢出率通常控制在 1% 以下。
二、路由机制的演进:从 Switch Transformer 到 Expert Choice
2.1 Switch Transformer:Top-1 路由的极简方案
Google 在 2021 年提出的 Switch Transformer 将 Top-k 设为 1,即每个 token 只路由到一个专家。这极大简化了路由逻辑和通信开销,但代价是表达能力下降和更高的负载不均衡风险。其路由公式为:
1
2
3
4 # Switch Transformer 路由
logits = h @ W_gate # [B, num_experts]
probabilities = torch.softmax(logits, dim=-1)
route_prob, route_idx = torch.max(probabilities, dim=-1)
Switch Transformer 还引入了 load balancing loss(负载均衡损失),强制 token 尽可能均匀地分布在各专家上:
1
2
3
4
5
6
7
8
9
10 # 负载均衡损失 (Switch Transformer)
# f_i: 每个专家被选中的 token 比例
# P_i: 每个专家的平均路由概率
def load_balancing_loss(gate_probs, expert_mask, num_experts):
# gate_probs: [B, num_experts] softmax probabilities
# expert_mask: [B, num_experts] one-hot selected
f = expert_mask.float().mean(dim=0) # 每个专家的负载比例
P = gate_probs.mean(dim=0) # 每个专家的平均概率
loss = num_experts * (f * P).sum()
return loss
该损失的最小值在所有专家均匀分布时取得,值为 1.0。训练时将其加权(通常权重为 0.01)加入总损失,引导路由器学会均衡分配。
2.2 Expert Choice 路由:让专家选 token
传统的 token-choice 路由中,路由器为每个 token 选择专家,容易导致”热门专家”过载。Expert Choice 反转了这一逻辑:让每个专家从所有 token 中选择自己想处理的 Top-k 个。这天然实现了负载均衡——每个专家的容量严格相等。
1
2
3
4
5
6
7
8 # Expert Choice 路由
def expert_choice_routing(h, W_gate, num_experts, capacity):
logits = h @ W_gate # [B, num_experts]
# 转置后对每个专家选 top-k token
transposed = logits.t() # [num_experts, B]
weights, indices = torch.topk(transposed, capacity, dim=-1)
weights = torch.softmax(weights, dim=-1)
return weights, indices
Expert Choice 在训练效率上显著优于 token-choice,但也带来了一个新问题:某些 token 可能不被任何专家选中,导致信息丢失。这在推理时尤为棘手,因为无法保证每个输入 token 都被处理。
2.3 DeepSeek 的共享专家机制
DeepSeek-MoE 引入了共享专家(Shared Expert)的创新设计:将 N 个专家拆分为
1 | N_s |
个共享专家和
1 | N_r |
个路由专家。共享专家始终对所有 token 激活,负责提取通用特征;路由专家则按需激活,负责捕获 specialized 特征。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32 # DeepSeek-MoE 前馈层
class DeepSeekMoEFFN(nn.Module):
def __init__(self, dim, d_ff, num_shared=2, num_routed=64, top_k=6):
super().__init__()
self.shared_experts = nn.ModuleList([
FFN(dim, d_ff) for _ in range(num_shared)
])
self.routed_experts = nn.ModuleList([
FFN(dim, d_ff) for _ in range(num_routed)
])
self.gate = nn.Linear(dim, num_routed)
self.top_k = top_k
def forward(self, h):
# 共享专家:所有 token 都经过
shared_out = sum(expert(h) for expert in self.shared_experts)
# 路由专家:Top-K 选择
logits = self.gate(h)
weights, indices = torch.topk(logits, self.top_k, dim=-1)
weights = torch.softmax(weights, dim=-1)
routed_out = torch.zeros_like(h)
for i in range(self.top_k):
expert_idx = indices[:, i] # [B]
w = weights[:, i:i+1] # [B, 1]
for e in range(len(self.routed_experts)):
mask = (expert_idx == e)
if mask.any():
routed_out[mask] += w[mask] * self.routed_experts[e](h[mask])
return shared_out + routed_out
这种设计有效减少了路由专家之间的知识冗余,让每个专家能更专注于特定领域的特征。DeepSeek-V2 的 236B 模型使用了 160 个路由专家 + 2 个共享专家,每次激活 6 个路由专家,总激活参数约 21B。
三、MoE 训练的核心工程挑战
3.1 All-to-All 通信与专家并行
当专家数量远大于 GPU 数量时,单卡无法容纳所有专家,必须引入专家并行(Expert Parallelism)。这意味着不同专家分布在不同 GPU 上,token 需要通过 All-to-All 通信发送到持有目标专家的 GPU。这是 MoE 训练中最复杂的部分。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17 # 伪代码:专家并行下的前向传播
def moe_forward_expert_parallel(h, gate, experts, num_experts_per_gpu):
# 1. 本地计算路由
weights, indices = gate(h)
# 2. All-to-All dispatch:将 token 发到目标 GPU
# h: [local_tokens, dim] -> [dispatched_tokens, dim]
dispatched = all_to_all_dispatch(h, indices)
dispatched_weights = all_to_all_dispatch(weights, indices)
# 3. 本地专家计算
local_output = experts(dispatched) * dispatched_weights
# 4. All-to-All combine:将结果送回原 GPU
output = all_to_all_combine(local_output, indices)
return output
All-to-All 通信的延迟是 MoE 训练效率的瓶颈。在实际工程中,通常会采用以下优化:
- 通信重叠:将 All-to-All 通信与注意力层的计算重叠执行
- 桶式排序:按目标 GPU 排序 token,减少通信碎片
- FP16 通信:使用半精度进行 All-to-All 传输以减半通信量
- 分组专家:将专家分为若干组,组内共享路由以减少跨组通信
3.2 负载均衡:不只是均匀分配那么简单
负载均衡损失虽然能让 token 均匀分布,但过度追求均匀可能损害模型效果——某些专家可能天然需要处理更多 token。DeepSeek-V2 提出了无辅助损失的负载均衡策略,通过为每个专家引入一个偏置项(bias)来动态调节路由偏好:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19 # DeepSeek 无辅助损失负载均衡
class BiasAdjustedGate(nn.Module):
def __init__(self, dim, num_experts, top_k):
self.gate = nn.Linear(dim, num_experts)
self.bias = nn.Parameter(torch.zeros(num_experts))
self.top_k = top_k
# bias 不参与梯度回传,而是按更新频率调整
self.update_freq = 1000 # 每 1000 步更新一次
def forward(self, h):
logits = self.gate(h) + self.bias # 加入偏置
weights, indices = torch.topk(logits, self.top_k, dim=-1)
weights = torch.softmax(weights, dim=-1)
return weights, indices
def update_bias(self, expert_loads):
# 负载高的专家降低 bias,负载低的提高 bias
mean_load = expert_loads.mean()
self.bias.data += self.update_speed * (mean_load - expert_loads)
这种动态偏置调整避免了辅助损失对主任务的梯度干扰,在 DeepSeek-V2 的实验中取得了比传统辅助损失更好的效果。
3.3 精度问题:MoE 训练中的梯度噪声
MoE 模型的路由操作是离散的(Top-K 选择),导致梯度无法直接回传到门控网络。标准的解决方案是直通估计(Straight-Through Estimator),但引入的梯度噪声在训练初期可能导致路由不稳定。以下是一些常见的缓解措施:
| 策略 | 原理 | 适用场景 |
|---|---|---|
| 路由 z-loss | 惩罚路由 logits 的绝对值,防止路由概率过度集中 | 训练初期不稳定 |
| 温度退火 | softmax 温度从高到低退火,初期鼓励探索 | 路由坍缩问题 |
| 专家丢弃 | 训练时随机丢弃部分专家,迫使路由器学会冗余分配 | 过拟合与泛化 |
| JS 散度正则 | 用 Jensen-Shannon 散度替代 KL 散度做负载均衡 | 更温和的均衡约束 |
四、MoE 推理部署:从显存优化到动态路由
4.1 显存挑战与量化方案
MoE 模型最大的推理挑战是显存占用。Mixtral 8x7B 虽然每次只激活 2 个专家,但所有 8 个专家的权重都必须驻留在显存中。对于更大的 MoE 模型(如 DeepSeek-V2 的 236B),单卡显存根本无法容纳全部权重。
工程上通常采用以下方案组合:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22 # MoE 推理显存优化方案
# 方案 1:专家权重量化
# 将专家 FFN 权重从 FP16 量化到 INT4
# Mixtral 8x7B: 93GB (FP16) -> 25GB (INT4)
# 方案 2:专家卸载到 CPU/SSD
# 仅保留激活频率最高的专家在 GPU 上
# 低频专家按需从 CPU 内存或 SSD 加载
# 方案 3:分组量化 + 专家共享
# 量化后同质专家共享码本,进一步压缩
# 方案 4:专家剪枝
# 训练后剪掉使用率极低的专家
# 需要在少量数据上做路由统计评估
def prune_unused_experts(model, threshold=0.01):
expert_usage = collect_expert_usage_stats(model, calibration_data)
pruned = [i for i, usage in enumerate(expert_usage)
if usage < threshold]
print(f"Pruning {len(pruned)} experts: {pruned}")
return remove_experts(model, pruned)
4.2 MoE 与 KV Cache 的协同优化
在自回归生成中,MoE 架构的 KV Cache 与稠密模型相同——注意力层的参数不变,只有 FFN 层被替换为 MoE。这意味着 KV Cache 的优化策略(如 PagedAttention、KV Cache 量化)可以直接复用。但 MoE 推理有一个独特的优化机会:专家级批处理。
当多个请求的 token 被路由到同一个专家时,可以将它们合并为一个 batch 进行计算,大幅提高 GPU 利用率。vLLM 等推理框架已实现了这一优化:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22 # 专家级动态批处理(vLLM 中的 MoE 优化)
def expert_batched_forward(tokens, router, experts):
# 1. 路由所有 token
weights, expert_indices = router(tokens)
# 2. 按 expert 分组
expert_groups = defaultdict(list)
for i, (w, ei) in enumerate(zip(weights, expert_indices)):
for e, w_e in zip(ei, w):
expert_groups[e].append((i, w_e))
# 3. 每个 expert 批量处理其分配到的 token
outputs = [torch.zeros_like(tokens) for _ in tokens]
for expert_id, items in expert_groups.items():
token_idx = [item[0] for item in items]
token_weights = [item[1] for item in items]
batched_tokens = tokens[token_idx]
expert_output = experts[expert_id](batched_tokens)
for i, (idx, w) in enumerate(zip(token_idx, token_weights)):
outputs[idx] += w * expert_output[i]
return outputs
4.3 推理时专家切换的开销
当使用专家卸载方案时,推理过程中切换专家会引入额外的 PCIe 传输延迟。一个典型的优化策略是专家预取:利用路由器的输出提前预判下一步可能需要的专家,在当前 token 计算期间异步加载下一批专家权重到 GPU。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25 # 专家预取策略
class ExpertPrefetcher:
def __init__(self, experts, gpu_cache_size=4):
self.experts = experts # 存储在 CPU 上的所有专家
self.gpu_cache = {} # GPU 上缓存的活跃专家
self.cache_size = gpu_cache_size
self.prefetch_queue = []
def predict_next_experts(self, current_indices, history):
# 基于历史路由模式预测下一批可能激活的专家
# 简单实现:取当前 top-k 之外的最高概率专家
predicted = []
for indices in history[-3:]:
predicted.extend(indices.tolist())
return list(set(predicted))[:self.cache_size]
def prefetch(self, predicted_experts):
for e in predicted_experts:
if e not in self.gpu_cache:
self.gpu_cache[e] = self.experts[e].to('cuda')
if len(self.gpu_cache) > self.cache_size:
# LRU 淘汰
evict = min(self.gpu_cache,
key=lambda k: self.access_time.get(k, 0))
del self.gpu_cache[evict]
五、MoE 模型架构对比与选型建议
当前主流的 MoE 模型各有侧重,下面从架构设计、训练规模和部署友好度三个维度进行对比:
| 模型 | 总参数 | 激活参数 | 专家数 | Top-K | 特色设计 |
|---|---|---|---|---|---|
| Mixtral 8x7B | 46.7B | 12.9B | 8 | 2 | 标准 Top-2 路由 |
| DeepSeek-V2 | 236B | 21B | 160+2共享 | 6 | 共享专家 + 无辅助损失均衡 |
| Qwen2-MoE | 57B (A14B) | 14B | 60+4共享 | 4 | 细粒度专家 + 共享专家 |
| Switch-Base | 7.4B | 1.1B | 128 | 1 | Top-1 极简路由 |
| Grok-1 | 314B | ~86B | 8 | 2 | 大颗粒度专家 |
从工程选型角度,以下几点值得重点关注:
- 细粒度专家(如 DeepSeek 的 160 专家)比粗粒度专家(如 Mixtral 的 8 专家)在效果上更优,因为更细的分工让每个专家能高度特化。但代价是路由逻辑更复杂、通信开销更大。
- 共享专家是当前 MoE 设计的标配,几乎所有的最新 MoE 模型都采用了这一机制。它有效解决了通用知识存储与 specialized 能力之间的平衡问题。
- 部署友好度:专家数量越多,显存压力越大但单次计算量不变。如果部署环境显存有限,优先选择激活参数比适中的模型(如 Qwen2-57B-A14B)。
六、未来展望:从静态路由到动态架构
当前 MoE 架构的路由策略大多在训练阶段固定,推理时无法动态调整。一些前沿研究方向正在突破这一限制:
动态专家数量:根据输入复杂度自适应调整 Top-K 值。简单 token 只用 1 个专家,复杂推理 token 可激活更多专家。这需要在训练时引入可学习的路由预算机制。
层级 MoE:在注意力层也引入 MoE 机制,不同的注意力头可以被路由到不同的”注意力专家”,进一步增强模型的容量。但这也带来了更高的通信开销和实现复杂度。
跨层专家共享:不同层的专家可以共享权重,减少总参数量。这种”深度共享”策略在小规模 MoE 模型上已展示出良好效果,但在大规模模型上的有效性仍有待验证。
MoE 架构正在从研究原型走向工业标配。理解其路由机制、训练技巧和部署优化,是每一位从事大模型工程的技术人员的必备知识。随着 DeepSeek、Qwen 等国产 MoE 模型的快速迭代,我们有理由相信,稀疏激活将成为下一代基础模型的核心范式。
汤不热吧