# Lecture 4 · 混合专家模型 MoE
> **CS336: Language Modeling from Scratch** · Stanford · Spring 2025
> 📅 Apr 10 · 💻 [不可执行讲义（PDF）](https://github.com/stanford-cs336/spring2025-lectures/blob/main/nonexecutable/2025%20Lecture%204%20-%20MoEs.pdf) · 讲者 Tatsunori Hashimoto
---
## 承上启下
**上一讲（[[Lecture 3 · 架构与超参数]]）**：
- 现代 Transformer 已收敛到一套稳定配方：Pre-Norm + RMSNorm + SwiGLU + RoPE + 无 bias，这些是「稠密模型内部」的共识。
- 超参数大多落在平坦盆地里：FFN 倍数（ReLU/GeLU 用 $4\times$，GLU 系用 $\frac{8}{3}\times$）、head 维度 64–128、深宽比约 100–200、词表 30k–250k。
- 训练稳定性技巧（z-loss、QK-norm、soft-capping）都在管 softmax 的输入尺度——这个思路本讲还会在 router 上再遇到一次。
- GQA/MQA 通过共享 K/V head 压缩 KV cache，是「推理成本反过来塑造架构」的例子。
**本讲（混合专家模型 MoE）**：
- 上一讲优化的是**单个稠密模型**的架构细节——无论怎么改，每个 token 都要走遍所有参数。
- 本讲换一个思路：**同样的计算预算，能不能装下更多参数？**
- 答案是 **MoE（Mixture of Experts，混合专家模型）**：把 FFN 层拆成多个「专家」，每个 token 只激活其中几个——每 token 计算量不变，总参数量成倍增长。
> [!note] 核心权衡
> 稠密模型是「小参数 × 全激活」，MoE 是「大参数 × 稀疏激活」——在相同 FLOPs 下，MoE 的总参数量可以是稠密模型的 4–8 倍甚至更多，性能通常更好；代价是**路由设计**、**负载均衡**与**All-to-All 通信**这三座新大山。
---
## 1. 为什么需要 MoE？
### 1.1 缩放定律的困境
**Chinchilla 缩放定律**（Hoffmann et al., 2022，详见 [[Lecture 9 · 缩放定律 I：基础]]）告诉我们：想训练算力最优的模型，参数量 $N$ 与数据量 $D$ 要等比例同增（约 20 tokens/参数）。而 [[Lecture 2 · PyTorch 与资源核算]] 的训练计算量公式是
$$C \approx 6ND$$
$N$、$D$ 各翻一倍，FLOPs 就涨四倍——按这条路线继续堆稠密模型，成本增长快于收益。MoE 的问题意识由此而来：**能不能只增加 $N$（容量），不等比例增加每个 token 实际消耗的计算？**
```mermaid
graph TD
A["想要更好的模型"] --> B["缩放定律：需要更大 N 和更多 D"]
B --> C["计算成本 ∝ N·D"]
C --> D["训练成本爆炸"]
D --> E{"有没有办法
只增加 N
但不增加 FLOPs？"}
E -->|稀疏激活| F["MoE"]
style A fill:#e1f5fe,stroke:#0277bd
style D fill:#ffcdd2,stroke:#c62828
style F fill:#c8e6c9,stroke:#2e7d32
classDef default fill:#fff,stroke:#666
```
**现实案例**（推理 FLOPs 近似正比于激活参数量，以 Llama 3.1 405B 为基准 1×）：
| 模型 | 类型 | 总参数 | 激活参数 | FLOPs（推理时） | 性能 |
| :--- | :--- | ---: | ---: | :--- | :--- |
| Llama 3.1 405B | 稠密 | 405B | 405B | 基准 1× | 基准 |
| DeepSeek-V3 | MoE | 671B | 37B | ≈ 0.09× | **更优** |
| Mixtral 8×7B | MoE | 47B | 13B | ≈ 0.03× | 接近 Llama 2 70B |
> [!tip] 直觉
> DeepSeek-V3 的总参数量是 Llama 3.1 405B 的 1.66 倍，但每个 token 只激活 37B 参数——推理 FLOPs 只有稠密 405B 的约 9%，性能却更好。要习惯 MoE 时代的「双参数量」记法：**总参数决定容量与显存占用，激活参数决定每 token 的 FLOPs**。MoE 省的是「算」，不省「存」——671B 参数无论激活多少，都得放在显存/内存里。
### 1.2 MoE 的核心思想
把 Transformer 的 **FFN 层**拆成 $N$ 个「专家（expert）」，每次只用其中 $K$ 个：
```mermaid
flowchart TB
subgraph Dense["稠密 FFN（传统）"]
D1["输入 x"] --> D2["一个大 FFN
所有参数都激活"]
D2 --> D3["输出"]
end
subgraph MoE["MoE FFN（稀疏）"]
M1["输入 x"] --> Router["路由函数 g(x)
选出 top-K 专家"]
Router --> E1["专家 1"]
Router --> E2["专家 2"]
Router -.不激活.-> E3["专家 3"]
Router -.不激活.-> EN["专家 N"]
E1 --> Combine["加权组合"]
E2 --> Combine
Combine --> M2["输出"]
end
style D2 fill:#ffcdd2,stroke:#c62828
style Router fill:#fff9c4,stroke:#f57f17
style E1 fill:#c8e6c9,stroke:#2e7d32
style E2 fill:#c8e6c9,stroke:#2e7d32
style E3 fill:#e0e0e0,stroke:#757575
style EN fill:#e0e0e0,stroke:#757575
style Combine fill:#e1f5fe,stroke:#0277bd
```
**前向传播对比**——稠密 FFN 激活全部参数，MoE 只对路由选中的专家求值：
$$\text{稠密：}\; y = \mathrm{FFN}(x) \qquad\qquad \text{MoE：}\; y = \sum_{i \in \mathrm{TopK}(x)} g_i(x)\, E_i(x)$$
符号说明：$E_i(\cdot)$ 是第 $i$ 个专家（一个独立的小 FFN），$g_i(x) \ge 0$ 是路由函数给专家 $i$ 的门控权重（未选中的专家 $g_i = 0$，对应的 $E_i(x)$ 根本不计算），$\mathrm{TopK}(x)$ 是被选中的 $K$ 个专家的下标集合。
> [!warning] 关键权衡
> - **优势**：相同 FLOPs 下参数量可达稠密模型的数倍，同算力预算下 loss 更低（Switch Transformer、OLMoE 等系统实验反复验证）。
> - **代价**：需要设计路由函数、解决负载均衡、处理通信开销（并行训练时专家分布在不同设备上）；训练稳定性与微调也更棘手。
---
## 2. MoE 架构细节
### 2.1 专家的粒度
**问题**：应该在 Transformer 的哪个部件用 MoE？每个专家应该多大、多少个？
**业界共识**（2025）：
```mermaid
graph LR
A["Transformer Block"] --> B["Multi-Head Attention
（保持稠密）"]
B --> C["Add & Norm"]
C --> D["MoE FFN Layer
（替换原 FFN）"]
D --> E["Add & Norm"]
style B fill:#e0e0e0,stroke:#757575
style D fill:#c8e6c9,stroke:#2e7d32
```
| 设计选择 | 说明 |
| :--- | :--- |
| **只替换 FFN** | FFN 占 Transformer 参数的约 2/3，是稀疏化收益最大的部位；attention 的 MoE 变体（如 MoA）存在但不主流，且 attention 稀疏化对稳定性影响更大 |
| **哪些层用 MoE** | 早期 GShard 隔层替换；现代模型（Mixtral、DeepSeek 系）基本每层都用，但 DeepSeek-V3 保留**前 3 层为稠密 FFN**（浅层功能通用，稀疏化收益小且伤稳定性） |
| **专家大小** | 每个专家 = 一个完整的小 SwiGLU FFN（gate_proj + up_proj + down_proj） |
**细粒度专家（fine-grained experts）**是 DeepSeekMoE 引入的重要维度：把每个专家切小为原来的 $\frac{1}{m}$，同时把专家数与激活数都乘 $m$（$N \to mN$，$K \to mK$）。激活参数量不变，但可选的专家**组合数**从 $\binom{N}{K}$ 暴增到 $\binom{mN}{mK}$，专家分工可以更细、冗余更少。
> [!example] 组合数演算
> $N=16, K=2$ 时组合数 $\binom{16}{2} = 120$；切细 4 倍（$m=4$）后 $N=64, K=8$，组合数 $\binom{64}{8} \approx 4.4 \times 10^9$。同样的激活计算量，表达「专家组合」的自由度完全不在一个量级——这是 DeepSeek 系列坚持用大量小专家（V3 达 256 个）的核心理由。
**代码示例**（简化的 MoE FFN）：
```python
import torch
import torch.nn as nn
class ExpertFFN(nn.Module):
"""单个专家 = 一个标准的 SwiGLU FFN"""
def __init__(self, d_model: int, d_ff: int):
super().__init__()
self.up_proj = nn.Linear(d_model, d_ff, bias=False)
self.gate_proj = nn.Linear(d_model, d_ff, bias=False)
self.down_proj = nn.Linear(d_ff, d_model, bias=False)
def forward(self, x):
# SwiGLU: (SiLU(W_gate(x)) ⊙ W_up(x)) W_down —— 激活函数作用在门控分支上
return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))
class MoELayer(nn.Module):
"""MoE FFN 层：N 个专家 + 路由函数"""
def __init__(self, d_model: int, d_ff: int, num_experts: int, top_k: int):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
# N 个独立的专家
self.experts = nn.ModuleList([
ExpertFFN(d_model, d_ff) for _ in range(num_experts)
])
# 路由网络：输入 x -> N 个专家的分数
self.router = nn.Linear(d_model, num_experts, bias=False)
def forward(self, x):
# x: (batch, seq_len, d_model)
batch_size, seq_len, d_model = x.shape
x_flat = x.view(-1, d_model) # (batch * seq_len, d_model)
# 1. 路由：计算每个 token 对每个专家的分数
router_logits = self.router(x_flat) # (batch * seq_len, num_experts)
router_weights = torch.softmax(router_logits, dim=-1)
# 2. Top-K 选择：每个 token 只选 top_k 个专家
top_k_weights, top_k_indices = torch.topk(router_weights, self.top_k, dim=-1)
top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True) # 重归一化
# 3. 专家计算（简化版：逐 token 循环，实际实现会批量化）
output = torch.zeros_like(x_flat)
for i in range(x_flat.shape[0]):
for k in range(self.top_k):
expert_idx = top_k_indices[i, k]
expert_output = self.experts[expert_idx](x_flat[i:i+1])
output[i] += top_k_weights[i, k] * expert_output.squeeze(0)
return output.view(batch_size, seq_len, d_model)
```
> [!note]
> 真实实现（如 Megatron-LM、MegaBlocks）会先按专家把 token 分桶，再对每个专家做一次批量矩阵乘，彻底避免逐 token 循环；token 数不齐导致的「不规则矩阵乘」正是 [[Lecture 6 · Kernel 与 Triton]] 里 MegaBlocks 这类块稀疏 kernel 要解决的问题。多 GPU 上还要做 expert parallelism（见 §6）。
### 2.2 共享专家（Shared Experts）
**问题**：某些功能（基础语法、高频词处理）是**所有 token 都需要的**。如果只有路由专家，每个专家都得各自学一份这种通用功能——参数冗余，还挤占了专业化的容量。
**解决方案**（DeepSeekMoE / Qwen-MoE）：除了 $N$ 个「路由专家」，额外加 $M$ 个「共享专家」，**每个 token 都无条件经过**：
$$y = \underbrace{\sum_{j=1}^{M} E_j^{\text{shared}}(x)}_{\text{稠密激活，所有 token 都走}} + \underbrace{\sum_{i \in \mathrm{TopK}(x)} g_i(x)\, E_i^{\text{routed}}(x)}_{\text{稀疏激活，top-}K}$$
**DeepSeek-V3 配置**：256 个路由专家（每 token 激活 8 个）+ 1 个共享专家。
> [!tip] 直觉
> 共享专家负责「通识」，路由专家负责「专科」。DeepSeekMoE 的消融显示：去掉共享专家后，路由专家之间的参数冗余明显上升（很多专家被迫重复学习通用功能）。用公司打比方：共享专家是所有项目共用的行政与 IT 部门，路由专家是各业务线——没有前者，每条业务线都得自建一套后勤。
---
## 3. 路由函数设计
路由函数 $g(x)$ 决定**每个 token 去哪些专家**——这是 MoE 的核心设计点，也是所有麻烦（均衡、稳定性、通信）的源头。
### 3.1 Top-K 路由（主流方案）
**思路**：用一个可学习的线性层给 $N$ 个专家打分，softmax 后取分数最高的 $K$ 个，并在选中集合内重新归一化权重：
$$s = \mathrm{softmax}(W_r\, x) \in \mathbb{R}^N, \qquad \mathcal{K} = \mathrm{TopK}(s, K)$$
$$g_i(x) = \begin{cases} \dfrac{s_i}{\sum_{j \in \mathcal{K}} s_j} & i \in \mathcal{K} \\[2mm] 0 & \text{否则} \end{cases} \qquad\Longrightarrow\qquad y = \sum_{i \in \mathcal{K}} g_i(x)\, E_i(x)$$
符号说明：$W_r \in \mathbb{R}^{N \times d_{\text{model}}}$ 是路由权重矩阵（整个 MoE 层里唯一的「决策者」，参数量小得可以忽略），$s_i$ 是 token 对专家 $i$ 的归一化分数，$\mathcal{K}$ 是选中的 $K$ 个专家下标。重归一化保证选中专家的权重和为 1。
**超参数**：
| 参数 | 典型值 | 说明 |
| :--- | ---: | :--- |
| **K（top-k）** | 2–8 | DeepSeek-V3 = 8，Mixtral = 2 |
| **N（专家总数）** | 8–256 | DeepSeek-V3 = 256，Mixtral = 8 |
| **激活比例** | K/N | DeepSeek-V3 = 8/256 ≈ 3.1% |
```mermaid
flowchart LR
X["输入 token x"] --> R["W_router · x"]
R --> S["Softmax"]
S --> T["Top-K 选择"]
T --> E1["专家 1
权重 0.6"]
T --> E2["专家 2
权重 0.4"]
T -.丢弃.-> E3["专家 3
权重 0.05"]
E1 --> O["0.6·E₁(x) + 0.4·E₂(x)"]
E2 --> O
style R fill:#fff9c4,stroke:#f57f17
style S fill:#e1f5fe,stroke:#0277bd
style T fill:#ffccbc,stroke:#d84315
style E1 fill:#c8e6c9,stroke:#2e7d32
style E2 fill:#c8e6c9,stroke:#2e7d32
style E3 fill:#e0e0e0,stroke:#757575
```
> [!note] Top-K 是离散选择，梯度怎么流？
> $\mathrm{TopK}$ 本身不可导，但**被选中专家的门控权重 $g_i(x)$ 对路由参数 $W_r$ 可导**：如果专家 $i$ 的输出对降低 loss 有贡献，梯度会推高 $s_i$，让它下次更容易被选中。未选中的专家收不到任何梯度——这是一种有偏但实践中极其有效的启发式。理论上更「干净」的做法（REINFORCE 等随机梯度估计）反而因方差大、不稳定而被淘汰。
> [!warning] Top-K 的代价
> 路由是可学习的，梯度会让模型「钻空子」——训练早期表现稍好的专家会被更多选中、得到更多训练、变得更好……正反馈之下，少数专家过载、其余闲置。这就是**专家崩塌**，需要显式的负载均衡机制（见 §4）。
### 3.2 其他路由策略（对比）
| 路由方式 | 原理 | 优势 | 劣势 | 代表工作 |
| :--- | :--- | :--- | :--- | :--- |
| **Top-K** | 可学习打分 + topk | 灵活，性能好 | 需要负载均衡 | Switch, Mixtral, DeepSeek |
| **Hash 路由** | 固定哈希函数按 token id 分配 | 负载天然均衡、无需学习 | 无法按语义分工，性能上限低 | Hash Layers (Roller et al., 2021) |
| **线性分配** | 把路由建模为最优分配问题求解 | 均衡有保证 | 求解开销大，推理时不好用 | BASE Layers (2021) |
| **随机路由** | 随机采样 K 个专家 | 简单，天然均衡 | 无法学习专业化 | 早期实验 |
| **RL 路由** | 把离散选择当动作，用策略梯度训练 | 理论上无偏 | 方差大、训练不稳定 | Bengio et al. 早期条件计算 |
RL 路由的失败很有启发性：路由本质上是一个离散决策问题（可类比 CS221 [[Lecture 8 · 强化学习]] 里的动作选择），但工业界最终选择了「假装可导」的 top-k 启发式——**简单、有偏、稳定**胜过**优雅、无偏、高方差**。
**业界现状**（2025）：**Top-K + 均衡机制**是唯一被大规模反复验证的方案。
### 3.3 路由的「专业化」现象
训练后统计各专家实际接收的 token，会发现**专家自发专业化**：某些专家集中处理特定语言、代码或数字/标点等 token 类型（下图为示意，编号虚构）：
```mermaid
graph TD
T1["中文 token
你好"] --> E42["专家 42
中文倾向"]
T2["代码 token
def func"] --> E137["专家 137
代码倾向"]
T3["数学 token
∫ x dx"] --> E201["专家 201
数学符号倾向"]
T4["通用 token
the cat"] --> Shared["共享专家
通识"]
style E42 fill:#ffccbc,stroke:#d84315
style E137 fill:#c8e6c9,stroke:#2e7d32
style E201 fill:#b3e5fc,stroke:#0277bd
style Shared fill:#ffe0b2,stroke:#e65100
```
> [!tip] 直觉
> 为什么会专业化？梯度优化形成正反馈：擅长某类 token 的专家 → 路由分数高 → 更常接到那类 token → 在那类 token 上练得更好。共享专家的存在让路由专家不必浪费容量在通识上，进一步促进分化。
> [!warning] 易错点
> 「专业化」不等于人类直觉的「学科分工」。Mixtral 论文专门分析过：其专家**并没有**按主题（数学、生物……）分化，分化更多体现在 token 级、语法级的浅层模式上（如缩进、标点、高频词）。OLMoE 与 DeepSeek 的分析则在更细粒度的设置下观察到更清晰的领域倾向。结论：专业化真实存在，但它是优化「顺手」形成的统计分工，别把专家过度拟人化成「数学家」「程序员」。
---
## 4. 负载均衡
### 4.1 问题：专家崩塌（Expert Collapse）
Top-K 路由是可学习的，放任梯度自由更新会导致**少数专家承担绝大多数 token**：
$$\text{理想（8 专家）：每个专家 } 12.5\% \qquad \text{崩塌：专家 1 占 } 78\%,\ \text{专家 2 占 } 15\%,\ \text{其余各} \sim 1\%$$
**后果**：
- 参数利用率崩坏——闲置专家学不到东西，MoE 退化成一个昂贵的小稠密模型；
- 性能退化——过载专家被迫身兼百职，失去专业化收益；
- 系统灾难——专家并行时，某几块 GPU 过载而其余空转，吞吐由最慢者决定。
### 4.2 解决方案：辅助损失（Auxiliary Loss）
在语言建模交叉熵之外，额外加一个**负载均衡损失**惩罚不均衡的路由。Switch Transformer (2021) 的经典形式：
$$\mathcal{L}_{\text{aux}} = \alpha \cdot N \sum_{i=1}^{N} f_i \, P_i$$
符号说明：$f_i$ 是实际被路由到专家 $i$ 的 token **占比**（batch 内统计量，不可导），$P_i$ 是专家 $i$ 在所有 token 上的**平均路由概率**（softmax 输出的均值，可导），$N$ 是专家数，$\alpha$ 是平衡系数（典型值 0.01）。完全均衡时 $f_i = P_i = \frac{1}{N}$，损失取最小值 $\alpha$。
```python
def load_balancing_loss(router_probs, expert_indices, num_experts):
"""
router_probs: (batch * seq_len, num_experts) - 路由概率
expert_indices: (batch * seq_len, top_k) - 被选中的专家
"""
# 1. 每个专家被选中的频率（fraction of tokens）
expert_mask = F.one_hot(expert_indices, num_experts).sum(dim=1) # (B*T, num_experts)
f_i = expert_mask.float().mean(dim=0) # (num_experts,) - 每个专家的负载比例
# 2. 每个专家的平均路由概率
P_i = router_probs.mean(dim=0) # (num_experts,)
# 3. 辅助损失 = num_experts * Σᵢ (f_i · P_i)
# 直觉：f_i 高（被选多）且 P_i 高（分数高）的专家会被惩罚
aux_loss = num_experts * (f_i * P_i).sum()
return aux_loss
# 总损失
total_loss = ce_loss + alpha * load_balancing_loss(...)
# ↑ 平衡系数，通常 0.01
```
**为什么是 $f_i \cdot P_i$ 这个形式？** 两个巧妙之处：
- $f_i$ 是「谁真的过载」的准确统计，但它来自不可导的计数；$P_i$ 可导但只是「意图」。乘积形式让梯度**通过 $P_i$ 流动**、方向由 $f_i$ 加权——过载专家（$f_i$ 大）的路由概率被压得最狠；
- 在 $\sum_i f_i$ 与 $\sum_i P_i$ 固定、且 $f$ 与 $P$ 正相关的条件下，$\sum_i f_i P_i$ 在均匀分布处取最小——所以最小化它就是在推动均衡。
> [!tip] 直觉
> 辅助损失本质是给主目标加一个「反垄断税」：谁接的活越多、要价越高，谁被课的税越重。$\alpha$ 控制税率——太小压不住崩塌，太大则强行均分、牺牲专业化。它是 MoE 训练里最重要的超参数之一。
### 4.3 其他均衡策略
| 方法 | 原理 | 优劣 |
| :--- | :--- | :--- |
| **Expert capacity** | 每个专家设硬上限 $C = \frac{T}{N} \times c_f$（$T$ 为 token 总数，$c_f$ 为容量因子），超出的 token 被丢弃（走残差直通）或溢出到次选专家 | 强制均衡、显存可预算，但丢 token 损失信息 |
| **噪声注入** | 路由 logits 加高斯噪声，鼓励探索 | 简单，但噪声尺度难调 |
| **Sinkhorn 归一化** | 迭代算法把路由矩阵行列和都拉平 | 理论优雅，计算开销大 |
**现代共识**（2025）：辅助损失为主，capacity 限制为辅（训练时容量因子常取 1.0–1.5）；DeepSeek-V3 则进一步走向无辅助损失的 bias 调节（见 §8.3）。
---
## 5. 训练稳定性
MoE 比稠密模型更难训——router 引入了新的 softmax、新的离散性、新的对称性问题。
### 5.1 Router Z-Loss
**问题**：路由 logits 可能数值爆炸（某个专家的分数持续走高），softmax 饱和后梯度异常，甚至 NaN。这与 [[Lecture 3 · 架构与超参数]] 中输出层 z-loss 面对的是同一类问题——softmax 对 logits 整体尺度不敏感，交叉熵不会自动约束它。
**解决方案**（ST-MoE, 2022）：惩罚路由 logits 的 log-partition 函数的平方：
$$\mathcal{L}_z = \beta \cdot \frac{1}{B} \sum_{b=1}^{B} \left( \log \sum_{j=1}^{N} e^{\ell_j^{(b)}} \right)^2$$
符号说明：$\ell_j^{(b)}$ 是第 $b$ 个 token 对专家 $j$ 的路由 logit（softmax 前），$B$ 是 batch 内 token 数，$\beta$ 是系数（典型值 0.001）。$\mathrm{logsumexp}$ 是 logits 的光滑最大值，压住它就间接压住了最大 logit。
```python
def router_z_loss(router_logits):
"""
router_logits: (batch * seq_len, num_experts) - 路由打分（softmax 前）
"""
# Z-loss = log(Σⱼ exp(logitⱼ))²
# 等价于惩罚 logsumexp 的平方
log_z = torch.logsumexp(router_logits, dim=-1) # (B*T,)
z_loss = (log_z ** 2).mean()
return z_loss
# 总损失
total_loss = ce_loss + alpha * load_balancing_loss + beta * router_z_loss
# ↑ 通常 beta = 0.001
```
> [!note]
> 于是完整的 MoE 训练目标是三项之和：$\mathcal{L} = \mathcal{L}_{\text{CE}} + \alpha\,\mathcal{L}_{\text{aux}} + \beta\,\mathcal{L}_z$——主目标学语言，$\mathcal{L}_{\text{aux}}$ 管「分得匀」，$\mathcal{L}_z$ 管「分得稳」。工程上还会叠加 QK-norm 与梯度裁剪（clip_grad_norm = 1.0）等稠密模型的常规手段。
### 5.2 专家参数初始化
MoE 对初始化更敏感：如果 $N$ 个专家初始化完全相同且都收到相同梯度（例如稠密门控下），对称性会让它们难以分化——分化只能依赖路由打分的微小差异与选择的随机性，速度很慢。
**标准做法**是让每个专家独立随机初始化，从一开始就打破对称：
```python
# 每个专家独立初始化，用不同的随机种子
for i, expert in enumerate(self.experts):
torch.manual_seed(base_seed + i)
expert.reset_parameters()
```
（Upcycling 场景刻意反其道而行——所有专家从同一份稠密 FFN 复制而来，靠路由的随机性慢慢分化，这正是它的主要风险，见 §7。）
### 5.3 FP32 路由权重
**经验规则**：即使模型主体用 BF16/FP16 训练，**路由网络的权重与 logits 保持 FP32**。路由要在 $N$ 个接近的分数里排序选 top-k，低精度的舍入误差足以翻转选择结果，造成训练轨迹抖动。
```python
class MoELayer(nn.Module):
def __init__(self, ...):
self.router = nn.Linear(d_model, num_experts, bias=False)
self.router.to(torch.float32) # 强制 FP32
def forward(self, x):
router_logits = self.router(x.float()) # 转 FP32 计算
# ... 后续用 FP32 的 router_logits
```
### 5.4 微调阶段的过拟合
ST-MoE 的另一个重要发现：**稀疏模型在小数据微调时比同激活量的稠密模型更容易过拟合**——参数多、有效数据少，路由还可能在分布偏移下漂移。缓解手段包括：微调时只更新部分参数（如冻结路由或只调非专家参数）、加大微调阶段的 dropout、或用更大的微调数据。这是 MoE「训练便宜、用起来娇贵」的一面，在做对齐（[[Lecture 15 · 对齐 I：SFT 与 RLHF]]）时需要记住。
> [!tip] 直觉
> 本节四个技巧对应 MoE 的四个新脆弱点：router 的 softmax（z-loss）、专家的对称性（独立初始化）、离散选择对噪声的敏感（FP32 路由）、参数量与微调数据的失衡（过拟合）。共同主线：**稀疏性把「连续优化」变成了「连续优化 + 离散分配」的混合问题**，所有额外的护栏都是在替离散那一半擦屁股。
---
## 6. 并行训练策略
MoE 的稀疏性带来新的并行挑战——传统的数据并行 + 张量并行不够用（并行策略的系统性介绍见 [[Lecture 7 · 并行 I：基础]] 与 [[Lecture 8 · 并行 II：分布式训练实战]]）。
### 6.1 Expert Parallelism（专家并行）
**思路**：总参数量太大，单卡装不下所有专家——把 $N$ 个专家**分布到不同设备**，每个设备只持有一部分专家；token 则「走出去」找专家。
```mermaid
graph TB
subgraph "GPU 0"
E1["专家 1-64"]
end
subgraph "GPU 1"
E2["专家 65-128"]
end
subgraph "GPU 2"
E3["专家 129-192"]
end
subgraph "GPU 3"
E4["专家 193-256"]
end
T["Token batch"] --> R["路由决策"]
R -->|All-to-All| E1
R -->|All-to-All| E2
R -->|All-to-All| E3
R -->|All-to-All| E4
E1 -->|All-to-All| O["输出"]
E2 -->|All-to-All| O
E3 -->|All-to-All| O
E4 -->|All-to-All| O
style T fill:#e1f5fe,stroke:#0277bd
style R fill:#fff9c4,stroke:#f57f17
style O fill:#c8e6c9,stroke:#2e7d32
```
**通信模式**——每个 MoE 层前向要做两次 All-to-All（反向再来两次）：
1. **All-to-All（dispatch）**：按路由结果把每个 token 的激活发到其专家所在的 GPU；
2. **本地专家计算**：各 GPU 对收到的 token 批量跑自己的专家；
3. **All-to-All（combine）**：把专家输出送回 token 原属的 GPU 做加权求和。
单向通信量约为 $B \times L \times K \times d_{\text{model}}$ 个激活元素（每个 token 要发给 $K$ 个专家）——它随 batch 与序列长度线性增长，且无法像梯度 All-Reduce 那样容易与计算重叠。
> [!warning] 瓶颈
> All-to-All 是带宽密集型操作，且跨节点带宽（InfiniBand）比节点内 NVLink 低一个量级以上。token dispatch 一旦跨节点，通信时间可能超过专家计算本身——这是 MoE 扩展的主要系统限制，也是 §8 里 DeepSeek 各种「设备级路由」设计的动机。
### 6.2 混合并行（Hybrid Parallelism）
大规模 MoE 训练把多种并行叠起来用：
| 并行方式 | 分割维度 | 用于 | 通信原语 |
| :--- | :--- | :--- | :--- |
| **Data Parallelism** | batch | 所有层 | All-Reduce（梯度） |
| **Tensor Parallelism** | 矩阵列/行 | 注意力、FFN 内部 | All-Reduce（前向/反向） |
| **Expert Parallelism** | 专家数 | MoE 层 | All-to-All |
| **Pipeline Parallelism** | 层深度 | 整个模型 | P2P（send/recv） |
**DeepSeek-V3 实际配置**（技术报告）：
- 2048 块 H800 GPU；
- **16 路流水线并行**（自研 DualPipe 调度，用计算-通信重叠掩盖 All-to-All）；
- **64 路专家并行**，跨 8 个节点——256 个路由专家分摊后**每块 GPU 持有 4 个专家**；
- 数据并行用 **ZeRO-1**（切分 optimizer state）；
- **不使用张量并行**——TP 的 All-Reduce 太频繁，在 H800 的带宽条件下不划算。
```mermaid
graph LR
subgraph "DeepSeek-V3 并行策略堆叠"
A["2048 块 H800"] -->|"PP = 16"| B["流水线 16 段
DualPipe 调度"]
B -->|"EP = 64"| C["每 EP 组 64 卡
每卡 4 个专家"]
C -->|"ZeRO-1 DP"| D["切分 optimizer state
的数据并行"]
end
style A fill:#e1f5fe,stroke:#0277bd
style B fill:#c8e6c9,stroke:#2e7d32
style C fill:#fff9c4,stroke:#f57f17
style D fill:#ffccbc,stroke:#d84315
```
> [!tip] 调优经验
> 专家并行度**不必等于专家数**——每卡放多个专家反而让本地矩阵乘更「胖」、效率更高（V3 是 256 专家 / EP64 = 每卡 4 个）。真正的关键是**把 All-to-All 尽量按在节点内**：路由层面用设备/节点受限路由（§8.2、§8.3），调度层面用 DualPipe 这类重叠技术把通信藏进计算里。GPU 带宽层级的量化直觉见 [[Lecture 5 · GPU]]。
---
## 7. 从稠密模型到 MoE：Upcycling
**问题**：从零训练 MoE 需要海量计算——能否**复用已有的稠密模型**？
**Upcycling（稀疏再利用，Komatsuzaki et al., 2022）**：把训练好的稠密模型**改造成 MoE**，再继续训练：
1. 把原稠密 FFN 复制 $N$ 份，作为 $N$ 个专家的初始化；
2. 插入随机初始化的 router，通常从保守的 top-k 设置起步；
3. 用足够长的继续预训练让专家分化；
4. 再进入指令微调 / 对齐阶段。
```mermaid
flowchart LR
A["Dense LM"] --> B["复制 FFN 权重"]
B --> C["插入 Router"]
C --> D["MoE LM"]
D --> E["继续预训练"]
E --> F["SFT / 对齐"]
style A fill:#e1f5fe,stroke:#0277bd
style D fill:#fff9c4,stroke:#f57f17
style F fill:#c8e6c9,stroke:#2e7d32
```
> [!tip] 直觉
> Upcycling 不是「免费变强」，而是把稠密模型已学到的通用表示当作**初始化**，避免 MoE 从随机专家开始摸索。它的先天矛盾是：初始化时所有专家完全相同（对称），而 MoE 的价值恰恰来自专家**不同**——所以成败取决于继续训练能否用足够的数据、合适的路由正则与负载均衡把对称性打破（对照 §5.2：从零训练时会刻意独立初始化来避免这个问题）。
### 7.1 例子：MiniCPM 与 Qwen MoE
官方讲义给了两个典型例子：
| 模型 | 初始化 | MoE 设置 | 讲义中的结论 |
| :--- | :--- | :--- | :--- |
| **MiniCPM MoE** | MiniCPM dense model | top-k=2，8 experts，约 4B active params | 用约 520B tokens 继续训练后，相比 base model 有收益 |
| **Qwen MoE** | Qwen 1.8B | top-k=4，60 experts，4 shared experts | 与 DeepSeekMoE 类似，是较早被确认成功的 upcycling 案例之一 |
> [!warning] 关键点
> Upcycling 的难点不在「把 FFN 复制 N 份」，而在于**避免所有专家学成同一个函数**。如果 router 过早塌缩，或继续训练的数据量不足，得到的只是一个参数更多、利用率更差的模型。MiniCPM 的 520B token 继续训练量说明：打破对称是要付出真金白银的算力的。
---
## 8. DeepSeek MoE：从 v1 到 v3
Lecture 4 最后用 DeepSeek 系列把前面的概念串起来：共享专家、细粒度专家、设备级路由、通信均衡、aux-loss-free balancing 都不是孤立技巧，而是在回答同一个问题：**如何让稀疏激活既有效（数学上均衡、分工细）又跑得动（系统上通信可控）**。
### 8.1 DeepSeek MoE v1
| 维度 | 设计 |
| :--- | :--- |
| 参数规模 | 16B total，约 2.8B active |
| 专家形态 | 2 个 shared experts + 64 个 routed experts（每个专家为标准 FFN 的 1/4 大小，激活 6 个） |
| 路由 | 标准 top-k routing |
| 均衡 | expert-level + device-level auxiliary loss |
**理解方式**：
- shared experts 承接所有 token 共需的通用能力，消除路由专家间的功能冗余；
- 细粒度 routed experts（1/4 尺寸 × 4 倍数量）用组合数优势换更细的分工；
- device-level balancing 说明 DeepSeek 从 v1 起就把**系统侧的均衡**（每台设备的负载）当作与数学目标同级的约束，而不只是优化一个漂亮的 loss。
### 8.2 DeepSeek MoE v2
| 维度 | 设计 |
| :--- | :--- |
| 参数规模 | 236B total，约 21B active |
| 专家形态 | 2 个 shared experts + 160 个 routed experts，细粒度比例约 1/10 |
| 激活 | 6 个 routed experts active |
| 新增机制 | top-M device routing（每个 token 最多访问 M=3 台设备）；communication balancing loss |
**Top-M device routing** 的动机很朴素：如果一个 token 的 top-k 专家散落在太多设备上，All-to-All 的扇出会非常难看。先把 token 可访问的**设备集合**限制在 M 台（按设备上专家的最高分选设备），再在这些设备内部选专家——把模型质量和通信成本放进同一个优化框架。
**Communication balancing loss** 进一步平衡系统侧指标：每台设备**收到**多少 token、**发出**多少输出、每个 MoE 层是否出现局部热点——本质是把 §4 的均衡思想从「专家维度」推广到「设备维度」。
> [!tip] 直觉
> MoE 的扩展瓶颈经常不是 FLOPs，而是跨设备 token dispatch。路由策略如果只追求 loss、不管通信拓扑，训练吞吐会被 All-to-All 直接打穿。DeepSeek 的演进路线就是不断把「通信」显式写进路由的约束里——算法为系统让路，这门课反复出现的主题。
### 8.3 DeepSeek MoE v3
| 维度 | 设计 |
| :--- | :--- |
| 参数规模 | 671B total，约 37B active |
| 专家形态 | 1 个 shared expert + 256 个细粒度 routed experts（前 3 层保持稠密） |
| 激活 | 8 个 routed experts active |
| 路由 | sigmoid 打分；按「分数 + per-expert bias」选 top-k；门控权重在选中集合内归一化 |
| 均衡 | per-expert bias 的 auxiliary-loss-free balancing + 很轻的 sequence-wise auxiliary loss；node-limited routing（每 token 至多 4 个节点） |
v3 路由的两个细节值得展开：
- **sigmoid 替代 softmax 打分**：softmax 让专家分数互相竞争（一个升、其余必降），sigmoid 让每个专家的分数独立——细粒度大 $N$ 下更稳定，也让 bias 调节不会牵一发动全身；
- **auxiliary-loss-free balancing**：给每个专家维护一个**只影响 top-k 选择、不参与门控权重**的 bias $b_i$，按负载在线调整。
$$s_i = \sigma(w_i^\top x), \qquad \mathcal{K} = \mathrm{TopK}(s_i + b_i,\ K), \qquad g_i = \frac{s_i}{\sum_{j \in \mathcal{K}} s_j}$$
注意 $b_i$ 只出现在选择步骤里，输出加权仍用原始分数 $s_i$——均衡干预不污染模型的函数值。
```python
# 伪代码：per-expert bias 的在线更新直觉
if expert_load[i] < target_load:
bias[i] += step_size # 专家吃得太少，下次更容易被选中
else:
bias[i] -= step_size # 专家吃得太多，下次降低优先级
```
最值得记住的一点：**v3 没有把负载均衡问题「取消掉」，只是把它从显式 auxiliary loss（会给主目标注入干扰梯度）迁移到不产生梯度的在线 bias 调节里**，再留一个很轻的序列级辅助项兜底。均衡压力还在，只是换了个不伤主目标的施力点。
---
## 9. 补充：DeepSeek-V3 还用了什么？
严格说，下面两点已经超出 MoE 本身，但官方 Lecture 4 在结尾把它们放进 DeepSeek-V3 的整体架构里讲，所以值得放在这里作索引。
### 9.1 MLA：Multi-head Latent Attention
**核心想法**：把 K/V 表示成一个低维潜变量（latent）的函数，缓存潜变量而非完整 K/V，从而压缩 KV cache——与 [[Lecture 3 · 架构与超参数]] 里 GQA 的目标相同，但走「低秩压缩」而非「共享 head」的路线。
$$c_t^{KV} = W^{DKV} h_t \quad (\text{降维，只缓存这个}), \qquad k_t^C = W^{UK} c_t^{KV}, \quad v_t^C = W^{UV} c_t^{KV} \quad (\text{用时再升维})$$
符号说明：$h_t$ 是第 $t$ 个 token 的隐状态，$W^{DKV}$ 把它压到低维 $c_t^{KV}$（DeepSeek-V2 中维度仅为 $4 d_{\text{head}}$），$W^{UK}, W^{UV}$ 在计算注意力时再把潜变量投影回各 head 的 K/V。
**收益**：
- DeepSeek-V2 报告 KV cache 缩减约 **93%**，同时质量优于 GQA 同级配置——因为压缩矩阵是学出来的，而不是像 GQA 那样硬性共享；
- query 侧也可做类似压缩，降低训练期激活内存；
- 对长上下文推理尤其重要（KV cache 是长上下文的主要内存瓶颈，见 [[Lecture 10 · 推理]]）。
**复杂之处**：RoPE 的旋转依赖各 head 的完整 key，与「只缓存低维潜变量」不兼容。解法是**解耦**：留出一小部分维度作为不经压缩、可正常旋转的 RoPE key，与潜变量恢复出的内容 key 拼接使用。
### 9.2 MTP：Multi-Token Prediction
MTP 的想法是在主模型之外挂一个轻量模块，训练时额外预测更远的未来 token：
$$\text{普通 LM：} x_t \to \hat{x}_{t+1} \qquad\qquad \text{MTP：} x_t \to \hat{x}_{t+1}, \hat{x}_{t+2}, \dots$$
DeepSeek-V3 的实现相当保守——讲义特别指出它实际只做 one-token-ahead 的 MTP 变体（额外预测一个 token）。它的价值主要在于：
- 每个位置多一个预测目标，**训练信号更密**；
- MTP 模块天然可以充当 speculative decoding 的草稿模型（draft model），加速推理；
- 推理时可以直接丢弃 MTP 模块，不增加主模型的部署成本。
---
## 总结
```mermaid
mindmap
root((MoE))
动机
同 FLOPs 装更多参数
总参数与激活参数解耦
架构
只替换 FFN
细粒度专家
共享专家
路由
Top-K 打分
离散选择的梯度启发式
专家专业化
负载均衡
辅助损失
expert capacity
per-expert bias
训练稳定
router z-loss
FP32 路由
独立初始化
微调易过拟合
系统
专家并行
两次 All-to-All
设备受限路由
DeepSeek 演进
v1 共享加细粒度
v2 通信均衡
v3 aux-loss-free
MLA 与 MTP
```
### 本讲核心
| 问题 | MoE 的回答 |
| :--- | :--- |
| 同样 FLOPs 下能否扩大参数量？ | 用稀疏激活：每个 token 只走 top-k 专家 |
| 路由不可导怎么办？ | 工业界主要用 top-k heuristic + softmax/sigmoid + 辅助均衡 |
| 专家负载不均怎么办？ | auxiliary loss、per-device balancing、per-expert bias |
| 为什么 MoE 难训？ | router 不稳定、专家塌缩、fine-tuning 易过拟合、All-to-All 通信复杂 |
| 为什么 MoE 仍然流行？ | 在相同推理 FLOPs 下，更多参数和更强模型质量通常是值得的 |
### 一句话
> **MoE 的本质是把「模型容量」与「每个 token 的计算量」解耦；真正难的不是多放几个专家，而是让路由、负载均衡、通信和训练稳定性同时成立。**
### 和后续课程的连接
- **[[Lecture 5 · GPU]]**（下一讲）：理解 GPU 内存层级、带宽与矩阵乘效率，才能量化 MoE 的系统瓶颈——All-to-All 为什么贵、每卡多个专家为什么反而快。
- **[[Lecture 6 · Kernel 与 Triton]]**：MegaBlocks / 块稀疏 matmul 这类 kernel 优化，本质上是为 MoE 的不规则计算（各专家 token 数不齐）服务。
- **[[Lecture 7 · 并行 I：基础]] / [[Lecture 8 · 并行 II：分布式训练实战]]**：Expert Parallelism 依赖 All-to-All，与 DP/TP/PP 一起组成大规模 MoE 训练栈。
- **[[Lecture 9 · 缩放定律 I：基础]] / [[Lecture 11 · 缩放定律 II：细节]]**：MoE 改变的是 active compute 与 total parameters 的关系，是缩放律里的关键设计变量。
---
## 复习自测
> [!question]- Q1：概念辨析——MoE 省的是「算」还是「存」？总参数与激活参数分别决定什么？
> MoE 省算不省存。每 token 的训练/推理 FLOPs 正比于**激活参数**（如 DeepSeek-V3 的 37B），所以计算成本像个小模型；但**总参数**（671B）无论激活与否都要放进显存/内存，部署门槛、通信拓扑都由它决定。这也是为什么 MoE 必须配专家并行——不是为了算得快，而是单卡根本装不下。
> [!question]- Q2：Top-K 选择本身不可导，为什么 router 还能被梯度训练出来？
> 因为被选中专家的门控权重 $g_i(x) = s_i / \sum_{j \in \mathcal{K}} s_j$ 对路由参数可导：某专家的输出对降 loss 有贡献时，梯度会推高其分数，使其更容易再次被选中。未选中专家收不到梯度——这是有偏的启发式，但比 REINFORCE 等无偏随机估计稳定得多，被工业界普遍接受。
> [!question]- Q3：Switch Transformer 的辅助损失 $\alpha N \sum_i f_i P_i$ 里，为什么要同时用不可导的 $f_i$ 和可导的 $P_i$？均衡时它取什么值？
> $f_i$（实际负载占比）是准确的均衡度量但来自计数、不可导；$P_i$（平均路由概率）可导但只是意图。乘积让梯度经 $P_i$ 流动、力度由 $f_i$ 加权——过载专家的路由概率被压得最狠。完全均衡时 $f_i = P_i = 1/N$，损失取最小值 $\alpha \cdot N \cdot N \cdot \frac{1}{N^2} = \alpha$。
> [!question]- Q4：为什么专家并行的每个 MoE 层前向需要两次 All-to-All？瓶颈为何出现在跨节点？
> 第一次（dispatch）把每个 token 的激活按路由结果发到专家所在 GPU，第二次（combine）把专家输出送回 token 原属 GPU 做加权求和——token 的「出差」必须有去有回。单向通信量约 $B \times L \times K \times d_{\text{model}}$，随 batch 和序列长度线性增长；而跨节点带宽比节点内 NVLink 低一个量级以上，dispatch 一旦大量跨节点，通信就会盖过计算。这正是 top-M device routing / node-limited routing 的动机。
> [!question]- Q5：共享专家解决什么问题？DeepSeek-V3 的 per-expert bias 相比辅助损失又好在哪里？
> 共享专家承接所有 token 共需的通用功能，去掉它时每个路由专家都得重复学一份通识，参数冗余、专业化受挤压（DeepSeekMoE 消融验证）。per-expert bias 只作用于 top-k 的**选择**步骤且不参与梯度，负载失衡时在线加减 bias 即可调节流量——均衡压力不再以辅助损失梯度的形式干扰主目标的优化，门控权重也不被污染。
---
## 参考资料
- [CS336 Spring 2025 Schedule](https://stanford-cs336.github.io/spring2025/)
- [Lecture 4 PDF: Mixtures of Experts](https://github.com/stanford-cs336/spring2025-lectures/blob/main/nonexecutable/2025%20Lecture%204%20-%20MoEs.pdf)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://arxiv.org/abs/2006.16668)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://arxiv.org/abs/2202.08906)
- [Mixtral of Experts](https://arxiv.org/abs/2401.04088)
- [DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models](https://arxiv.org/abs/2401.06066)
- [DeepSeek-V3 Technical Report](https://arxiv.org/abs/2412.19437)
- [Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints](https://arxiv.org/abs/2212.05055)
- [Hash Layers for Large Sparse Models](https://arxiv.org/abs/2106.04426)
- [OLMoE: Open Mixture-of-Experts Language Models](https://arxiv.org/abs/2409.02060)