第 2 周 · Day 2：混合专家（MoE）

时长 7 小时：§1 思想与形式（1.5h）· §2 参数、计算与显存（1h）· §3 训练难题：均衡、容量、稳定性（2h）· §4 新设计（0.5h）· §5 系统层面（0.5h）· §A–§C（1.5h）

对应正课 Lecture 3 的混合专家部分。学习计划要求能说清”MoE 的总参数量和激活参数量的区别”、负载均衡损失，以及”Mixtral 8x7B 为什么不等于 56B 的推理成本”。

配套代码 w2day2_checks.py（纯 numpy）：top-2 MoE 前向、Mixtral 的精确参数计数、负载均衡损失及其梯度方向、容量因子与丢 token、解码时被读到的专家比例、细粒度专家的组合数、z-loss、无辅助损失的偏置均衡

§B 参考实现（code/ 目录）：w2d2_moe.py（MoE 层，接进 Week 1 的 mini-GPT；稠密 vs MoE 对照、负载监控、去掉均衡损失的消融）。

今天的目标：能画出一个 MoE 层的数据流；能精确数出 Mixtral 的 46.7B / 12.9B；能从 softmax 雅可比推出负载均衡损失的梯度为什么能把负载推均匀；能说清 MoE 在推理时省了什么、没省什么。

§0 今天的主线

Day 1 说：参数越多，损失越低。但在稠密模型里，参数翻倍，每个 token 的计算也翻倍（Week 1 Day 2 §5.2：推理每 token $2N$ ）。

能不能让参数多、而每个 token 的计算不变？ MoE 的回答：把 FFN 复制成很多份”专家”，每个 token 只激活其中少数几个。参数量和计算量从此解耦。

今天要用到的旧结论：FFN 是键值记忆，激活是稀疏的（Week 1 Day 2 §1.3）；SwiGLU（Week 1 Day 2 §1.4）；softmax 雅可比 $\mathrm{diag}(p)-pp^\top$ （Week 0 Day 1 §5.2.1）；roofline 与解码的访存瓶颈（Week 1 Day 4 §1.5）；Day 1 的缩放律。

§1 思想与形式

1.1 动机：条件计算

Week 1 Day 2 §1.3 把 FFN 看成一张键值记忆表： $d_{ff}$ 个”槽”，每个 token 只激活其中一部分（随机输入下约一半，训练后更稀疏）。既然每个 token 只用到一部分知识，为什么要为每个 token 计算全部？

直觉：一家大医院有很多科室，病人只去相关的一两个，不会每个科室都走一遍。医院的”总能力”随科室数增长，每个病人花的时间却不变。MoE 就是把 FFN 拆成很多科室，再加一个分诊台（路由器）。

1.2 一段历史

Jacobs et al.（1991）：专家混合的原始想法。
Shazeer et al.（2017）：稀疏门控 MoE，把它用进 LSTM 语言模型，参数做到 1370 亿。
GShard（2020）、Switch Transformer（Fedus et al. 2021）：放进 Transformer 的 FFN 位置，Switch 用 top-1 路由，参数做到 1.6 万亿。
Mixtral 8x7B（2023）：开放权重的 MoE，让它成为主流选择。
DeepSeekMoE / DeepSeek-V2 / V3（2024）：细粒度专家、共享专家、无辅助损失均衡（§4）。

1.3 形式化

一个 MoE 层替换 Transformer 块里的 FFN（注意力不变）。有 $E$ 个专家 $\{f_1,\dots,f_E\}$ （每个都是一个 FFN，现代模型用 SwiGLU）和一个路由器 $W_r\in\mathbb R^{d\times E}$ ：

$z=W_r^\top x\in\mathbb R^E\qquad(\text{每个专家一个打分})$

$\mathcal T=\mathrm{TopK}(z,k)\qquad(\text{选出打分最高的 }k\text{ 个})$

$g_i=\frac{e^{z_i}}{\sum_{j\in\mathcal T}e^{z_j}}\ \ (i\in\mathcal T),\qquad y=\sum_{i\in\mathcal T}g_i\,f_i(x)$

门控权重只在选中的 $k$ 个上做 softmax（Mixtral 的做法；Switch 的 $k=1$ 直接用完整 softmax 的概率）。每个 token 只算 $k$ 个专家。

w2day2_checks.py 第 1 组（8 个专家，top-2）：6 个 token 共调用专家 12 次（ $=Tk$ ），而不是 48 次。与”所有专家都算、按完整 softmax 加权”的稠密混合相比，输出相对差 0.30：top-k 是一个真正的近似，不是免费的。

为什么选择本身不可导也能训？ $\mathrm{TopK}$ 是离散的，梯度过不去。但门控权重 $g_i$ 对选中专家的 $z_i$ 是可导的：被选中的专家如果表现好，梯度会让它的 $g_i$ 变大、它的 $z_i$ 变大，下次更可能被选中。路由器靠这条路学习。

§2 参数、计算与显存

2.1 Mixtral 8x7B：精确计数

Mixtral 的结构：32 层， $d=4096$ ，GQA（32 个 query 头、8 个 KV 头，每头 128 维），每层 8 个 SwiGLU 专家（ $d_{ff}=14336$ ）、top-2，词表 32000，输入输出嵌入不绑定，无 bias。w2day2_checks.py 第 2 组：

项目	每层
注意力（Q、K、V、O）	41.9M
每个专家（ $3\cdot d\cdot d_{ff}$ ）	176.2M
路由器（ $d\times E$ ）	32,768

$\text{总参数}=46{,}702{,}792{,}704\approx46.7\text{B},\qquad\text{激活参数}=12{,}879{,}925{,}248\approx12.9\text{B}$

与 Mistral 公布的数字完全一致。

2.2 为什么 “8x7B” 不是 56B

只有 FFN 复制了 8 份，注意力、嵌入、归一化都是共享的。专家占总参数的 96.6%，但”一个 7B 模型 × 8”会把共享部分也数 8 遍。激活参数是 12.9B 而不是 $2\times7=14$ B，也是同一个原因。

三个数要分清：

	由什么决定	Mixtral
计算（每 token FLOPs）	激活参数	约 $2\times12.9\text{B}=26$ GFLOPs，接近 13B 稠密模型
显存（装下模型）	总参数	bf16 约 87 GiB，接近 47B 稠密模型
质量	介于两者之间	Mistral 报告在多数基准上与 Llama 2 70B 相当或更好

MoE 的买卖：用显存换计算。显存不省：所有专家都得装着。

2.3 解码时的访存：batch 大小决定省不省

Week 1 Day 4 §1.5 说过：解码是访存密集的，时间主要花在”把权重从显存读出来”。MoE 在解码时要读多少权重？

取决于 batch 里的 token 选中了多少个不同的专家。 第 5 组（top-2 of 8，路由均匀时，至少被一个 token 选中的专家比例）：

batch	1	2	4	8	16	32
被读到的专家	25.0%	43.8%	68.4%	90.0%	99.0%	100.0%

batch 为 1 时只读 2/8 的专家，访存按激活参数算，MoE 很划算；batch 大了几乎所有专家都要读，访存按总参数算。 这时 MoE 省的是计算，不是访存。这是 MoE 推理部署比稠密模型复杂的一个根本原因。

§3 训练的难题

3.1 路由坍缩：强者愈强

训练初期，某个专家偶然表现好一点 → 路由器更多地选它 → 它得到更多训练、变得更好 → 路由器更偏爱它。正反馈。 最终少数几个专家承担几乎全部 token，其余专家从不被训练。MoE 退化成一个小的稠密模型，还浪费了大量显存。

3.2 负载均衡损失：推导它为什么有效

Switch Transformer 加了一个辅助损失（ $E$ 个专家，一个 batch 共 $T$ 个 token）：

$\mathcal L_{\text{aux}}=\alpha\cdot E\cdot\sum_{i=1}^Ef_i\,P_i,\qquad f_i=\frac1T\sum_t\mathbb 1[\text{token }t\text{ 选了 }i],\qquad P_i=\frac1T\sum_tp_{t,i}$

$f_i$ 是专家 $i$ 实际分到的 token 比例（离散计数，不可导）； $P_i$ 是路由器给专家 $i$ 的平均概率（可导）。

为什么均匀时最小？ 当 $f_i=P_i=1/E$ 时， $\mathcal L_{\text{aux}}=\alpha E\cdot E\cdot\frac1{E^2}=\alpha$ 。w2day2_checks.py 第 3 组：均匀路由时 0.0100，轻度偏斜时 0.0131，几乎坍缩到一个专家时 0.0766。

梯度推导（这是理解它的关键）。 $f$ 当常数，对 token $t$ 的路由 logit $z_{t,i}$ 求导。 $P_j$ 里只有 $p_{t,j}$ 依赖 $z_t$ ，用 softmax 雅可比 $\partial p_{t,j}/\partial z_{t,i}=p_{t,j}(\delta_{ij}-p_{t,i})$ （Week 0 Day 1 §5.2.1）：

$\frac{\partial\mathcal L_{\text{aux}}}{\partial z_{t,i}}=\frac{\alpha E}T\sum_jf_j\,p_{t,j}(\delta_{ij}-p_{t,i})=\frac{\alpha E}T\,p_{t,i}\Big(f_i-\sum_jp_{t,j}f_j\Big)$

$\delta_{ij}$ 把求和塌缩到 $j=i$ ，和 Week 1 Day 1 §2.3 推注意力反向时的那一步一模一样。

读这个式子：括号里是”专家 $i$ 的负载”减去”按当前路由概率加权的平均负载”。负载高于平均的专家，梯度为正，梯度下降把它的 logit 压低；低于平均的，梯度为负，logit 被抬高。第 3 组的数值梯度：负载高于 $1/E$ 的专家 5、6、7 梯度全为正，低于的专家 0–4 全为负，最忙的专家 7 梯度最大。

$\alpha$ 取多少？ Switch 用 0.01。太大会让路由器只顾均衡、不顾把 token 送给最合适的专家；太小压不住坍缩。

3.3 容量因子：给每个专家设上限

为了高效并行，每个专家一次处理的 token 数最好是固定的。容量 $=\text{CF}\times T/E$ ，CF 是容量因子。超出容量的 token 被丢弃：不经过任何专家，只走残差直通。

第 4 组（ $T=4096$ ，8 个专家）：

负载分布	CF=1.0	CF=1.25	CF=2.0
完全均匀	0.0%	0.0%	0.0%
轻度不均（最忙 1.5 倍平均）	10.0%	3.1%	0.0%
严重不均（Zipf）	30.2%	23.9%	11.8%

取舍：CF 小，丢 token；CF 大，每个专家要预留更多计算和显存，大部分时候浪费。Switch 用 1.0–1.25。负载均衡损失与容量因子是一对：前者让负载尽量均匀，后者兜底。后来的”无丢弃”实现（如 MegaBlocks 用块稀疏矩阵乘）取消了容量上限，代价是实现复杂。

3.4 训练不稳定与 z-loss

MoE 比稠密模型更容易出现训练不稳定。一个来源是路由器 logit 越来越大：softmax 越来越尖、越来越接近不可导的硬选择，bf16 下 exp 也更容易溢出。ST-MoE（Zoph et al. 2022）加了路由器 z-loss：

$\mathcal L_z=\frac1T\sum_t\Big(\log\sum_ie^{z_{t,i}}\Big)^2$

它惩罚 logit 的整体量级。第 7 组：logit 标准差 1、5、20 时，z-loss 分别是 6.2、59.6、988.9，路由概率最大值均值从 0.355 升到 0.949。这和 $1/\sqrt{d_k}$ （Week 0 Day 4 §5.5）、YaRN 的温度（Week 1 Day 5 §3.4）是同一类手段：控制 softmax 的尺度。

§4 新设计

4.1 细粒度专家与共享专家

DeepSeekMoE 的两个改动：

细粒度：把每个专家切成 $m$ 份更小的专家，激活数也乘 $m$ 。计算量不变，但可能的”专家组合”暴涨。第 6 组：8 选 2 只有 28 种组合；切成 1/8 大小、64 选 16，有约 $4.9\times10^{14}$ 种。每个 token 能拼出更精细的知识组合，专家也更专一。
共享专家：几个所有 token 都走的专家，承担通用知识（语法、常见模式），让路由专家不必各自重复学这些。

DeepSeek-V3 把这套设计推到 671B 总参数、37B 激活参数。

4.2 无辅助损失的负载均衡

辅助损失的梯度会干扰主任务（§3.2 里 $\alpha$ 的两难）。DeepSeek-V3 的做法：给每个专家一个偏置 $b_i$ ，只在”选谁”时加上，不进入门控权重：

$\mathcal T=\mathrm{TopK}(z+b,\ k),\qquad g_i\ \text{仍由原始的}\ z\ \text{计算}$

每步根据负载调偏置：过载的专家 $b_i$ 下调，欠载的上调。它不是梯度下降，所以不会把梯度注入主损失。第 8 组模拟（随机初始化的路由器本身就不均衡）：最忙专家的负载从平均的 1.149 倍降到 1.008 倍（10 步 1.055，50 步后稳定），最闲的从 0.853 升到 0.990。DeepSeek-V3 仍保留了一个权重极小的序列级均衡损失作为补充。

§5 系统层面：专家并行

专家太多，一张卡放不下，于是把不同专家放在不同的卡上（专家并行）。一个 MoE 层的前向变成：

1. 每张卡算自己那批 token 的路由
2. all-to-all：把每个 token 发到它选中的专家所在的卡      ← 通信
3. 各卡上的专家处理收到的 token
4. all-to-all：把结果发回 token 原来所在的卡              ← 通信
5. 按门控权重加权求和

两次 all-to-all 是 MoE 的主要系统开销，每次都要搬动约 $T\times k\times d$ 个数。负载不均时，最忙的那张卡决定整层的速度（木桶效应），所以负载均衡不只是”让专家都学到东西”，也是系统吞吐问题。这些是 Lecture 4 硬件优化与分布式训练（Day 4 §4）的内容。

§A MoE：API 用法与边界

A.1 HuggingFace 的 MoE 配置

cfg = AutoConfig.from_pretrained("mistralai/Mixtral-8x7B-v0.1")
cfg.num_local_experts      # E = 8
cfg.num_experts_per_tok    # k = 2
cfg.router_aux_loss_coef   # 均衡损失的系数 α
out = model(input_ids, labels=labels, output_router_logits=True)
out.aux_loss               # 只有 output_router_logits=True 时才会计算并加进 out.loss

微调 MoE 时，如果不传 output_router_logits=True，均衡损失不会生效，路由可能在微调中逐渐坍缩。
不同模型族的字段名不同（num_experts、moe_intermediate_size、n_shared_experts 等），先读 config。
加载要按总参数准备显存（§2.2）。放不下时用 device_map="auto" 把部分层卸载到 CPU，速度会显著下降。

A.2 自己实现时的几个坑

最简单的实现：对每个专家用布尔掩码挑出属于它的 token，算完再用 index_add_ 写回。正确但慢（逐专家的 Python 循环）；够学习用。
高效的实现：按专家编号对 token 排序，变成若干段连续的矩阵乘（grouped GEMM）。
top-k 之后要重新归一化门控权重（在选中的 $k$ 个上做 softmax），否则输出尺度会随路由分布漂移。
均衡损失只在训练时加；统计 $f_i$ 时，top- $k$ 路由要把 $k$ 个选择都计入。
torch.topk 遇到并列时的选择不保证稳定；初始化路由器时别用全零，否则所有 token 都去同样的专家。

§B 实操

跑通 w2day2_checks.py，8 组结论都看懂。第 2 组自己按结构重数一遍；第 3 组对照 §3.2 的梯度推导。
推导（1h）：(a) 推出负载均衡损失对路由 logit 的梯度；(b) 证明均匀时 $\mathcal L_{\text{aux}}=\alpha$ ；(c) 数出 Mixtral 的总参数与激活参数。
实现 MoE 层（2h）：top-k 路由、SwiGLU 专家、门控重新归一化、均衡损失、z-loss。验证 $E=1,k=1$ 时与稠密 SwiGLU FFN 完全一致。
接进 mini-GPT 做对照（2h，边跑边读）：
- 稠密 vs MoE（激活计算相同：每个专家的 $d_{ff}$ 取稠密的一半、top-2）
- 训练中每 250 步打印每个专家的负载分布
- 消融：去掉均衡损失，看负载是否坍缩，验证 loss 是否变差
读论文：Switch Transformer 的第 2 节（路由、容量、均衡损失），对照今天 §3。

参考实现：code/w2d2_moe.py。

§C 自测

动机与直觉层

MoE 解决的是什么问题？“参数与计算解耦”是什么意思？
为什么 top-k 选择不可导，路由器还能学？
什么是路由坍缩？为什么会发生？
容量因子为什么存在？调大调小各有什么代价？
为什么说 MoE “用显存换计算”？batch 大小如何影响它在解码时的访存优势？
细粒度专家和共享专家各解决什么问题？

推导层

数出 Mixtral 8x7B 的总参数和激活参数。为什么不是 56B 和 14B？
写出 Switch 的负载均衡损失，推出它对路由 logit 的梯度，解释梯度的符号。
为什么 $f_i$ 不可导也没关系？
z-loss 惩罚的是什么？它和 $\sqrt{d_k}$ 缩放有什么共同点？
top-2 of 8 的路由、batch 为 8 时，期望有多少比例的专家被读到？

工程层

微调 HF 的 MoE 模型时，怎样确保均衡损失生效？
专家并行的两次 all-to-all 分别在做什么？负载不均为什么会拖慢整层？
无辅助损失的均衡为什么不影响主损失的梯度？

出关标准

能画出 MoE 层的数据流，写出路由与门控的公式
能精确数出 Mixtral 的 46.7B / 12.9B，说清计算、显存、质量各由什么决定
能推出负载均衡损失的梯度，解释它为什么能把负载推均匀
能说出容量因子与丢 token 的取舍
能用 roofline 解释 MoE 解码时 batch 大小的影响
能说出 z-loss、细粒度专家、共享专家、无辅助损失均衡各自的动机
实现 MoE 层，接进 mini-GPT，完成稠密 vs MoE 与去掉均衡损失的对照