MOE模型小窥

Kyrie Chen 2025-02-18

Mixture of Experts (MoE) 是一种把“大参数量”和“低单次计算量”拆开的模型架构。它的基本想法并不复杂：模型里放入多个子网络作为专家（Experts），每个 token 只被路由到少数几个专家上计算，再把专家输出组合回来。这样一来，模型可以拥有很大的总参数量，但每次前向传播只激活其中一小部分参数。

对大语言模型来说，这个特性很有吸引力。模型规模通常是提升能力的重要因素，但稠密模型会让每个 token 都经过所有参数，训练和推理成本都很高。MoE 则通过稀疏激活，让模型在相近 FLOPs 下获得更大的容量。它的代价也很明确：所有专家参数仍然要常驻显存，路由会带来负载均衡问题，多机训练时还会引入大量 AllToAll 通信。

本文主要参考 Hugging Face 的《混合专家模型（MoE）详解》，结合自己对训练和推理系统的理解，对 MoE 做一个从原理到工程实践的梳理。

为什么是 MoE

MoE 最核心的收益是更高的预训练效率。给定固定计算预算时，一个更大的稀疏模型往往能比一个较小的稠密模型更快达到相近质量。直觉上说，MoE 把参数容量扩展到了更多专家上，但每个 token 只使用 Top-1 或 Top-2 个专家，因此训练时的计算量没有按总参数量线性增长。

MoE 的推理优势需要分开看。它的计算量可能低于同等总参数量的稠密模型，因为每个 token 只激活部分专家；但它的显存需求并不低，因为服务时通常仍要加载全部专家。例如 Mixtral 8x7B 的总参数量约 46.7B，但每个 token 只经过 2 个专家，同时注意力层等非专家部分是共享的，所以它的激活计算量明显小于一个 46.7B 稠密模型，却仍然需要接近 47B 参数级别的显存预算。

因此，MoE 更适合高吞吐、多卡或多机环境；如果只有少量显存、低并发推理，稠密模型反而更简单。

简史：从局部专家到 Transformer MoE

MoE 的概念最早可以追溯到 1991 年的 Adaptive Mixture of Local Experts。当时的思想更接近集成学习：多个局部专家分别处理输入空间中的不同区域，再由一个门控网络决定每个专家的权重。

2010 到 2015 年前后，有两条线索推动了后来的稀疏大模型：一条是把专家作为深层网络中的组件，而不是完整模型外部的集成器；另一条是条件计算，也就是让不同输入动态激活不同网络部分。二者结合后，就自然形成了“在大模型内部放入稀疏激活模块”的路线。

2017 年，Shazeer 等人在《Outrageously Large Neural Networks》中把稀疏 MoE 用到大规模 LSTM 翻译模型里，展示了可以在不按比例增加计算量的情况下扩展到很大的参数规模。后来 GShard、Switch Transformer、GLaM、ST-MoE 等工作把 MoE 推进到 Transformer 时代，并系统研究了并行训练、路由稳定性、负载均衡和微调泛化问题。

MoE 的基本结构

现代 Transformer MoE 通常不是把整个模型复制成多个专家，而是把 Transformer block 里的 FFN 层替换成 MoE 层。注意力层、LayerNorm、Embedding 等部分仍然共享，专家通常就是若干个独立的 FFN。

MoE 作为一种基于 Transformer 架构的模型，主要有两部分组成：

  • 稀疏 MoE 层: 这些层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干“专家”(例如 8 个)，每个专家本身是一个独立的神经网络。在实际应用中，这些专家通常是前馈网络 (FFN)，但它们也可以是更复杂的网络结构，甚至可以是 MoE 层本身，从而形成层级式的 MoE 结构。
  • 门控网络或路由: 这个部分用于决定哪些令牌 (token) 被发送到哪个专家。例如，在下图中，“More”这个令牌可能被发送到第二个专家，而“Parameters”这个令牌被发送到第一个专家。有时，一个令牌甚至可以被发送到多个专家。令牌的路由方式是 MoE 使用中的一个关键点，因为路由器由学习的参数组成，并且与网络的其他部分一同进行预训练。

稀疏性

稀疏性的概念采用了条件计算的思想。在传统的 Dense 模型中，所有的参数都会对所有输入数据进行处理。相比之下，稀疏性允许我们仅针对整个系统的某些特定部分执行计算。这意味着并非所有参数都会在处理每个输入时被激活或使用，而是根据输入的特定特征或需求，只有部分参数集合被调用和运行。

一个 MoE 层主要包括三部分：

  • 专家（Experts）：多个可学习的子网络，通常是 FFN。不同专家会在训练中逐渐处理不同分布的 token，但这种“专化”不是人工指定的。
  • 路由器（Router/Gating Network）：根据 token 的 hidden state 计算它应该发往哪些专家，常见方式是 Top-1 或 Top-2 路由。
  • 组合器（Combiner）：把一个或多个专家的输出按路由权重加权组合，再送回原来的 token 顺序。

可以把它写成下面这个形式：

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

这里 $E_i$ 表示第 $i$ 个专家，$G(x)_i$ 表示路由器给这个专家的权重。如果某个专家权重为 0，就不需要执行对应专家的计算，这就是 MoE 稀疏性的来源。

一个最朴素的路由器可以写成：

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

实际模型中通常还会加上 Top-K 截断、噪声或其他正则项。Shazeer 等人的 Sparsely-Gated MoE 使用过 Noisy Top-K Gating，其流程可以概括为：先给 router logits 加噪声，再只保留最大的 K 个专家，最后做 Softmax。

  1. 添加一些噪声
\[H(x)_i=(x \cdot W_g)_i + \text{Standard Normal}() \cdot Softplus((x \cdot W_{noise})_i)\]
  1. 选择保留前 K 个值
\[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}\]
  1. 应用 Softmax 函数
\[G(x)=Softmax(KeepTopK(H(x), k))\]

当时用较低的 $k$值（例如 1 或 2）时，我们可以比激活多个专家时更快地进行训练和推理。但为了让门控学会如何进行有效的路由选择，因此至少需要选择两个专家。

路由、容量与负载均衡

MoE 不是“专家越多越好”那么简单。它真正难的地方在于：如何让 token 均匀、稳定、高效地分配到专家上。

如果 router 总是偏爱少数专家，就会出现热门专家过载。热门专家因为拿到更多样本而训练得更快，之后又更容易被 router 选择，形成自我强化。冷门专家则训练不足，参数容量被浪费。为了解决这个问题，MoE 训练通常会加入辅助损失（auxiliary loss），鼓励路由概率和实际 token 分布更加均匀。

另一个关键概念是专家容量（expert capacity）。由于深度学习编译和分布式通信通常要求张量形状相对固定，我们不能让每个专家接收任意数量的 token。因此需要为每个专家设置容量上限：

\[Expert\ Capacity = \frac{tokens\ per\ batch}{number\ of\ experts} \times capacity\ factor\]

其中 capacity factor 越大，越不容易溢出 token，但显存和通信成本也越高；capacity factor 越小，效率更高，但可能丢弃或绕过更多 token。Switch Transformer 证明了较小容量因子也可以工作得不错，但实际取值仍要结合 batch size、专家数、网络拓扑和是否 dropless 训练来调。

常见路由策略的取舍大致如下：

  • Top-1：每个 token 只去一个专家，通信和计算更省，代表工作是 Switch Transformer。
  • Top-2：每个 token 去两个专家，表达力和训练稳定性更强，但通信和专家计算更贵，GShard、GLaM、Mixtral 都采用过类似策略。
  • Noisy Top-K：通过噪声增加早期探索，避免过早塌缩到少数专家。
  • Dropping/Dropless：前者允许溢出 token 被丢弃或走残差，后者尽量不丢 token，但系统成本更高。

ST-MoE 还提出 Router z-loss 来约束 router logits 的尺度。它的目标不是直接均衡专家，而是防止路由器输出过大的 logits，降低指数运算和低精度训练中的数值不稳定。

训练与微调

MoE 的预训练效率通常很好，但微调不一定比稠密模型省心。Hugging Face 的综述里也强调，稀疏模型在微调阶段更容易出现过拟合，尤其是在小数据任务上。原因之一是每个 token 只更新少量专家，专家之间的数据分布又不完全一致，微调时很容易把路由和专家同时推向不稳定区域。

实践上可以关注几类手段：

  • 更强正则：稀疏层可以使用更高 dropout，或保留 aux loss 来缓解过拟合与负载失衡。
  • 学习率和 batch size：一些研究观察到，稀疏模型微调可能更适合较小 batch 和较高学习率，但需要按任务验证。
  • 冻结策略：只更新非 MoE 层有时能接近全量微调效果，同时降低训练成本；只冻结非专家层通常效果更差，因为专家层承载了大量参数容量。
  • 指令微调：MoE 在多任务和指令微调中可能收益更大，Flan-MoE 相关结果显示指令调优能明显改善 MoE 的下游表现。

训练稳定性可以按这个顺序排查：先看各专家 token 数和溢出率，再看 router logits 是否过大，再看 loss 曲线和梯度范数，最后再进入通信和 kernel 级 profile。

并行与系统

MoE 的性能上限常由通信与放置决定，而不是单纯由算力决定。一次 MoE 前向大致包含四步：

  1. 路由与分桶：Router 计算每个 token 的 Top-K 专家，按专家 ID 把 token 重新分桶。
  2. AllToAll：各 GPU 交换属于彼此专家的 token，这一步往往是瓶颈。
  3. 专家计算：每个 GPU 上的本地专家并行执行 FFN；如果叠加 TP，则专家内部矩阵乘也会继续切分。
  4. 组合与回传：专家输出按原 token 顺序拼回，并由 Combiner 聚合。

MoE 常见并行方式包括：

  • 数据并行（DP）：不同设备处理不同 batch，模型副本相同。
  • 张量并行（TP）：把单个大矩阵或算子切到多个设备上。
  • 流水并行（PP）：按层切分模型。
  • 专家并行（EP）：把不同专家放到不同设备上，token 根据路由跨设备发送。

在实际工程里，常见组合是 EP 负责专家横向扩展，TP/PP 解决单层或整模型过大问题，DP 放在外层扩吞吐。单机多卡依赖 NVLink/NVSwitch 会舒服很多；多机环境下跨节点 AllToAll 昂贵，需要更谨慎地设置专家分组、micro-batch、capacity factor 和通信计算重叠。

容量与放置的一些经验：

  • Capacity factor 过小会导致溢出，过大则放大显存与 AllToAll 压力。
  • 固定专家到设备的映射，减少训练中不必要的迁移。
  • 热门专家尽量放在互联更好的 GPU 或通信域内。
  • 用 Nsight Systems/Compute、NCCL_DEBUG=INFO 和 nccl-tests 先确认瓶颈，再决定是否分块、分桶或调整专家数。

提高容量因子 (Capacity Factor, CF) 可以增强模型的性能，但这也意味着更高的通信成本和对保存激活值的显存的需求。在设备通信带宽有限的情况下，选择较小的容量因子可能是更佳的策略。一个合理的初始设置是采用 Top-2 路由、1.25 的容量因子，同时每个节点配置一个专家。在评估性能时，应根据需要调整容量因子，以在设备间的通信成本和计算成本之间找到一个平衡点。

一句话记忆：把专家放在带宽好的地方，把通信藏到计算背后。

推理与部署

MoE 推理的难点不只是“每次只算几个专家”。服务系统还要处理动态路由带来的 batch 不规则、专家冷热不均、跨卡通信和显存常驻问题。

本地部署时，MoE 最大的问题往往是显存：即使每个 token 只激活部分专家，所有专家权重也通常需要加载。云端高吞吐服务则更关心吞吐和通信效率：连续批处理（continuous batching）会把不同请求的 token 混在一起，router 又会把 token 分发到不同专家，这要求调度系统能高效地做重排、AllToAll 和回填。

几类常见的部署优化方向：

  • 量化：对专家权重做 MoE-aware quantization，减少显存占用，但要特别关注 router 和激活分布的稳定性。
  • 蒸馏：把稀疏 MoE 蒸馏回较小稠密模型，牺牲部分 MoE 收益换取简单部署。
  • 专家合并：将多个专家权重聚合或裁剪，减少推理时需要加载的专家数量。
  • 任务级路由：把句子、任务或用户请求路由到特定专家子集，降低 token 级路由的系统复杂度。

现有 MoE 模型汇总

注：下表中带“需核实”的条目表示公开资料不足或厂商未披露完整细节，仅作参考。

模型 发布 规模 训练 备注
Deepseek DeepSeek 于 2024 年 12 月 10 日发布并持续更新 16B（激活 2.4B）、236B（激活 22B） 未披露 当前最优秀的 MoE 系列大模型之一
Qwen2.5-54B-A14B Alibaba 于 2024 年 5 月发布 54B（激活 14B） 包含大规模文本与代码数据，采用 MoE 与稠密层结合 针对聊天生成任务优化，HF 地址中写的是 57B，需核实是否为笔误
Mixtral Mistral AI 于 2024 年 1 月发布 46.7B（8x7B） 稀疏 MoE 架构 8 个专家，每次选择 2 个
Arctic Snowflake 于 2024 年 4 月发布 480B（128x3.66B，激活 17B） 动态数据课程，包含代码和文本数据，使用约 1000 张 GPU 训练约 3 周 开源模型，强调企业任务与代码能力
DBRX Databricks 于 2024 年 3 月发布 132B（36B 活跃参数） 12T 文本和代码数据，使用 3072 张 H100 训练约 3 个月 16 个专家，每次选择 4 个，编程任务表现突出
Grok-1 xAI 于 2023 年 11 月发布 推测数十亿参数量（未公开） 可能包含 X 平台数据，细节未公开 与 X 平台深度整合，需核实
Grok-2 xAI 于 2024 年 8 月发布 未知 据称包含多模态数据与改进 MoE 架构 官方未披露技术细节，需核实
Grok-3 xAI 于 2025 年 2 月 18 日发布 未知 据称使用大规模 GPU 训练 官方未披露技术细节，需核实
OLMoE AllenAI 于 2024 年 9 月 24 日发布 7B（1B 激活参数） 英文语料约 5T tokens，使用 256 个 H100 训练约 10 天 同等规模下表现较好的开源 MoE 模型

评测与排错

评测 MoE 不能只看任务分数，还要看效率和稳定性指标。一个 MoE 模型如果 benchmark 分数不错，但专家负载严重倾斜、溢出率很高、AllToAll 占据大部分时间，实际训练和服务都可能很痛苦。

建议至少跟踪这些指标：

  • 任务指标：MMLU、GSM、HumanEval、MT-Bench 或业务自己的离线评测。
  • 效率指标：tokens/s、MFU、通信占比、AllToAll 时间、显存峰值。
  • 路由指标：每个专家 token 数、路由熵、Top-K 分布、溢出率、被丢弃 token 比例。
  • 稳定性指标：router logits 范围、aux loss、z-loss、梯度范数、loss spike。

常见问题可以按下面思路定位：

  • 热门专家过载：先调 aux loss、router 温度、噪声或 capacity factor。
  • 溢出严重：增大 capacity factor，或检查 batch size 是否太小导致随机波动过大。
  • 训练发散：重点看 router logits、混合精度策略、z-loss 和梯度裁剪。
  • 通信瓶颈：看 AllToAll 时间线，调整专家放置、通信域、分桶大小和通信计算重叠。
  • 推理尾延迟高：检查是否存在少数专家过热、跨节点路由过多或 batch 重排成本过高。

结语：从原理到实践

MoE 的核心是稀疏激活：通过 router 将每个 token 路由到少量专家，在不按比例增加 FLOPs 的前提下扩大模型容量。围绕这一点，Top-K、capacity factor、aux loss、z-loss 和噪声调度决定了训练是否稳定，也决定了专家能否真正被用起来。

但 MoE 最终能不能跑好，很大程度上取决于系统。一次前向或反向的主瓶颈常常不是 FFN 矩阵乘，而是跨设备 token 交换。工程上需要用 EP、TP、PP、DP 的组合解决规模问题，再用拓扑感知放置、分桶、分块和通信计算重叠把 AllToAll 的代价压下去。

一句话总结：MoE 让参数规模和单次计算量部分解耦，但要兑现它的潜力，既要在路由和正则上管好专家，也要在并行和拓扑上摆好专家。

参考与延伸阅读

  • Hugging Face: 混合专家模型（MoE）详解
  • Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
  • GShard, Switch Transformer, GLaM, ST-MoE, MegaBlocks
  • Mixtral 8x7B、OLMoE、DeepSpeed-MoE、Megatron-Core MoE、Fairseq-MoE