第 16 讲 · 混合专家（MoE）· Mixture of Experts

对应 Maps to：CMU 10-423 / 10-623 / 10-723 Generative AI（Spring 2026；Matt Gormley、Pat Virtue），Lecture 16 · Mon 16-Mar · Mixture of Experts（slides：lecture16-moe）。 阅读材料：Shazeer et al. (2017), Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer；Fedus et al. (2021), Switch Transformers；Jiang et al. (2024), Mixtral of Experts；Muennighoff et al. (2024), OLMoE；Clark et al. (2022), Unified Scaling Laws for Routed Language Models。 依据说明：本讲按原课第 16 讲公开的 slides 展开：MoE 的时间线；LLM 里的线性层（GPT-3 各部分的参数、每层的参数怎样算）；稀疏门控 MoE；把一个线性层 / 前馈层拆成专家；稠密 MoE；稀疏门控（top-k、计算门控与 MoE 的例子、带噪声的 top-k、Mixtral 的 SwiGLU 专家、初始化）；专家并行（两个容量的例子）；平衡专家的利用率（负载均衡项、路由器 z-loss）；激活参数；Mixtral vs Llama-2；OLMoE 的超参数；性能 vs 代价；专家数怎么选；路由模型的规模定律。所有实验是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture16/（4 个脚本 + 1 个模块；都在 2 分钟以内） 预计阅读时间：2.5 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标 Learning goals

学完这一讲，你应该能够：

1. 写出稠密 MoE 和稀疏门控 MoE 的公式，手算一个 top-k 门控的例子；说明”把一个前馈层拆成几个专家”在什么条件下和原来等价。
2. 区分全部参数和激活参数，算出 Mixtral、OLMoE 的这两个数；说明 MoE 用什么换了什么（显存换计算）。
3. 解释为什么需要负载均衡：路由器会”偏心”；写出负载均衡项和路由器 z-loss，说出它们各管什么。
4. 描述专家并行：容量、容量因子、被丢弃的 token、all-to-all 通信。
5. 说出 MoE 适合什么、不适合什么（“记事实”vs”多算几步”），以及专家数的取舍。

0.2 先修与本讲范围 Prerequisites & scope

0.3 本讲地图 Map

§1  动机：前馈层占了大部分参数
§2  从前馈层到专家：稠密 MoE、稀疏门控 MoE                            ← 代码 01
§3  激活参数：Mixtral、OLMoE                                         ← 代码 01
§4  负载均衡：负载均衡项、路由器 z-loss                                ← 代码 02
§5  性能 vs 代价、专家数、路由模型的规模定律                            ← 代码 03
§6  专家并行：容量与通信                                              ← 代码 04
§7–11  小结、误区、练习、延伸阅读、下一讲

本讲的语言模型实验都在第 15 讲的 Zipf 事实世界里（4096 个实体 × 4 个属性，熵精确已知，超额损失 = 和真实分布的 KL，精确可算）：这是一个”要记很多事实”的任务，正好用来看”参数多、计算不多”的 MoE 有没有用。

1. 动机：前馈层占了大部分参数 · The linear layer in LLMs

slides 先放了一张 MoE 的时间线（Cai et al., 2024 的综述）：从 1991 年 Jacobs 等人的”自适应的局部专家混合”，到 2017 年 Shazeer 等人把稀疏门控 MoE 用到 LSTM 语言模型上，再到 GShard、Switch Transformer、Mixtral、DeepSeek、OLMoE……MoE 已经是大模型的标准做法之一。

为什么是前馈层？ slides：“Transformer LLM 一半以上的参数通常在前馈层里”。GPT-3（vocab 50,257、$d_{\text{model}}$ 12,288、96 个头、96 层、$d_{\text{ff}}$ 49,152）——代码 01 第 1 部分用 P1 的 param_count 重算 slides 的表：

每层的参数怎样算（slides 的问题）：LayerNorm 是 2d（缩放和平移）；嵌入是 V × d；注意力是 $W_Q, W_K, W_V, W_O$ 四个 d × d 的矩阵加偏置，$4d^2 + 4d$ ≈ 6.04 亿；前馈层是 $W_1$（d × $d_{\text{ff}}$）、$W_2$（$d_{\text{ff}}$ × d）加偏置，$2 d\, d_{\text{ff}} + d_{\text{ff}} + d$ ≈ 12.08 亿（$d_{\text{ff}} = 4d$ 时正好是注意力的两倍）。前馈层占全部参数的 66%。要让模型”大很多、但每个 token 的计算不变”，动的就是这一块。

2. 从前馈层到专家 · From FFN to experts

2.1 把一个前馈层拆成专家

slides 先拆一个线性层：$z = Wx + b$ 可以按输出拆成 3 块，$z_i = W_i x + b_i$，拼起来就是原来的 z。再拆一个前馈层：$y = U^\top\sigma(Wx + b) + c$，拆成 3 个专家 $y_i = U_i^\top\sigma(W_i x + b_i) + c_i$。slides 的问题：W = stack(W₁, W₂, W₃)、b = stack(b₁, b₂, b₃)、U = stack(U₁, U₂, U₃) 时，y = y₁ + y₂ + y₃ 需要 c = ____？

答案：c = c₁ + c₂ + c₃。隐藏单元之间互不影响（σ 是逐元素的），所以按隐藏单元把前馈层拆开、各自算完再相加，和原来完全一样；只有输出偏置 c 是”加一次”的，必须在专家之间分掉。代码 01 第 2 部分：c₁ = c、c₂ = c₃ = 0 时和原来相差 3.0e-08；每个专家都带完整的 c 时相差 0.49。

注意 slides 的提醒：MoE 里的每个专家不是一个线性层，而是一个完整的前馈层（有一个隐藏层的神经网络）。

2.2 稠密 MoE

slides：“稠密的混合专家给每个专家一个（非零的）发言权”。它和 §2.1 的关系：把一个宽前馈层拆成 $N_e$ 个专家、再给每个专家一个依赖于输入的权重 $G(x)_i$。所有专家都要算，计算量没有省。

2.3 稀疏门控 MoE

Shazeer et al. (2017) 最早把它用在 RNN 上，现在 Transformer 上很流行：

$-\infty$ 经过 softmax 变成 0：只有 k 个专家的权重非 0，其余专家根本不用算。

手算一个例子（代码 01 第 3 部分，4 个专家、k = 2）：x = [1, −1, 0.5]，$x \cdot W_g$ = [1.05, 1.6, −0.4, −0.6]；top-2 留下专家 0 和 1，softmax([1.05, 1.6]) 得 G(x) = [0.366, 0.634, 0, 0]；$y = 0.366\,E_0(x) + 0.634\,E_1(x)$。MoELayer 算出来的和手算完全一样（相差 0）。

带噪声的 top-k（Shazeer et al.）：

噪声让路由有一点随机性：一个 token 有时会被送给它的”第三选择”，那个专家也就有机会学、有机会变好。

Mixtral 的专家是 SwiGLU（第 4 讲）：$E_i(x) = \mathrm{SwiGLU}(x) = (\mathrm{Swish}(xW + b)\odot(xV + c))\,W_2$，三个矩阵。

初始化（slides）：把 $W_g$ 和 $W_{\text{noise}}$ 初始化为全 0——实际上不给任何信号，只有一点噪声。代码 01 第 4 部分把 10000 个随机 token 送进一个刚初始化的路由器（8 个专家、k = 2）：

不带噪声时，全 0 的路由器给所有专家同样的 logits，topk 总是选下标最小的两个——专家 2–7 一个 token 都没有，一开始也得不到任何梯度。零初始化 + 噪声 = 一开始均匀地、随机地分配，让每个专家都有机会学起来。

3. 激活参数 · Active parameters

slides：在有 MoE 层的 Transformer 里，top-k 的 k 通常选得很小——Mixtral：k = 2、$N_e$ = 8；OLMoE：k = 8、$N_e$ = 64。激活参数 = 路由器为一个 token 选中、真正参与计算的参数。对一个 MoE（粗略地）：

• GPU 显存的需求 ∝ 全部参数（所有专家都要放在显存里，随时可能被选中）；
• FLOPs 计算的需求 ∝ 激活参数。

代码 01 第 5 部分按公开的配置数（每层：注意力 + 路由器 + $N_e$ 个 SwiGLU 专家；嵌入和输出层另算）：

和论文报告的一致。Mixtral 存 46.7B 个参数，每个 token 只用 12.9B（28%）。

Mixtral vs Llama-2（slides 的图，Jiang et al., 2024）：Mixtral 8x7B 在大多数基准上追平或超过 Llama-2 70B，而它每个 token 的激活参数只有 Llama-2 70B 的约 1/5。OLMoE 的超参数（slides）：64 个小专家、每个 token 选 8 个，负载均衡项权重 0.01、路由器 z-loss 权重 0.001——§4 的两个正则项。

为什么 “8x7B” 不是 56B：只有前馈层被复制成 8 份，注意力、嵌入是所有专家共用的；所以全部参数是 46.7B 而不是 8 × 7 = 56B。

4. 负载均衡 · Balancing expert utilization

4.1 问题

slides：不加约束时，门控会集中到少数几个在训练早期就受欢迎的专家上。原因是一个正反馈：某个专家一开始碰巧分到的 token 多一点 → 它被训练得多一点、变得好一点 → 路由器更愿意把 token 给它 → 它分到的更多……其他专家越来越没活干，最后”死掉”：参数白白占着显存。

4.2 两个正则项

slides（OLMoE 的做法，变体很多）：

• $f_i$ = 这一批里被路由到专家 i 的 token 比例（数出来的，不可导）；$P_i$ = 路由器给专家 i 的平均概率（可导）。
• 为什么是 $f_i P_i$：梯度只通过 $P_i$ 传回去，而它的系数是 $f_i$——哪个专家现在越忙，就越用力地压低路由器给它的概率。全部均匀时 $f_i = P_i = 1/N_e$，$L_{\text{LB}} = 1$；如果 $f = P$，由柯西-施瓦茨不等式 $N_e\sum f_i^2 \ge (\sum f_i)^2 = 1$，均匀时最小。
• 路由器 z-loss（Zoph et al., 2022）：$x^{(b)}$ 是第 b 个 token 的路由器 logits。它惩罚大的 logits，让路由器的 softmax 不要太尖、数值上更稳（在 BF16 里，大的 logits 指数一下就溢出，舍入误差也大）。

4.3 我们的实验（代码 02）

2 层的 MoE 语言模型，每层 8 个专家（宽 64）、k = 2，训练 2M 个 token：

（负载 = 每个专家分到的”分配”比例，均匀是 0.125；均衡度 = 负载分布的熵 / log 8。）

• 不带噪声：一开始所有 token 都去专家 0、1（第 0 步最大负载 0.50）；训练中路由器学到了一些区分，第 100 步降到 0.22，但始终不均衡，损失也最差（1.254）。
• 带噪声但不加正则：第 1 层还算均衡，第 2 层有一个专家的负载是 0.000——它死了。
• 加负载均衡：两层都均衡（0.99、0.97）。损失略升（1.126 → 1.145）：负载均衡是在”让每个 token 去最合适的专家”和”让每个专家都有活干”之间折中；换来的是没有死掉的专家和更少的丢弃（下面）。
• z-loss 把路由器的 logsumexp² 从 10.2 压到 0.2，损失又略升（1.173）。在我们的 FP32 小模型里看不出它”让训练更稳”的好处——那是大模型、低精度训练时的问题。

容量：给每个专家设一个容量（§6），超出的分配被丢掉。容量 × 1 时四种设置分别丢掉 24.6% / 23.4% / 10.8% / 6.4% 的分配，容量 × 2 时 0.1% / 2.6% / 0.0% / 0.0%——负载越均衡，丢得越少。

专家的分工（带噪声 + 负载均衡，第 1 层，每种输入 token 的分配去了哪些专家）：

不同种类的 token 明显去了不同的专家：值的分配有 90% 在专家 0 和 7，属性有 48% 在专家 4。但不是一一对应。OLMoE 的分析也发现，专家的分工更像”词的种类”，而不是人们想象的”领域”（数学专家、代码专家……）。

5. 性能 vs 代价 · Performance vs cost

5.1 slides 的结论

• 性能 vs 代价（OLMoE 的图）：同样的激活参数（同样的 FLOPs），MoE 比稠密模型好；MoE 在”性能”和”FLOPs 代价”之间给出了一个很好的取舍。
• 专家数怎么选：早期 LSTM 语言模型上的 MoE 偏爱非常多的专家（成千上万个）；近来 Transformer 语言模型上的工作偏爱相对少的专家（Mixtral 8 个；OLMoE 64 个小专家）。
• 路由模型的规模定律（Clark et al., 2022）：激活参数固定时，损失大致随专家数的对数下降，收益逐渐递减；模型越大，增加专家的收益越小。

5.2 我们的实验（代码 03）

三种模型，每层的前馈部分分别是：稠密 · 小（一个宽 128 的前馈层，激活参数 ≈ MoE）、MoE（8 个宽 64 的专家，k = 2：激活 128 宽、存了 512 宽）、稠密 · 大（一个宽 512 的前馈层，全部参数 ≈ MoE、计算量约 4 倍）。

• 同样的计算量，MoE 更好：激活参数和稠密 · 小差不多，4M token 时损失 0.630 vs 0.714——多出来的、平时不激活的专家参数在”记事实”（“值”的部分 0.583 vs 0.649）。
• 但比不上同样全部参数的稠密模型：稠密 · 大 0.553。每个 token 只用 1/4 的前馈参数，和”全部都用”相比还是有代价。
• 数据越多，参数越重要：稠密 · 小比稠密 · 大差 14%（2M）→ 29%（4M）。数据少时，没见过的事实再多参数也记不住（第 15 讲），参数的好处要等数据跟上才显现。
• 专家数 4 / 8 / 16（激活参数不变，2M token）：超额损失 1.160 / 1.145 / 1.118，全部参数 50,384 / 84,448 / 152,576。专家越多越好；在我们这么小的范围里还没看到”收益递减”（每翻一倍降 0.015、0.027）。
• ⏱ 时间：在 CPU 上，稠密 · 大每步 15 ms，MoE 17 ms——这么小的矩阵，路由和按专家分组的开销比省下的乘法还多。FLOPs 少不等于时间短：MoE 要在矩阵足够大、实现足够好（专门的分组矩阵乘法内核）时，才真正省时间。

一个提醒：我们的任务几乎全是”记事实”。Jelassi et al. (2024) 发现，MoE 对记忆类任务的帮助远大于推理类任务——多出来的专家参数增加了”知识容量”，但不增加”每个 token 能做多少步计算”。

6. 专家并行 · Expert parallelism

6.1 容量

slides：每个专家可以放在不同的设备（GPU）上；每个 token 被路由给 K 个专家。如果太多 token 被路由给同一个专家，就有负载均衡的问题。每块设备一次只能处理固定数量的 token——容量（capacity）：

超出容量的 token 直接经过残差连接传到下一层（Switch Transformer；slides）。容量因子太小，容易超额分配（丢 token）；太大，一些 GPU 就会闲着（slides）。

slides 的两个例子（代码 04 第 1 部分；3 块设备，每块容量 3 个 token；按 GShard 的规则先放所有 token 的第一选择，再放第二选择）：

(a) 里设备 2 只用了 1 个容量：“浪费的容量”。(b) 里 token 3 也被分给了设备 1，“但它已经满了！“——这份分配被丢掉，token 3 只剩设备 2 上的那个专家。

真实路由（代码 04 第 2 部分：两个训练好的 MoE，8 个专家放在 8 块设备上，一批 32,256 个 token）。以最不均衡的第 2 层为例：

没有丢弃时，空闲的比例只由容量因子决定（1 − 1 / 容量因子：× 2 时一半的容量空着）。负载均衡让我们能用更小的容量因子、而几乎不丢 token，空闲也就更少——这是它在工程上最直接的价值。

6.2 通信

专家并行的每个 MoE 层要做两次 all-to-all（P4）：dispatch（把每个 token 的隐藏向量发给放着它的专家的那块 GPU）和 combine（把专家的输出发回来）。反向传播再来两次。代码 04 第 3 部分：一个 Mixtral 规模的层（d = 4096、k = 2、8 个专家在 8 块 GPU 上、每块 GPU 8192 个 token），每块 GPU 每层发出约 235 MB：

专家并行最好放在一个节点里；跨节点时通信就占了相当一部分时间，专家越多、越分散，越难。这也是近来的模型偏爱”不太多”的专家的工程原因之一。

7. 小结 · Summary

1. 前馈层占 Transformer 大部分参数（GPT-3：66%）。把它拆成 $N_e$ 个专家、用路由器只选 k 个，就能让参数很多、每个 token 的计算量不变。
2. 稀疏门控：$G(x) = \mathrm{softmax}(\mathrm{topk}(xW_g + \text{噪声}, k))$；路由器全 0 初始化 + 噪声，让一开始的分配均匀而随机。
3. 显存 ∝ 全部参数，FLOPs ∝ 激活参数（Mixtral 46.7B / 12.9B）。
4. 路由器会偏心、专家会死；负载均衡项 $N_e\sum f_i P_i$ 让分配均匀（代价是损失略升），z-loss 压住路由器的 logits。
5. 同样的计算量，MoE 比稠密模型更会”记事实”，但比不上同样全部参数的稠密模型；专家并行要处理容量、丢弃和 all-to-all 通信。

8. 常见误区 · Common pitfalls

• “8x7B 就是 56B。” 只有前馈层被复制了；注意力、嵌入是共用的：46.7B。
• “MoE 的每个专家是一个线性层。” 是一个完整的前馈网络（slides 特意强调）。
• “MoE 省显存。” 省的是计算；显存要按全部参数准备。
• “负载均衡让模型更准。” 它主要让专家都被用上、少丢 token、GPU 利用率高；在我们的实验里损失反而略升。
• “FLOPs 少就一定快。” 小矩阵上，路由、分组、通信的开销可能比省下的乘法还多（我们的 CPU：MoE 17 ms vs 4 倍计算量的稠密模型 15 ms）。
• “专家会按领域分工（数学专家、代码专家）。” 实际观察到的分工更像”词的种类”（我们的：值、属性各自集中在少数几个专家）。
• “比较 MoE 和稠密模型时，其他都不用管。” 我们的第一版里，稠密层用 GELU、专家用 ReLU，就让 MoE 显得比 4 倍计算量的稠密模型还好——对比实验里，除了要比较的那一处，其余必须完全一样。

9. 练习 · Exercises

练习 1（数参数）。 一个 Transformer：d = 4096、32 层、注意力用 GQA（32 个查询头、8 个键值头）、前馈层是 SwiGLU（$d_{\text{ff}}$ = 14336）、词表 32000、嵌入和输出层不共享。(a) 稠密版本有多少参数？(b) 把每层的前馈层换成 8 个同样大小的专家、k = 2，全部参数和激活参数各是多少？（这就是 Mixtral。）

每层注意力：$2d^2 + 2d\cdot(8\cdot128) = 2\cdot4096^2 + 2\cdot4096\cdot1024 ≈ 41.9$M；每个 SwiGLU：$3\cdot4096\cdot14336 ≈ 176.2$M；嵌入 + 输出层：$2\cdot32000\cdot4096 ≈ 0.26$B。(a) 32 × (41.9 + 176.2)M + 0.26B ≈ 7.24B（Mistral 7B）。(b) 全部：32 × (41.9M + 8 × 176.2M + 路由器 4096 × 8) + 0.26B ≈ 46.7B；激活：32 × (41.9M + 2 × 176.2M) + 0.26B ≈ 12.9B。

练习 2（负载均衡项的梯度）。 $L_{\text{LB}} = N_e\sum_i f_i P_i$，$P_i = \frac{1}{M}\sum_{t}p_i(x_t)$，$p(x_t) = \mathrm{softmax}(x_t W_g)$。(a) 对 $p_i(x_t)$ 的梯度是多少？(b) 用它说明：一个过载的专家（$f_i$ 大），路由器会被推着”少选它”。

(a) $f_i$ 是数出来的、当作常数，$\partial L_{\text{LB}}/\partial p_i(x_t) = N_e f_i / M$。(b) 梯度下降让 $p_i(x_t)$ 朝 $-N_e f_i/M$ 的方向变：$f_i$ 越大，压得越用力；通过 softmax，压低 $p_i$ 也就抬高了其他专家的概率。所以过载的专家在下一步被选中的概率变小，负载被推向均匀。

练习 3（容量与丢弃）。 8 个专家、k = 2、一批 1024 个 token、容量因子 1.25。(a) 每个专家的容量是多少？(b) 如果路由完全均匀，会丢 token 吗？(c) 如果一个专家分到了全部分配的 25%，会丢掉多少份分配？

(a) 1.25 × 2 × 1024 / 8 = 320。(b) 均匀时每个专家分到 2 × 1024 / 8 = 256 份，小于 320，不丢。(c) 25% × 2048 = 512 份，超出 320 的 192 份被丢掉（占全部分配的 9.4%）。

练习 4（all-to-all 的量）。 推导 §6.2 的 235 MB：每块 GPU 8192 个 token、d = 4096、BF16、k = 2、8 块 GPU。如果改成 64 个专家在 64 块 GPU 上（每块一个专家）、其余不变，每块 GPU 发出的字节数怎样变？

dispatch：每个 token 的 d 维向量发给 k 个专家，$8192 \times 2 \times 4096 \times 2$ 字节 ≈ 134 MB，其中 (N−1)/N = 7/8 要发到别的 GPU，≈ 117 MB；combine 同样多，合计 ≈ 235 MB。64 块 GPU 时 (N−1)/N = 63/64，几乎全部要发出去，≈ 264 MB；而且对手从 7 个变成 63 个，每次发送更小、α（延迟）项更多，跨节点的比例也更大——专家越多，通信越难。

练习 5（为什么不能直接对 f 求导）。 $f_i$ 是”被选中的次数 / 总数”。为什么不能直接最小化 $\sum_i f_i^2$？slides 的 $f_i P_i$ 是怎样绕开这个问题的？

$f_i$ 由 argmax / top-k 决定，是分段常数：对路由器的参数几乎处处梯度为 0，没法用梯度下降。$f_i P_i$ 里，$f_i$ 只作为一个”权重”（谁忙谁的权重大），梯度从可导的 $P_i$（softmax 概率的平均）传回路由器——用一个可导的代理去推动不可导的负载。

练习 6（MoE 什么时候不划算）。 列出三个”MoE 比同样激活参数的稠密模型好不了多少、甚至更差”的情形。

(1) 任务靠”多算几步”而不是”多记东西”（推理、算术）：多出来的专家参数不增加每个 token 的计算深度（Jelassi et al., 2024）。(2) 数据太少：专家多了，每个专家分到的数据更少，训练不充分；我们的 2M token 时差距只有 4M 时的一半。(3) 推理时批很小、或显存紧张：所有专家都要放在显存里，单个请求时只用到很少几个，显存和带宽利用率很低；跨设备时 all-to-all 的通信更是额外开销。

10. 延伸阅读 · Further reading

• 起源与早期工作：Jacobs et al. (1991), Adaptive Mixtures of Local Experts；Shazeer et al. (2017), Sparsely-Gated MoE；Lepikhin et al. (2020), GShard；Fedus et al. (2021), Switch Transformers；Zoph et al. (2022), ST-MoE（路由器 z-loss）。
• 现代 MoE：Jiang et al. (2024), Mixtral of Experts；Muennighoff et al. (2024), OLMoE；Dai et al. (2024), DeepSeekMoE（细粒度专家、共享专家）；Cai et al. (2024), A Survey on Mixture of Experts。
• 规模定律与分析：Clark et al. (2022), Unified Scaling Laws for Routed Language Models；Jelassi et al. (2024), Mixture of Parrots: Experts improve memorization more than reasoning。

11. 下一讲预告 · Next lecture

第 17 讲：分布式训练。专家并行只是其中一种。一个 4000 亿参数的模型光是 BF16 权重就要 800 GB，一块 GPU 放不下；训练时还要存梯度和 Adam 的状态（每个参数 16 字节）。我们会用”一个列表模拟多块 GPU”的方式，真实地执行数据并行、张量并行（Megatron 的列切分 / 行切分）、流水线并行（GPipe、1F1B、它们的气泡）和 ZeRO 的三个阶段，并核对每一种都和单卡训练的结果一模一样。