Lec16 Systems for MoE Models

这一讲聚焦 Mixture-of-Experts (MoE) 模型背后的系统挑战：稀疏专家架构如何在不按比例增加计算量的情况下扩展模型容量；专家并行引入了怎样的通信模式；以及 GShard、DeepSpeed-MoE、DeepSeek V3 等系统在大规模下是如何解决路由、负载均衡和推理效率问题的。

动机：为什么要做稀疏模型？

稠密模型的扩展成本非常高昂。训练一个 175B 的 OPT 模型需要 992 张 A100 训练 56 天；540B 的 PaLM 则需要 6144 片 TPU v4 训练 57 天。计算成本大致与模型规模呈线性增长，因为每一个参数都参与每一次前向计算。

Model Model Size Hardware Days to Train
Megatron-LM GPT-2 8.3B 512 V100 GPU 9.2 days
OPT 175B 992 A100 GPU 56 days
MT-NLG 530B 2200 A100 GPU 60 days
PaLM 540B 6144 TPU v4 57 days

稀疏模型提供了另一条扩展路径：在保持每个 token 计算量恒定的同时增加总参数量。Mixture-of-Experts 是目前 LLM 中最主流的稀疏架构。MoE 的预训练速度远快于同等质量的稠密模型，而且由于每个 token 只激活一部分参数，其推理速度也快于总参数量相同的稠密模型。

Transformer Mixture-of-Experts 架构

Transformer MoE 的核心思想非常简单：把每个 Transformer block 中单一的稠密 FFN 替换成多个小的 expert FFN，再加上一个门控网络（router），由它为每个输入 token 选择要激活哪个（或哪些）专家。

这与经典的 Mixture-of-Experts 学习算法不同——后者学习的是多个预测模型的加权平均。在 Transformer MoE 中，专家本身就是模型内部的子网络，router 为每个 token 做硬选择或软选择。

Switch Transformer

Switch Transformer（Fedus 等人，JMLR 2022）是经典的 MoE 架构。每个 token 会根据 router 的输出被路由到恰好一个 FFN 专家：

  • 门控网络计算一个关于专家的概率分布：

\[G_\sigma(x) = \text{Softmax}(x \cdot W_g)\]

  • 输出是专家输出的加权组合：

\[y = \sum_{i=1}^{n} G(x)_i E_i(x)\]

  • Top-k gating 只保留 top-\(k\) 的专家得分，其余全部置为 \(-\infty\)：

\[\text{KeepTopK}(v, k)_i = \begin{cases} v_i & \text{if } v_i \text{ is in the top } k \text{ elements of } v \\ -\infty & \text{otherwise} \end{cases}\]

\[G(x) = \text{Softmax}(\text{KeepTopK}(H(x), k))\]

Switch Transformer 具体采用 \(k=1\)：每个 token 只走一个专家。相比 top-2 设计，这简化了路由并降低了通信开销。

共享专家 vs. 路由专家

DeepSpeed-MoE 提出、DeepSeek 延续并采用的一种改进是共享专家 + 路由专家的设计：

Expert Type Behavior Purpose
Shared Expert Always activated for every token Captures common knowledge shared across all inputs
Routed Expert Conditionally activated by the router Captures token-specific, specialized knowledge

共享专家相当于每个 token 都会经过的固定 FFN，而路由专家提供的是条件性的额外容量。最终输出是共享专家的输出加上加权后的路由专家输出之和：

\[\text{output} = \text{SharedFFN}(h) + \sum_{i \in \text{TopK}} g_i \cdot \text{Expert}_i(h)\]

其中路由权重由 \(\text{Softmax}(\text{TopK}(h_t \cdot W))\) 计算得到。

专家到底学到了什么

Encoder 专家往往会专注在某些 token 组或浅层概念上：标点、专有名词、哨兵 token、动词、视觉描述、计数与数字等。Decoder 专家的专业化程度要弱得多。在多语言场景下，专家并不会按语言聚类——负载均衡和 token 路由机制会阻止这种聚集。

激活的专家在不同层之间会变化

不同 token 在不同层会激活不同的专家。token “red” 在第 1 层可能路由到 Expert 3，在第 2 层则路由到 Expert 1；而 “fox” 在这两层可能分别走到 Expert 2 和 Expert N。路由模式是动态的，并通过端到端训练学出来。

MoE 的参数量

MoE 模型的总参数量并不像朴素乘法那样多。Mixtral 8x7B 的参数量是 47B 而不是 56B，因为只有 FFN 层被复制成专家，attention、embedding 和 layer norm 在所有专家之间是共享的。

过拟合风险

MoE 模型在小数据集上容易过拟合。由于参数量远多于同级的稠密模型，稀疏模型可能会在小训练集上直接记住样本，而稠密模型反而泛化得更好。在较大规模任务上 MoE 表现良好。这在把 MoE 模型微调到下游任务时是一个重要的考量。

MoE 的训练与推理

专家并行（GShard）

专家并行（Lepikhin 等人，ICLR 2021）是分布 MoE 模型的标准做法：

  • 每个专家只存在于一台 worker 设备上 —— 专家被切到不同的 GPU 上
  • 其他所有组件（attention、embedding、layer norm）在每台设备上都复制一份
  • 使用 all-to-all 通信 把 token 派发到正确的专家设备，再把结果收集回来

前向过程分为四步：

  1. Route：每台设备运行本地的 router，判断每个 token 应当去哪个专家
  2. All-to-All Dispatch：把 token 发送到承载其被分配专家的设备
  3. Expert Compute：每台设备在接收到的 token 上运行本地的 expert FFN
  4. All-to-All Combine：把专家输出发回原始设备

GShard 还引入了 interleaving：MoE 层与标准稠密 FFN 层交替（每隔一层使用 MoE）。这样在稀疏路由带来的容量提升和通信开销之间取得了平衡。

MoE 中 token 的计算路径

在基于专家并行的推理过程中，每个 token 的计算路径会在 MoE 层处跨越设备边界。Token 先在它的「主」设备上做 embedding 和 attention，随后被派发到远程设备完成专家计算，再回到原设备。每台设备维护自己 attention 层的 KV-cache，但每一个 MoE 层上的专家计算都需要跨设备通信。

MoE 训练中的负载均衡

如果不加干预，router 可能会坍缩到把所有 token 都发给少数几个热门专家，导致大部分专家训练不充分。负载均衡损失正是为了防止这一点：

\[L_{ExpBal} = \alpha_1 M \sum_{i=1}^{\text{\#experts}} f_i P_i\]

其中：

  • \(f_i = \frac{\text{\#tokens to expert } i}{\text{\#tokens}}\) 是路由到专家 \(i\) 的 token 占比
  • \(P_i = \frac{1}{\text{\#tokens}} \sum_{t=1}^{\text{\#tokens}} s_{i,t}\) 是专家 \(i\) 的平均路由权重
  • \(M\) 是专家数量
  • \(\alpha_1\) 是控制均衡损失强度的超参数

乘积 \(f_i \cdot P_i\) 会惩罚那些同时拥有高 token 占比和高路由概率的专家。该损失通过 \(P_i\)（路由权重）是可微的，从而促使 router 更均匀地分发 token。

MoE 推理的挑战

MoE 推理性能主要由三个因素决定：

  • 总体模型大小 —— 决定整体内存占用
  • 激活专家的数量 —— 决定每个 token 的计算量
  • 总体内存带宽 —— 由于每次只有一部分专家被激活，推理往往是内存带宽瓶颈

默认实现会把所有专家都放在 GPU 显存中，即便每个 token 只激活少数几个专家，也需要很大的显存容量。

MoE 推理优化

MoE 推理系统设计的目标是：最小化每台设备的关键数据路径，并最大化聚合内存带宽。关键优化策略包括：

  • 把 critical data path 相同的 token 分组并一起路由，减少每台设备的数据访问量，获得最大聚合带宽
  • 优化通信调度，与并行策略协同
  • 优化 Transformer 和 MoE 相关的 kernel，提升单设备性能

专家并行与张量并行（DeepSpeed-MoE）

DeepSpeed-MoE（Rajbhandari 等人，2022）对 MoE 模型组合了多种并行策略：

  • 专家并行 / Expert Slicing：把分配到同一组专家的所有 token 放在同一个关键数据路径下，并把具有不同关键路径的 token 组在不同设备上并行处理
  • 张量并行 / Tensor Slicing：把非专家参数（attention）切分到多台设备上，通常在节点内部切分
  • 数据并行：把上述配置再复制到更多设备组上

这种多维并行方式对专家参数和非专家参数采用不同的处理方式，让每种参数都用上最合适的并行策略。

MoE Kernel 优化

MoE 专用的 kernel 优化包括：

  • 将 gating 函数融合到单个 kernel 中 —— 避免 softmax、top-k 和 dispatch 等操作各自启动 kernel
  • 稠密的 token-to-expert 映射表 —— 支持高效的批处理计算

这些优化可以将 MoE kernel 相关延迟降低 6 倍以上。

All-to-All 通信优化

专家并行需要在所有专家并行设备之间做 all-to-all 通信，延迟会随着设备数线性增长。两个关键优化：

  • 层级化 all-to-all 通信：先在节点内通信，再做跨节点通信，从而减少通信跳数
  • 与并行策略协同的通信：根据模型的并行策略调度通信，降低整体开销

MoE 训练与推理结果

DeepSpeed-MoE 表明，一个 1.3B+MoE-128 模型（总参数量 52B）可以达到和 6.7B 稠密模型相同的验证损失，同时训练吞吐提升 5 倍（128 张 A100 上 372 vs. 70 samples/sec）。

在推理方面，DeepSpeed 在所有 GPU 规模（8、16、32、64 张 GPU）下相比 PyTorch baseline 都实现了显著更低的延迟和更高的吞吐，吞吐量几乎随 GPU 数量线性扩展。

DeepSeek MoE

细粒度专家设计

DeepSeek MoE（Dai 等人，2024）引入了细粒度专家：把原始 FFN 切分成 \(k\) 个更小的专家，总共得到 \(kN\) 个专家（\(N\) 是原先的专家数）。更细的粒度允许更灵活的路由——模型每个 token 激活 \(kM\) 个专家（\(M\) 是基线激活数），每个专家都比标准 MoE 设计更小。

整个架构把共享专家和路由专家结合起来，并对激活的路由专家做 top-k 加权平均。

DeepSeek V3 MoE (670B)

DeepSeek V3 是一个 670B 参数的 MoE 模型，配置如下：

Parameter Value
Vocabulary 129,280
Hidden dimension 7,168
Number of layers 61
Dense layers (lowest) 3
Number of attention heads 128
FFN intermediate dimension 18,432
MoE dimension 2,048
Shared experts 1
Routed experts 256
Activated experts per token 8
Expert groups 8
Limited groups 4

架构采用 Multi-head Latent Attention (MLA) 代替标准的多头注意力。每个 MoE 层有 1 个共享专家 FFN 处理每一个 token，外加 256 个路由专家，每个 token 激活其中 8 个。路由使用 \(\text{TopK}(\text{Softmax}(h_t \cdot W), 8)\)。

每个专家使用 SwiGLU 激活：

\[\text{FFN}_{SwiGLU}(x) = (\text{Swish}(x \cdot W_1) \odot (x \cdot W_2)) \cdot W_3\]

前 3 层是稠密层（标准 FFN），其余 58 层使用 MoE。这种设计让最底层先构建通用表示，再引入稀疏路由。

DeepSeek MoE 的负载均衡

DeepSeek V3 使用了两个层级的负载均衡损失：

专家级均衡损失（避免路由坍缩）：

\[L_{ExpBal} = \alpha_1 \sum_{i=1}^{\text{\#experts}} f_i P_i\]

其中 \(f_i = \frac{\text{\#experts}}{\text{\#activated\_experts}} \cdot \frac{\text{\#tokens to expert } i}{\text{\#tokens}}\)，\(P_i = \frac{1}{\text{\#tokens}} \sum_{t=1}^{\text{\#tokens}} s_{i,t}\)（平均路由权重）。

设备级均衡损失（在设备之间均衡计算量）：

\[L_{DevBal} = \alpha_2 \sum_{j=1}^{\text{\#groups}} f_j P_j\]

其中 \(f_j\) 是组 \(j\) 中专家的平均 \(f\)，\(P_j\) 是组 \(j\) 中专家 \(P\) 的求和。

两级设计同时保证了单个专家的利用率和跨设备计算的均衡。

DeepSeek 为 MoE 开源的库

DeepSeek 开源了两个加速 MoE 的库：

  • DeepEP：专为 MoE 和专家并行优化的通信库，优化 all-to-all 的 dispatch 和 combine 操作
  • EPLB（Expert Parallelism Load Balancer）：负责在专家并行设备之间做动态负载均衡

DeepSpeed MoE 代码示例

DeepSpeed 提供了简单的 API 用于添加 MoE 层：

1
2
3
4
5
6
7
8
9
10
11
12
13
import torch
import deepspeed
import deepspeed.utils.groups as groups
from deepspeed.moe.layer import MoE

WORLD_SIZE = 4
EP_WORLD_SIZE = 2
EXPERTS = 8

fc3 = torch.nn.Linear(84, 84)
fc3 = MoE(hidden_size=84, expert=self.fc3,
           num_experts=EXPERTS, ep_size=EP_WORLD_SIZE, k=1)
fc4 = torch.nn.Linear(84, 10)

MoE 类包装现有的 nn.Module（专家模板），负责复制、门控以及专家并行通信。关键参数包括：

  • hidden_size：模型隐藏维度（输入输出维度）
  • expert：基础专家模块（例如一个 MLP 或 torch.nn.Linear）
  • num_experts：每层的总专家数
  • ep_size：专家并行组中的 rank 数
  • k：top-k gating 取值（1 或 2）
  • capacity_factor：控制训练时的专家容量
  • drop_tokens：是否丢弃超出专家容量的 token

在内部，Experts 类为每个本地专家深拷贝一份专家模块，给专家参数设置 param.allreduce = False（因为它们走的是专家并行而不是数据并行），并将输入按专家做切分：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class Experts(nn.Module):
    def __init__(self, expert, num_local_experts=1, expert_group_name=None):
        super().__init__()
        self.deepspeed_experts = nn.ModuleList(
            [copy.deepcopy(expert) for _ in range(num_local_experts)])
        for expert in self.deepspeed_experts:
            for param in expert.parameters():
                param.allreduce = False
                param.group_name = expert_group_name

    def forward(self, inputs):
        chunks = inputs.chunk(self.num_local_experts, dim=1)
        expert_outputs = []
        for chunk, expert in zip(chunks, self.deepspeed_experts):
            expert_outputs += [expert(chunk)]
        return torch.cat(expert_outputs, dim=1)

Lec16 要点回顾

  • MoE 扩展容量但不扩展计算量：通过只激活每个 token 所需的少量专家，MoE 模型可以拥有万亿级参数，同时保持与一个小得多的稠密模型相当的 per-token FLOPS
  • 专家并行是最自然的分布策略：把专家分散到不同设备，其他部分全部复制，再用 all-to-all 派发 token——但在大规模场景下，all-to-all 会成为瓶颈
  • 负载均衡至关重要：没有辅助损失的话，router 会坍缩到少数几个热门专家，浪费容量并造成设备间计算不均衡
  • DeepSeek V3 把细粒度设计推向极致：256 个路由专家里只激活 8 个、用共享专家承载通用知识、两级均衡损失、以及 670B 模型中的 MLA attention
  • 多维并行（EP + TP + DP） 在大规模下是必需的，每种策略都被应用到最适合它的那部分模型上

Final Summary

Topic Key Idea
Sparse vs. Dense MoE increases model capacity without proportionally increasing per-token compute
Switch Transformer routes each token to one expert via top-1 gating with softmax over expert scores
Shared vs. Routed Experts shared experts capture common knowledge; routed experts specialize per token
Expert Parallelism (GShard) split experts across devices, replicate non-expert params, use all-to-all communication
Load Balancing Loss auxiliary loss penalizing uneven token distribution across experts to prevent routing collapse
MoE Inference memory-bandwidth-bound; optimize by grouping tokens, fusing kernels, hierarchical all-to-all
DeepSpeed-MoE combines expert parallelism, tensor parallelism, and data parallelism with optimized kernels
DeepSeek V3 (670B) fine-grained experts (256 routed, 8 activated), shared expert, two-level balance loss, SwiGLU
DeepEP / EPLB DeepSeek libraries for optimized MoE communication and dynamic expert load balancing

核心要点：MoE 是在稠密模型计算预算之外继续扩展 LLM 的主流手段。系统层面的挑战从并行化一段统一的计算，转变为在多设备之间编排动态的、与 token 相关的路由——这需要借助 all-to-all 通信的专家并行、用于防止路由坍缩的负载均衡机制，以及能够分别处理专家与非专家参数的多维并行策略。

References

  1. Fedus et al. “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.” JMLR 2022.
  2. Lepikhin et al. “GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.” ICLR 2021.
  3. Rajbhandari et al. “DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale.” 2022.
  4. Dai, D. et al. “DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models.” 2024.
  5. DeepSeek-V3 inference code: https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py
  6. DeepEP: https://github.com/deepseek-ai/DeepEP
  7. EPLB: https://github.com/deepseek-ai/EPLB

本文基于 CMU 11-868 LLM Systems 的课程材料整理，授课教师 Lei Li（Systems for MoE Models, Lec16）。