Lecture 4 · 混合专家模型 MoE

CS336: Language Modeling from Scratch · Stanford · Spring 2025
📅 Apr 10 · 💻 不可执行讲义（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 计算量不变，总参数量成倍增长。

---

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 实际消耗的计算？

graph TD
    A["想要更好的模型"] --> B["缩放定律：需要更大 N 和更多 D"]
    B --> C["计算成本 ∝ N·D"]
    C --> D["训练成本爆炸"]
    D --> E{"有没有办法<br/>只增加 N<br/>但不增加 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

1.2 MoE 的核心思想

把 Transformer 的 FFN 层拆成 $N$ 个「专家（expert）」，每次只用其中 $K$ 个：

flowchart TB
    subgraph Dense["稠密 FFN（传统）"]
        D1["输入 x"] --> D2["一个大 FFN<br/>所有参数都激活"]
        D2 --> D3["输出"]
    end

    subgraph MoE["MoE FFN（稀疏）"]
        M1["输入 x"] --> Router["路由函数 g(x)<br/>选出 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$ 个专家的下标集合。

---

2. MoE 架构细节

2.1 专家的粒度

问题：应该在 Transformer 的哪个部件用 MoE？每个专家应该多大、多少个？

业界共识（2025）：

graph LR
    A["Transformer Block"] --> B["Multi-Head Attention<br/>（保持稠密）"]
    B --> C["Add & Norm"]
    C --> D["MoE FFN Layer<br/>（替换原 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}$，专家分工可以更细、冗余更少。

代码示例（简化的 MoE FFN）：

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)

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 个共享专家。

---

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%

flowchart LR
    X["输入 token x"] --> R["W_router · x"]
    R --> S["Softmax"]
    S --> T["Top-K 选择"]
    T --> E1["专家 1<br/>权重 0.6"]
    T --> E2["专家 2<br/>权重 0.4"]
    T -.丢弃.-> E3["专家 3<br/>权重 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

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 类型（下图为示意，编号虚构）：

graph TD
    T1["中文 token<br/>你好"] --> E42["专家 42<br/>中文倾向"]
    T2["代码 token<br/>def func"] --> E137["专家 137<br/>代码倾向"]
    T3["数学 token<br/>∫ x dx"] --> E201["专家 201<br/>数学符号倾向"]
    T4["通用 token<br/>the cat"] --> Shared["共享专家<br/>通识"]

    style E42 fill:#ffccbc,stroke:#d84315
    style E137 fill:#c8e6c9,stroke:#2e7d32
    style E201 fill:#b3e5fc,stroke:#0277bd
    style Shared fill:#ffe0b2,stroke:#e65100

---

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$。

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$ 在均匀分布处取最小——所以最小化它就是在推动均衡。

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。

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

5.2 专家参数初始化

MoE 对初始化更敏感：如果 $N$ 个专家初始化完全相同且都收到相同梯度（例如稠密门控下），对称性会让它们难以分化——分化只能依赖路由打分的微小差异与选择的随机性，速度很慢。

标准做法是让每个专家独立随机初始化，从一开始就打破对称：

# 每个专家独立初始化，用不同的随机种子
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，低精度的舍入误差足以翻转选择结果，造成训练轨迹抖动。

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）时需要记住。

---

6. 并行训练策略

MoE 的稀疏性带来新的并行挑战——传统的数据并行 + 张量并行不够用（并行策略的系统性介绍见 Lecture 7 · 并行 I：基础 与 Lecture 8 · 并行 II：分布式训练实战）。

6.1 Expert Parallelism（专家并行）

思路：总参数量太大，单卡装不下所有专家——把 $N$ 个专家分布到不同设备，每个设备只持有一部分专家；token 则「走出去」找专家。

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 那样容易与计算重叠。

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 的带宽条件下不划算。

graph LR
    subgraph "DeepSeek-V3 并行策略堆叠"
        A["2048 块 H800"] -->|"PP = 16"| B["流水线 16 段<br/>DualPipe 调度"]
        B -->|"EP = 64"| C["每 EP 组 64 卡<br/>每卡 4 个专家"]
        C -->|"ZeRO-1 DP"| D["切分 optimizer state<br/>的数据并行"]
    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

---

7. 从稠密模型到 MoE：Upcycling

问题：从零训练 MoE 需要海量计算——能否复用已有的稠密模型？

Upcycling（稀疏再利用，Komatsuzaki et al., 2022）：把训练好的稠密模型改造成 MoE，再继续训练：

1. 把原稠密 FFN 复制 $N$ 份，作为 $N$ 个专家的初始化；
2. 插入随机初始化的 router，通常从保守的 top-k 设置起步；
3. 用足够长的继续预训练让专家分化；
4. 再进入指令微调 / 对齐阶段。

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

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 案例之一

---

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 的均衡思想从「专家维度」推广到「设备维度」。

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$——均衡干预不污染模型的函数值。

# 伪代码：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 模块，不增加主模型的部署成本。

---

总结

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 的关系，是缩放律里的关键设计变量。

---

复习自测

---

参考资料

• CS336 Spring 2025 Schedule
• Lecture 4 PDF: Mixtures of Experts
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• ST-MoE: Designing Stable and Transferable Sparse Expert Models
• Mixtral of Experts
• DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models
• DeepSeek-V3 Technical Report
• Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints
• Hash Layers for Large Sparse Models
• OLMoE: Open Mixture-of-Experts Language Models