Transformer 解剖：从 Attention 到推理系统

第 12 章 Mixture of Experts：稀疏激活的工程账

作者杨艺韬 · 5,063 字

第 12 章 Mixture of Experts：稀疏激活的工程账

第 11 章我们看到密集模型在 Scaling 上撞到的几堵墙——数据耗尽、算力上限、边际递减。MoE（Mixture of Experts，专家混合）是工业界绕开其中一堵墙的关键技术：参数量再翻 10 倍，推理算力却几乎不变。

DeepSeek-V3 就是个标志性的例子：总参数 671B（比 Llama-3 70B 大近 10 倍），但每个 token 推理时只激活 37B（比 Llama-3 70B 还小一半）。同样的算力跑出大得多的模型容量——这就是 MoE 卖的核心账。

读完这章你能：

解释 MoE 的核心想法：「条件计算」（conditional computation）；
写出路由器（router）和 Top-K 门控的数学；
理解 MoE 训练里最头疼的问题——负载均衡——以及主流的解决方案；
对比 Mixtral 8x7B、DeepSeek-V2/V3、Qwen MoE、GShard / Switch Transformer 等不同 MoE 设计；
估算一个 MoE 模型的「总参数」「激活参数」「推理 FLOPs」三个不同量。

12.1 一个直觉：稀疏神经网络

先讲一个直觉的故事。

人脑大约有 860 亿神经元。但任何时刻活跃的神经元只占很小一部分——研究估计大约 1-5%。当你在读这段文字时，大量负责「运动」「视觉空间」「记忆其他场景」的神经元都在休眠。大脑用稀疏激活节省能量——这是亿万年进化给生物计算找到的最优路径。

稠密的 Transformer 不是这样工作的。每个 token 流过 FFN 时，所有的 $4d$ 个隐层神经元都被激活、都参与计算——无论这个 token 是英文「the」还是中文「的」、是代码 function 还是数学符号 ∑。这意味着大量计算花在了「与当前 token 无关」的神经元上。

MoE 的核心想法：给 FFN 增加结构，让每个 token 只激活其中一部分。

具体地：把一个大的 FFN 拆成 N 个小的「专家」（experts），每个专家是一个独立的 FFN（同样是两层 MLP）。再加一个路由器（router）——它接收输入 token，决定「这个 token 应该交给哪几个专家处理」。每个 token 只激活 N 个专家中的 K 个（典型 K=2 或 K=8），其他的根本不参与计算。

flowchart LR
  X["输入 token x"] --> R[Router]
  R -->|gate score| E1[Expert 1]
  R -->|gate score| E2[Expert 2]
  R -->|gate score| E3[Expert 3]
  R -->|...| EN[Expert N]
  R --> TOPK[选 Top-K=2 个分数最高的]
  TOPK --> ACTIVE["激活: E2, E3"]
  ACTIVE --> SUM[加权求和]
  SUM --> OUT[输出]
  E1 -.未激活.-> OFF[不计算]
  E4 -.未激活.-> OFF

这件事叫条件计算（conditional computation）：根据输入条件动态决定哪些参数要被计算。它把「模型容量」和「单次推理算力」解耦——你可以有 100 个专家（容量大），但每次只算 2 个（算力小）。

12.2 数学：MoE Layer 的形式

把一个稠密 FFN 替换成 MoE 时，公式是这样：

1. Router 计算每个专家的 gate score：

g(x) = \text{softmax}(W_g \cdot x), \quad g \in \mathbb{R}^N

$W_g \in \mathbb{R}^{d_{\text{model}} \times N}$ 是路由器的权重矩阵。 $g_i$ 是专家 $i$ 被「选中」的分数。

2. 选 Top-K 专家：

\mathcal{T} = \text{TopK}(g, K)

$\mathcal{T}$ 是分数最高的 K 个专家的索引集合。

3. 重新归一化 Top-K 分数：

\tilde{g}_i = \frac{g_i}{\sum_{j \in \mathcal{T}} g_j}, \quad i \in \mathcal{T}

让 K 个被选中专家的权重和为 1。

4. 加权求和被选中专家的输出：

\text{MoE}(x) = \sum_{i \in \mathcal{T}} \tilde{g}_i \cdot E_i(x)

其中 $E_i$ 是专家 $i$ 的 FFN（结构和稠密 FFN 一样，只是单独参数）。

形状记号：

张量	形状
$x$	$(d_{\text{model}})$
$W_g$	$(d_{\text{model}}, N)$
$g$	$(N,)$
Top-K 索引	$(K,)$
每个专家的输出	$(d_{\text{model}})$
MoE 输出	$(d_{\text{model}})$

整个 MoE Layer 的输入输出形状和稠密 FFN 一致—— $(d_{\text{model}})$ 进， $(d_{\text{model}})$ 出。这意味着可以直接替换 Transformer Block 里的 FFN，其余部分（attention、norm、residual）一字不变。

12.3 一个最小代码实现

class MoELayer(nn.Module):
    def __init__(self, d_model, d_ffn, n_experts, top_k):
        super().__init__()
        self.n_experts = n_experts
        self.top_k = top_k
        # N 个独立专家
        self.experts = nn.ModuleList([
            SwiGLU(d_model, d_ffn) for _ in range(n_experts)
        ])
        # Router
        self.router = nn.Linear(d_model, n_experts, bias=False)

    def forward(self, x):
        # x: (B, T, D)
        B, T, D = x.shape
        x_flat = x.view(-1, D)                     # (B*T, D)
        n_tokens = x_flat.size(0)

        # 1. Gate scores
        logits = self.router(x_flat)                # (B*T, N)

        # 2. Top-K 选择
        topk_vals, topk_idx = logits.topk(self.top_k, dim=-1)   # (B*T, K)
        topk_weights = F.softmax(topk_vals, dim=-1)             # 归一化

        # 3. 把 token 按选中的专家分组并送过去
        out = torch.zeros_like(x_flat)
        for i in range(self.n_experts):
            # 找到所有选中专家 i 的 (token, slot) 对
            mask = (topk_idx == i)                              # (B*T, K)
            if mask.any():
                token_idx, k_slot = mask.nonzero(as_tuple=True)
                # 这些 token 的输入
                tokens_for_i = x_flat[token_idx]                # (n_i, D)
                # 经过专家 i
                expert_out = self.experts[i](tokens_for_i)      # (n_i, D)
                # 加权累加到输出
                weights = topk_weights[token_idx, k_slot].unsqueeze(-1)
                out[token_idx] += weights * expert_out

        return out.view(B, T, D)

这是一个最小可读版本，没有做负载均衡损失、没有做 All-to-All 通信优化——但展示了 MoE 的核心数据流。

for i in range(self.n_experts) 这个循环在小规模上能跑，但 N 大时（比如 256）开销大。生产实现用 torch.scatter + torch.gather + grouped GEMM 的组合，把这个循环并行化。

12.4 关键挑战：负载均衡

MoE 训练里最大的工程问题是负载均衡。

设想：你有 8 个专家，每次每个 token 选 Top-2。理想情况下每个专家应该平均分到 1/4 的 token。但实际训练中，模型常常退化到「只用少数几个专家」的模式——比如 8 个专家里 2 个被选 80%，剩下 6 个几乎没机会工作。

为什么？因为路由器是个软目标，模型在训练初期可能因为随机初始化让几个专家「先发优势」——这几个专家被路由更多 token，得到更多梯度，能力增长更快，又被更多选——形成正反馈。最终少数专家承担大部分工作，剩下的「死掉」（dead experts）。

负载不均衡有两个严重后果：

模型容量利用率低：6 个 dead experts 的参数白白浪费，相当于把 N=8 退化成 N=2。
硬件利用率低：训练时专家分布在不同 GPU 上，一个 GPU 跑爆了、其他 GPU 空闲——bottleneck 拖慢整体速度。

flowchart TB
  subgraph ideal ["理想（均匀）"]
    EA["Expert 1: 12.5%"]
    EB["Expert 2: 12.5%"]
    EC["Expert 3: 12.5%"]
    ED["..."]
    EE["Expert 8: 12.5%"]
  end
  subgraph collapse ["崩塌（头部聚集）"]
    BA["Expert 1: 45%"]
    BB["Expert 2: 35%"]
    BC["Expert 3: 8%"]
    BD["Expert 4-7: 各 2%"]
    BE["Expert 8: 0% (dead)"]
  end

Auxiliary Loss：硬性负载均衡

经典解决方案是给训练 loss 加一项辅助损失（auxiliary loss），鼓励均匀使用：

\mathcal{L}_{\text{aux}} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i

其中：

$f_i$ ：专家 $i$ 被选中的比例（actual usage frequency）
$P_i$ ：专家 $i$ 的平均 gate 概率
$\alpha$ ：辅助损失权重（典型 0.01）

直觉：当某个专家既「被频繁选中（ $f_i$ 大）」又「平均分数高（ $P_i$ 大）」时，这一项变大——模型被惩罚，下次会让其他专家分担。这一项让所有专家被均匀使用。

GShard、Switch Transformer、Mixtral 都用这个 auxiliary loss。

Auxiliary-Free Load Balancing

DeepSeek-V3 提出了一种无辅助损失的负载均衡——通过给每个专家加一个动态的 bias，主动调整路由器输出：

g_i' = g_i + b_i

如果某个专家最近被选中太多，就把它的 $b_i$ 调小（让它在下一步路由时少被选）；反之调大。bias 在训练中按梯度自动更新。

这种方案的好处：不会让 auxiliary loss 影响主任务的训练动态，效果在 DeepSeek-V3 实验里更稳定。这是 2024 年的新趋势。

Token Dropping vs Expert Capacity

另一个工程细节：每个 GPU 上的某个专家能处理多少 token 是有上限的（叫 expert capacity）——它的硬件资源固定。如果某个专家被分配的 token 超出 capacity，就要 drop 多余的（直接跳过它的 FFN，残差走过去）或者 overflow（流到其他专家）。

设 capacity factor $C_f = 1.25$ （典型值），则每个专家最多接受 $\frac{B \cdot T \cdot K}{N} \cdot C_f$ 个 token。超出的 drop。这给了 buffer 容忍轻度负载不均衡，但极端不均衡仍然会大量 drop。

12.5 主流 MoE 设计对比

把工业界主流 MoE 模型的关键参数放一起：

模型	总参数	激活参数	专家数	Top-K	共享专家	注意
GShard (Google)	600B	25B	2048	2	否	早期奠基（2020）
Switch Transformer	1571B	8B	2048	1	否	Top-1，训练简单但容易死
GLaM (Google)	1200B	95B	64	2	否
Mixtral 8x7B	47B	13B	8	2	否	开源 MoE 起点
Mixtral 8x22B	141B	39B	8	2	否
DeepSeek-MoE 16B	16.4B	2.8B	64 (fine-grained)	6	2	细粒度专家
DeepSeek-V2	236B	21B	160	6	2	MLA + DeepSeek MoE
DeepSeek-V3	671B	37B	256	8	1	最大开源 MoE
Qwen1.5-MoE-A2.7B	14.3B	2.7B	60	4	4	共享专家在 HF 实现里被合并成一个大块

几个观察：

1. 专家数从 8 涨到 256+。早期 Mixtral 用 8 个大专家（每个相当于 7B 参数），新一代 DeepSeek 用 256 个小专家（每个相当于 2.5B）。细粒度更多专家 的设计被实证为优——更细的分工让专家学得更专。

2. Top-K 也从 2 涨到 6-8。专家变细之后，每次需要更多专家协作才能覆盖 token 的语义复杂度。

3. Shared experts 出现。DeepSeek 引入 shared experts—— 一两个专家「always on」（不通过 router，每个 token 都过它们）。共享专家学习「通用语言模式」，剩下的 routed experts 学习「专业化技能」——这种分层设计让 routed experts 的工作更聚焦。

4. 总 / 激活参数比从 4 涨到 18。Mixtral 8x7B 是 47B 总 / 13B 激活（约 3.6×），DeepSeek-V3 是 671B / 37B（约 18×）。稀疏率越来越高——每次只激活 5%-10% 的总参数。

flowchart TB
  subgraph "Mixtral 8x7B<br/>粗粒度 MoE"
    direction LR
    M_T[8 大专家<br/>每个 7B]
    M_K[Top-2]
    M_R[激活率 25%]
  end
  subgraph "DeepSeek-V3<br/>细粒度 MoE"
    direction LR
    D_T[256 小专家<br/>每个 2.5B]
    D_K[Top-8]
    D_R[激活率 ~3%]
    D_S[+ 1 shared expert<br/>所有 token 都过]
  end

12.6 MoE 的工程难点：All-to-All 通信

到这里讲的都是单 GPU 视角。MoE 真正的工程复杂度是多 GPU 训练——专家分布在不同 GPU 上时，token 必须从「自己的 GPU」流到「目标专家的 GPU」，然后流回来。这需要一种特殊的集合通信叫 All-to-All。

设想 8 张 GPU、64 个专家、每张 GPU 装 8 个专家。一个 token 经过 router 后被分配给某 K 个专家——这些专家可能分布在任何 GPU 上。流程：

flowchart LR
  GPU1[GPU 1<br/>有 8 个 token<br/>每个 token 选 2 个专家] --> AA[All-to-All]
  GPU2[GPU 2<br/>同上] --> AA
  GPU3[GPU 3<br/>同上] --> AA
  AA --> ROUTE[各 GPU 收到<br/>『发往自己 GPU 上某个专家』<br/>的 token]
  ROUTE --> EXP[GPU 各自计算自己的<br/>专家 FFN]
  EXP --> AA2[All-to-All 反向]
  AA2 --> RESULT[每张 GPU 拿回<br/>自己原本 token 的输出]

每次 forward / backward 都要做两次 All-to-All（一次发出、一次收回）。这是 MoE 训练比稠密模型多出的主要开销。

All-to-All 的特点：带宽瓶颈。每张 GPU 同时给所有其他 GPU 发数据、同时收所有其他 GPU 的数据——网络带宽必须很高才不被打瘪。

主流大规模 MoE 训练用的硬件：

InfiniBand 800 Gbps 或 NVLink Switch（NVL72）—— 同机内 GPU 之间用 NVLink，跨机用 InfiniBand
软件上：DeepSpeed-MoE、Tutel、FastMoE 都对 All-to-All 做了大量优化（管道化、压缩、合并通信）
DeepSeek-V3 把通信和计算重叠（在做 All-to-All 的同时计算其他 token），把通信 bubble 压到很小

这套工程栈是大模型训练真正的「门槛」——开源代码再多，没有工程团队你训不起来 MoE 大模型。

12.7 推理时的 MoE：批处理友好度

训练时 MoE 主要痛点是 All-to-All 通信。推理时痛点不同——主要是批处理友好度。

稠密模型推理时，一次 batch 进来 64 条 query，每条经过 FFN 时全部 4d 神经元都激活、batch 内所有 token 共享同一组权重——非常 GPU 友好。

MoE 推理时，64 条 query 的 token 在路由后被分散到不同专家，每个专家只处理一小撮 token——单个 GEMM 变小，GPU 利用率下降。

具体地，假设有 64 个专家、每个 token 选 Top-2、batch 里有 200 个 token：

平均每个专家收到 $200 \times 2 / 64 \approx 6$ 个 token
这 6 个 token 的 FFN 计算量太小，GEMM kernel 启动开销占主导
GPU 算力利用率从稠密模型的 50% 掉到 5-10%

解决方案：

更大 batch：把 batch 做到几千 token，每个专家平均收到 100+ 个，单 kernel 计算量回到合理范围。
Expert Parallelism + 专家 Group：把同 GPU 上的几个专家合并成一组，统一调度。
Speculative MoE（实验性）：用更小模型预测路由结果，提前调度。

vLLM、TensorRT-LLM、SGLang 都对 MoE 推理做了专门优化。当前 MoE 推理在 large batch 下能基本追平稠密模型的 GPU 利用率，但小 batch 仍然不友好——所以 MoE 模型部署给「在线低延迟」场景比稠密模型麻烦。

12.8 MoE vs 稠密：什么时候选 MoE？

把两者的取舍放一起：

维度	稠密	MoE
模型容量	受限于训练算力	容量大 10x（同算力）
训练复杂度	简单	复杂（负载均衡、All-to-All）
训练算力	高	同等模型「能力」下更低
显存占用	高	高（专家全部要在显存里）
推理 FLOPs	高	低（只激活 K 个专家）
推理 batch 友好	友好	大 batch 友好，小 batch 不友好
训练稳定性	好	差（需要专门技巧）
工程门槛	中	高

什么时候选 MoE？

目标是「最大化能力 / 单次推理 FLOPs 比」：MoE 用更少推理算力换更大容量。这是 OpenAI / Anthropic / Google / DeepSeek 等顶级实验室的选择。
训练算力充足、能容忍工程复杂度：MoE 训练成本不低，且对工程团队要求高。
场景是大 batch 服务：API 后端、批量推理。

什么时候不选？

小模型（< 10B）：稀疏化的边际收益小，不如直接稠密。
小 batch 在线服务：MoE 在小 batch 下硬件利用率低。
工程团队规模小：MoE 的训练和推理调优工作量都比稠密大很多。

主流开源生态目前的选择：普通规模（7B、13B、70B）大多用稠密；超大规模（200B+）几乎全是 MoE。这个分界线大致和「单 GPU 能放下一份模型」的资源边界吻合。

12.9 MoE 的训练稳定性

MoE 训练有几个特别容易掉的坑：

坑一：Router 的 logits 爆。路由器的输出经过 softmax，logits 太大会让 softmax 一面倒地选某个专家，剩下的几乎为零——选择多样性消失，模型退化。解决办法：

给 router 输出加 entropy regularization（鼓励熵不要太低）
用 z-loss 约束 logits 的尺度（让 log(sum(exp(logits))) 不要太大）

坑二：Dead experts。某个专家从初始化开始就没被选过，永远拿不到梯度更新，永远学不到东西。解决办法：

Auxiliary loss / DeepSeek bias 强制均匀使用
大初始化 + warmup（让 router 输出在初期更随机）

坑三：训练不稳定 spike。MoE 模型经常在训练中途 loss 突然飙升（spike）。原因是某些专家进入「重新分工」阶段——一个专家因为分布漂移，被频繁切换 token 来源，loss 高度震荡。解决办法：

训练 checkpoint 频繁，spike 时回滚
DeepSeek 用更小 lr + 更稳定的优化器配置
路由器单独用更小 lr

坑四：跨语言 / 跨模态的不平衡。如果训练数据是「英文 + 中文 + 代码」混合，模型可能让某些专家「专攻英文」、某些专攻中文——某个 batch 内英文比例高时，对应专家被打爆，剩下的空闲。解决办法：

在数据采样时按比例混合
用 packing 把不同语言的 token 打散到 batch 里

这些经验在 DeepSeek 系列论文里有大量讨论。读 DeepSeek-V2、V3 技术报告会让你看到「真正训出 MoE 大模型」的复杂度。

12.10 关于 MoE 的几个误解

误解一：MoE = 把模型拆成独立子模型

不对。所有专家共享同一个 attention 层、同一个 embedding、同一组归一化层——只有 FFN 部分被分成专家。整个模型仍然是「一个模型」，只是 FFN 部分稀疏激活。

误解二：每个专家学习一个特定领域

部分对。研究表明专家会自然分工——有些更擅长代码、有些更擅长中文、有些更擅长数学。但这种分工是模糊的、统计的，不是「这个专家专做 Python、那个专家专做 SQL」这种明确边界。Top-K 路由让每个 token 多路出击，分工不会完全清晰。

误解三：MoE 显存比稠密小

不对。所有专家的参数都要在显存里——MoE 的总参数甚至比同等能力的稠密模型大。MoE 节省的是计算量，不是显存量。这就是为什么部署 MoE 模型的硬件门槛和总参数量绑定（DeepSeek-V3 671B 推理仍然需要多卡）。

误解四：稀疏激活 = 推理加速 K/N 倍

不完全。理论上 K=2、N=64 时计算量是稠密的 2/64=3%，但实际上：

Router 本身有计算（虽然小）
专家分组后 GEMM 变小，GPU 利用率下降
All-to-All 或者 routing 操作本身有开销

实际加速大约是稠密模型的 2-5 倍——比理论值低，但仍然可观。

12.11 一个新视角：MoE 和模型记忆容量

我们在 5.2 节提到，FFN 是 Transformer 的「记忆容量」（key-value memory）。MoE 把 FFN 拆成 N 个专家，相当于把记忆容量拆成 N 个独立的「记忆库」——每个 token 根据内容在某些库里查。

这个视角让 MoE 的设计变得直观：

专家越多 + 越细粒度 → 记忆库越多越细，能存的「专门知识」越多
Top-K 越大 → 每次查询能跨多个库聚合，覆盖更复杂语义
Shared experts → 一个「公共图书馆」存通用知识，所有 token 都先过它

DeepSeek-V3 的 256 routed + 1 shared 设计恰好对应这个直觉：1 个公共图书馆 + 256 个专科图书馆，每个 token 从 256 个里挑 8 个最相关的查。这也解释了为什么 DeepSeek 的细粒度设计在大规模上更胜——记忆库更细，知识存得更整齐。

本章小结

MoE 解耦了模型容量和单次推理算力——通过条件计算让总参数量远超激活参数量。
MoE Layer = N 个独立 FFN（专家）+ 路由器（Top-K 选择）——结构上替换稠密 FFN，其余不变。
负载均衡是核心难点：少数专家被打爆、其他 dead 是常见 failure mode。Auxiliary loss 或 DeepSeek-style bias 是主流解决方案。
细粒度专家 + 共享专家是 2024 趋势：DeepSeek 256 routed + 1 shared 比 Mixtral 8 大专家更胜。
All-to-All 通信是分布式训练的瓶颈——需要 InfiniBand / NVLink + 专门通信库。
MoE 推理在小 batch 下不友好——专家分组让 GEMM 变小，GPU 利用率掉。生产部署需要大 batch。
MoE 不省显存只省计算——所有专家的参数都要在显存里。
训练稳定性是另一道墙：router logits 爆、dead experts、loss spike——需要专门工程经验。
何时选 MoE：需要超大容量但要控制推理算力、有大 batch 场景、有强工程团队。
何时不选：小模型（<10B）、低延迟在线服务、工程团队小。

下一章我们看长上下文之战——从 4K 到 1M，Transformer 是如何把上下文窗口推到这种量级的。这一章会用到第 4 章的位置编码、第 5 章的 attention 计算复杂度，作为综合应用。

第 12 章 Mixture of Experts：稀疏激活的工程账

12.1 一个直觉：稀疏神经网络

12.2 数学：MoE Layer 的形式

12.3 一个最小代码实现

12.4 关键挑战：负载均衡

Auxiliary Loss：硬性负载均衡

Auxiliary-Free Load Balancing

Token Dropping vs Expert Capacity

12.5 主流 MoE 设计对比

12.6 MoE 的工程难点：All-to-All 通信

12.7 推理时的 MoE：批处理友好度

12.8 MoE vs 稠密：什么时候选 MoE？

12.9 MoE 的训练稳定性

12.10 关于 MoE 的几个误解

12.11 一个新视角：MoE 和模型记忆容量

本章小结

延伸阅读