一文看懂大模型中的 MoE：从朴素思想到 1.6 万亿参数落地

Jay Kay

1133人浏览 · 2025-09-12 16:22:23

Jay Kay · 2025-09-12 16:22:23 发布

一文看懂大模型中的 MoE：从朴素思想到 1.6 万亿参数落地

如果文章太长，记住一句话就行：
MoE = 把大模型“拆成”一群小专家，只叫醒最有学问的那两个，剩下继续睡觉，于是同样算力下模型可以更大、训得更快。

1. 为什么又要"拆"模型？

2019 年后，大模型基本遵循「力大砖飞」定律：参数翻倍 ⇒ 指标提升。
但 GPU 算力/显存增速远远赶不上参数膨胀速度。
稠密（Dense）Transformer 每一层、每一个 token 都要激活全部参数，成本 O(d²·L) 逃不掉。

MoE 给出的思路很朴素：

别每次把全班同学喊醒，只让最懂这道题的那 1~2 位同学回答即可。

于是同样 FLOPs 预算，可以把模型宽度（专家数）横向拉爆，而计算量几乎不变——这就是 条件计算（Conditional Computation）。

2. MoE 的两大核心零件

零件	作用	常见实现
专家（Experts）	实际干活的前馈网络	8×MLP、SwiGLU、甚至子-MoE
门控/路由（Gating / Router）	决定喊哪几位专家	Top-K Softmax + 噪声 + Load-Balance Loss

把 Transformer 里的 每个 FFN 层 整体替换成一个 MoE 层即可——其余自注意力、残差、Layernorm 完全不动。

3. 稀疏 vs 稠密——数字说话

模型	总参数量	每 token 激活量	训练速度	推理吞吐
稠密 12B	12 B	12 B	1×	1×
MoE 47B (8×7B)	47 B	12 B	2.5×	1.5×

参数变 4 倍，训练快 2.5 倍，推理反而更快（并行专家 + 共享 Attention）。

代价：显存必须吃下 47 B 参数——所以 MoE 是「显存换 FLOPs」的典型代表。

4. 容量因子（Capacity Factor）（面试常问）

先算平均每个专家应该吃多少 token：
avg = batch_size * seq_len / num_experts
再乘「容量因子 CF」得到槽位上限：
capacity = ceil(avg × CF)

CF	含义	利弊
1.0	刚好平均数	通信/显存最小，易溢出
1.25	多留 25% 缓冲	Switch、Mixtral 默认选择
2.0	非常宽松	几乎不丢 token，显存翻倍

溢出 token 要么：

直接丢弃（加 auxiliary loss 惩罚）；
残差连接跳回主分支。

5. 负载均衡——为什么一定要 aux loss？

门控天生「偏科」：一开始谁随机得好，后面就越来越被青睐 → 其他专家得不到梯度 → 模型崩。
所以所有 MoE 论文都会加一条 Load-Balance Loss，强制让硬路由计数 ≈ 软路由概率，公式极简：

aux_loss = mean( (num_tokens_sent_to_expert / total_tokens) *
                 (gate_prob_sum_over_tokens) ) * num_experts

训练时把 aux_loss × α（α≈1e-2）加进总 loss 即可。

6. 微调 MoE 的坑与解

问题	现象	缓解方案
过拟合	小任务验证集掉点	调高稀疏层 dropout；调小 batch/调高 lr
显存爆炸	47B 全量加载	只训专家层，冻结共享 Attention；或用 LoRA
泛化差	SuperGLUE 不如稠密	先多任务指令微调（Flan-T5 路线），再下游任务

7. 工程落地三板斧

GroupGEMM + permute/un-permute
把去同一专家的 token 物理排在一起 → 一个 torch.bmm 打满 TensorCore。
专家并行（Expert Parallelism）
8 张卡各放 1 个专家，token 通过 all-to-all 通信 scatter/gather，单卡显存瞬间降 8 倍。
CUDA Kernel 融合
Megablocks、FasterMoE 把「路由+permute+计算+un-permute」写进一个 kernel，掉帧减少 30-50%。

8. 开源玩家速览

模型	规模	特色
Switch-Transformer	1.6T/2048 专家	Top-1 + CF=1.0，T5 结构
GLaM	1.2T/64 专家	仅用 1/3 能耗打平 GPT-3
Mixtral 8×7B	47B/8 专家	开源可商用，Top-2，LLaMA 生态
OpenMoE	社区复刻	基于 Llama，PyTorch 友好

9. 10 行代码实现「带容量限制」的 MoE

class Top1RouterWithCapacity(nn.Module):
    def __init__(self, d_model, n_exp, cf=1.25):
        super().__init__()
        self.gate = nn.Linear(d_model, n_exp)
        self.cf = cf
        self.n_exp = n_exp

    def forward(self, x):                       # x: (B,T,C)
        B, T = x.shape[:2]
        flat = x.view(-1, x.size(-1))
        logit = self.gate(flat)                 # (B*T, n_exp)
        score, idx = logit.max(dim=-1)          # Top-1
        cap = int((B*T)/self.n_exp * self.cf)
        # 简易溢出丢弃
        kept = torch.randperm(B*T, device=x.device)[:cap*self.n_exp]
        mask = torch.zeros_like(idx, dtype=torch.bool)
        mask[kept] = True
        idx[~mask] = -1                         # 标记丢弃
        return idx, score, self._aux_loss(logit.softmax(-1), idx)

    def _aux_loss(self, prob, idx):
        # 负载均衡损失（略）
        return 0.01 * (prob.mean(0) * (idx>=0).float().mean()).sum()

把 idx 喂给 GroupGEMM 即可上线。

10. 总结 & 什么时候选 MoE

训练阶段：同样 FLOPs → 更大模型 → 更低验证 loss；
推理阶段：同样延迟 → 更多参数 → 更好效果；
代价：显存占用 ≈ 总参数量；需要多机多卡 + 高带宽。

如果你有多台 A100/H100、InfiniBand 任性地跑，MoE 是性价比最高的放大器；
只有单卡 24G 玩本地部署，还是乖乖稠量化 + LoRA 吧。

参考资料

[1] Fedus et al. Switch Transformer: Scaling to Trillion Parameter Models
[2] Lepikhin et al. GShard: Scaling Giant Models with Conditional Computation
[3] Mixtral 8×7B release note – mistral.ai
[4] MegaBlocks GitHub – Stanford
[5] FairSeq MoE – Meta