Hugging Face 图解 Mixture of Experts（MoE）原理、训练与推理权衡

这篇博客文章有第二个版本（2026年2月），在其中我们介绍了 transformers 库如何围绕 MoE 构建，使它们成为该库和 Hub 的“一等公民”。以下是文章链接：Transformers 中的混合专家模型（MoE）

随着 Mixtral 8x7B 的发布（公告、模型卡），一类 transformer 已成为开放 AI 社区中最热门的话题：混合专家模型，简称 MoE。在这篇博客文章中，我们将探讨 MoE 的构建模块、它们如何训练，以及在推理服务时需要考虑的权衡。

让我们开始吧！

目录

• 什么是混合专家模型？
• MoE 简史
• 什么是稀疏性？
• MoE 的 token 负载均衡
• MoE 与 Transformer
• Switch Transformer
• 使用路由器 Z-loss 稳定训练
• 专家学到了什么？
• 扩展专家数量如何影响预训练？
• 微调 MoE
• 何时使用稀疏 MoE 而非稠密模型？
• Making MoEs go brrr
  • 专家并行
  • 容量因子与通信成本
  • 服务技术
  • 高效训练
• 开源 MoE
• 令人兴奋的研究方向
• 一些资源

内容提要

MoE：

• 与稠密模型相比，预训练速度快得多
• 与参数量相同的模型相比，推理速度更快
• 由于所有专家都加载在内存中，因此需要高显存
• 面临许多微调方面的挑战，但近期工作表明MoE指令微调前景可期

让我们开始吧！

什么是混合专家模型（MoE）？

模型规模是提升模型质量最重要的维度之一。在固定的计算预算下，用更少的步数训练更大的模型，比用更多的步数训练更小的模型效果更好。

混合专家模型使模型能够在预训练时使用远少的计算量，这意味着你可以在与稠密模型相同的计算预算下，大幅扩展模型或数据集规模。特别是，MoE 模型在预训练期间应该能比其稠密对应模型更快地达到相同的质量。

那么，MoE 到底是什么？在 transformer 模型的语境下，MoE 由两个主要元素组成：

• 稀疏 MoE 层用于替代稠密前馈网络（FFN）层。MoE 层拥有一定数量的“专家”（例如 8 个），每个专家都是一个神经网络。在实践中，专家是 FFN，但它们也可以是更复杂的网络，甚至本身就是一个 MoE，从而形成层级式 MoE！
• 一个门控网络或路由器，用于决定哪些 token 被发送到哪个专家。例如，在下图中，token“More”被发送到第二个专家，而 token“Parameters”被发送到第一个网络。正如我们稍后将探讨的，我们可以将一个 token 发送到多个专家。如何将 token 路由到专家是使用 MoE 时的重大决策之一——路由器由可学习参数组成，并与网络的其余部分同时预训练。

那么，总结一下，在 MoE 中，我们将 transformer 模型的每个 FFN 层替换为 MoE 层，该层由一个门网络和一定数量的专家组成。

尽管与稠密模型相比，MoE 提供了高效预训练和更快推理等优势，但它们也伴随着挑战：

• 训练： MoE 能够实现显著更计算高效的预训练，但它们在微调过程中历来难以泛化，导致过拟合。
• 推理： 尽管 MoE 可能有许多参数，但在推理过程中只使用其中一部分。与具有相同参数数量的稠密模型相比，这导致推理速度快得多。然而，所有参数都需要加载到 RAM 中，因此内存需求很高。例如，对于像 Mixtral 8x7B 这样的 MoE，我们需要足够的 VRAM 来容纳一个稠密的 47B 参数模型。为什么是 47B 参数而不是 8 x 7B = 56B？这是因为在 MoE 模型中，只有 FFN 层被视为独立的专家，而模型的其余参数是共享的。同时，假设每个 token 只使用两个专家，推理速度（FLOPs）就像使用一个 12B 模型（而不是 14B 模型），因为它计算 2x7B 矩阵乘法，但有一些层是共享的（稍后会详细说明）。

现在我们对 MoE 有了一个大致的了解，让我们来看看导致其发明的研究进展。

MoE 简史

MoE 的根源来自 1991 年的论文 Adaptive Mixture of Local Experts。这个想法类似于集成方法，是为一个由独立网络组成的系统提供监督过程，每个网络处理训练案例的不同子集。每个独立的网络或专家，专门处理输入空间的不同区域。如何选择专家？一个门网络决定每个专家的权重。在训练过程中，专家和门网络都会被训练。

在 2010-2015 年间，两个不同的研究领域为后来的 MoE 进步做出了贡献：

• 专家作为组件：在传统的 MoE 设置中，整个系统包括一个门网络和多个专家。作为整个模型的 MoE 已在 SVM、高斯过程和其他方法中被探索。Eigen、Ranzato 和 Ilya 的工作探索了将 MoE 作为更深网络的组件。这使得 MoE 可以作为多层网络中的层，使模型能够同时既大又高效。
• 条件计算：传统网络将所有输入数据通过每一层处理。在这一时期，Yoshua Bengio 研究了基于输入 token 动态激活或停用组件的方法。

这些工作导致了在 NLP 背景下探索专家混合。具体来说，Shazeer 等人（2017 年，“等人”包括 Geoffrey Hinton 和 Jeff Dean，Google 的 Chuck Norris）通过引入稀疏性，将这个想法扩展到 137B LSTM（当时事实上的 NLP 架构，由 Schmidhuber 创建），允许即使在大规模下也能保持非常快的推理速度。这项工作专注于翻译，但面临许多挑战，如高通信成本和训练不稳定性。

MoE 使得训练数万亿参数的模型成为可能，例如开源的 1.6T 参数 Switch Transformers 等。MoE 也已在计算机视觉领域得到探索，但本博客文章将聚焦于 NLP 领域。

什么是稀疏性？

稀疏性利用了条件计算的思想。在稠密模型中，所有参数都用于所有输入，而稀疏性允许我们只运行整个系统的一部分。

让我们更深入地了解 Shazeer 对用于翻译的 MoE 的探索。条件计算的思想（网络的一部分在每个样本的基础上激活）允许人们在不增加计算量的情况下扩展模型规模，因此，这导致每个 MoE 层中使用了数千个专家。

这种设置带来了一些挑战。例如，尽管大批量通常对性能更好，但在 MoE 中，随着数据流经活跃的专家，批量大小实际上被减小了。例如，如果我们的批量输入包含 10 个 token，五个 token 可能最终进入一个专家，而另外五个 token 可能最终进入五个不同的专家，导致批量大小不均匀和利用率不足。下面的让 MoE 跑得飞快部分将讨论其他挑战和解决方案。

我们如何解决这个问题？一个学习到的门控网络（G）决定将输入的一部分发送给哪些专家（E）：

y=∑i=1nG(x)iEi(x) y = \sum_{i=1}^{n} G(x)_i E_i(x)

在这种设置中，所有专家都会对所有输入运行——这是一种加权乘法。但是，如果 G 为 0 会怎样？如果是这种情况，就不需要计算相应专家的操作，因此我们节省了计算量。典型的门控函数是什么？在最传统的设置中，我们只使用一个带有 softmax 函数的简单网络。该网络将学习将输入发送给哪个专家。

Gσ(x)=Softmax(x⋅Wg) G_\sigma(x) = \text{Softmax}(x \cdot W_g)

Shazeer 的工作还探索了其他门控机制，例如 Noisy Top-k Gating。这种门控方法引入了一些（可调的）噪声，然后保留前 k 个值。即：

1. 我们添加一些噪声

H(x)i=(x⋅Wg)i+StandardNormal()⋅Softplus((x⋅Wnoise)i) H(x)_i = (x \cdot W_{\text{g}})_i + \text{StandardNormal()} \cdot \text{Softplus}((x \cdot W_{\text{noise}})_i)

1. 我们只选取前 k 个

KeepTopK(v,k)i={viif vi is in the top k elements of v,−∞otherwise. \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}

1. 我们应用 softmax。

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

这种稀疏性带来了一些有趣的特性。通过使用足够低的 k（例如一或二），我们可以比激活许多专家时更快地训练和运行推理。为什么不直接选择顶级专家？最初的猜想是需要路由到多个专家，才能让门控学习如何路由到不同的专家，因此至少必须选择两个专家。Switch Transformers 部分重新审视了这一决定。

我们为什么要添加噪声？那是为了负载均衡！

为 MoE 平衡 token 负载

如前所述，如果我们的所有 token 都只发送给少数几个热门专家，会导致训练效率低下。在正常的 MoE 训练中，门控网络会收敛到主要激活相同的少数专家。这种自我强化会随着受青睐的专家训练得更快而被更多地选择。为了缓解这一问题，引入了一个辅助损失，以鼓励给予所有专家同等的重要性。该损失确保所有专家接收到的训练样本数量大致相等。以下部分还将探讨专家容量的概念，它引入了一个专家可以处理多少 token 的阈值。在transformers中，辅助损失通过aux_loss参数暴露出来。

MoE 与 Transformer

Transformer 是一个非常明确的案例，表明扩大参数数量可以提高性能，因此 Google 通过GShard探索这一点并不令人意外，该工作探索了将 Transformer 扩展到超过 6000 亿参数。

GShard 在编码器和解码器中都使用 top-2 门控，将每隔一个 FFN 层替换为 MoE 层。下一张图展示了编码器部分的样子。这种设置对于大规模计算非常有利：当我们扩展到多个设备时，MoE 层跨设备共享，而所有其他层则被复制。这在“让 MoE 飞起来”部分有进一步讨论。

为了在大规模下保持负载均衡和效率，GShard 作者除了引入与上一节讨论类似的辅助损失外，还引入了几个变化：

• 随机路由：在 top-2 设置中，我们总是选择排名第一的专家，但第二个专家按其权重成比例的概率被选中。
• 专家容量：我们可以设置一个专家可以处理多少 token 的阈值。如果两个专家都达到容量上限，该 token 被视为溢出，并通过残差连接发送到下一层（在其他项目中则完全丢弃）。这个概念将成为 MoE 最重要的概念之一。为什么需要专家容量？因为所有张量形状在编译时都是静态确定的，但我们无法提前知道有多少 token 会分配给每个专家，所以我们需要固定容量因子。

GShard 论文的贡献在于表达了适用于 MoE 的并行计算模式，但讨论这些超出了本博客文章的范围。

注意：当我们进行推理时，只有部分专家会被触发。同时，还有一些共享计算，例如自注意力，会应用于所有 token。这就是为什么当我们谈论 8 个专家的 47B 模型时，我们可以用 12B 稠密模型的计算量来运行。如果我们使用 top-2，则会使用 14B 参数。但鉴于注意力操作（以及其他操作）是共享的，实际使用的参数数量是 12B。

Switch Transformer

尽管 MoE 展现出了很大的前景，但它们在训练和微调时存在不稳定性。Switch Transformer 是一项非常令人兴奋的工作，深入探讨了这些主题。作者甚至在 Hugging Face 上发布了一个拥有 2048 个专家的1.6 万亿参数 MoE，你可以用 transformers 运行它。Switch Transformer 相比 T5-XXL 实现了 4 倍的预训练加速。

正如在 GShard 中一样，作者将 FFN 层替换为 MoE 层。Switch Transformers 论文提出了一种 Switch Transformer 层，它接收两个输入（两个不同的 token）并拥有四个专家。

与最初使用至少两个专家的想法相反，Switch Transformers 使用了一种简化的单专家策略。这种方法的效果是：

• 路由器计算量减少
• 每个专家的批次大小至少可以减半
• 通信成本降低
• 质量得以保持

Switch Transformers 还探讨了专家容量的概念。

专家容量=(每批 token 数专家数量)×容量因子 \text{Expert Capacity} = \left(\frac{\text{tokens per batch}}{\text{number of experts}}\right) \times \text{capacity factor}

上述建议的容量将批次中的 token 数量平均分配给专家数量。如果我们使用大于 1 的容量因子，就为 token 未完美平衡时提供了缓冲。增加容量会导致更昂贵的设备间通信，因此这是一个需要牢记的权衡。特别是，Switch Transformers 在低容量因子（1-1.25）下表现良好

Switch Transformer 的作者还重新审视并简化了各节中提到的负载均衡损失。对于每个 Switch 层，辅助损失在训练期间被添加到总模型损失中。这种损失鼓励均匀路由，并可以使用超参数进行加权。

作者还尝试了选择性精度，例如用 bfloat16 训练专家，而其余计算使用全精度。较低的精度降低了处理器之间的通信成本、计算成本以及存储张量的内存。在最初的实验中，专家和门网络都用 bfloat16 训练，导致训练更加不稳定。这尤其是由于路由器计算：由于路由器具有指数函数，因此具有更高的精度很重要。为了缓解不稳定性，路由也使用了全精度。

这个 notebook 展示了为摘要任务微调 Switch Transformers，但我们建议先查看微调部分。

Switch Transformers 使用编码器-解码器设置，其中他们做了 T5 的 MoE 对应版本。GLaM 论文探讨了通过训练一个匹配 GPT-3 质量的模型来提升这些模型的规模，同时只使用 1/3 的能量（是的，由于训练 MoE 所需的计算量更少，他们可以将碳足迹减少多达一个数量级）。作者专注于仅解码器模型以及少样本和单样本评估，而不是微调。他们使用了 Top-2 路由和更大的容量因子。此外，他们探索了将容量因子作为一个指标，可以根据想要使用的计算量在训练和评估期间进行更改。

使用路由器 Z-loss 稳定训练

前面讨论的平衡损失可能导致不稳定问题。我们可以使用许多方法来稳定稀疏模型，但会牺牲质量。例如，引入 dropout 可以提高稳定性，但会导致模型质量下降。另一方面，添加更多乘法组件可以提高质量，但会降低稳定性。

Router z-loss 在 ST-MoE 中被提出，通过惩罚进入门控网络的大 logits，显著提高了训练稳定性，且不会导致质量下降。由于该损失鼓励数值的绝对幅度更小，舍入误差得以减少，这对诸如门控这样的指数函数可能相当重要。我们建议查阅论文以了解详情。

专家学到了什么？

ST-MoE 的作者观察到，编码器专家专门处理一组 token 或浅层概念。例如，我们可能最终得到一个标点专家、一个专有名词专家等。另一方面，解码器专家的专门化程度较低。作者还在多语言设置下进行了训练。尽管人们可能会想象每个专家专门处理一种语言，但实际情况恰恰相反：由于 token 路由和负载均衡，没有任何一个专家专门处理任何给定的语言。

扩展专家数量如何影响预训练？

更多的专家会带来更好的样本效率和更快的加速，但这些收益是递减的（尤其是在 256 或 512 之后），并且推理时需要更多的显存。Switch Transformers 在大规模下研究的特性在小规模下也保持一致，即使每层只有 2、4 或 8 个专家。

微调 MoE

Mixtral 在 transformers 4.36.0 版本中受支持。你可以使用 pip install transformers==4.36.0 --upgrade 安装它

密集模型和稀疏模型的过拟合动态非常不同。稀疏模型更容易过拟合，因此我们可以在专家内部探索更高的正则化（例如 dropout）（例如，我们可以为密集层设置一个 dropout 率，为稀疏层设置另一个更高的 dropout 率）。

一个问题是是否在微调时使用辅助损失。ST-MoE 的作者尝试关闭辅助损失，即使多达 11% 的 token 被丢弃，质量也没有受到显著影响。Token 丢弃可能是一种有助于防止过拟合的正则化形式。

Switch Transformers 观察到，在固定的预训练困惑度下，稀疏模型在下游任务中比对应的密集模型表现更差，尤其是在 SuperGLUE 等推理密集型任务上。另一方面，对于 TriviaQA 等知识密集型任务，稀疏模型的表现异常出色。作者还观察到，较少数量的专家有助于微调。另一个证实泛化问题的观察是，模型在较小的任务上表现更差，但在较大的任务上表现良好。

可以尝试冻结所有非专家权重。也就是说，我们只更新 MoE 层。这会导致性能大幅下降。我们可以尝试相反的做法：只冻结 MoE 层中的参数，这几乎与更新所有参数的效果一样好。这有助于加速微调并减少内存占用。这可能有些反直觉，因为在 ST-MoE 项目中，80% 的参数都在 MoE 层中。他们对该架构的假设是，由于专家层每 4 层才出现一次，且每个 token 在每层最多只看到两个专家，因此更新 MoE 参数所影响的层数远少于更新其他参数。

微调稀疏 MoE 时需要考虑的最后一点是，它们有不同的微调超参数设置——例如，稀疏模型往往更受益于较小的批量大小和较高的学习率。

此时，你可能会有些沮丧，因为人们在微调 MoE 方面一直困难重重。令人兴奋的是，最近的一篇论文，MoEs Meets Instruction Tuning（2023 年 7 月），进行了以下实验：

• 单任务微调
• 多任务指令微调
• 多任务指令微调后进行单任务微调

当作者微调 MoE 和对应的 T5 时，T5 表现更好。当作者微调 Flan T5（T5 指令版）MoE 时，MoE 表现显著更好。不仅如此，Flan-MoE 相对于 MoE 的提升比 Flan T5 相对于 T5 的提升更大，这表明 MoE 可能比稠密模型从指令微调中获益更多。MoE 从更多任务数量中获益更多。与之前建议关闭辅助损失函数的讨论不同，该损失实际上防止了过拟合。

何时使用稀疏 MoE 与稠密模型？

专家模型适用于拥有多台机器的高吞吐量场景。在预训练计算预算固定的情况下，稀疏模型会更优。对于显存有限、吞吐量低的场景，稠密模型会更好。

注意：不能直接比较稀疏模型和稠密模型的参数数量，因为两者代表的东西截然不同。

让 MoE 飞起来

最初的 MoE 工作将 MoE 层呈现为分支结构，导致计算缓慢，因为 GPU 并非为此设计，并且由于设备之间需要互相发送信息，网络带宽成为瓶颈。本节将讨论一些现有工作，使这些模型的预训练和推理更加实用。MoE 飞飞飞。

并行

让我们简要回顾一下并行：

• 数据并行：相同的权重在所有核心上复制，数据在各核心间划分。
• 模型并行：模型在各核心间划分，数据在各核心上复制。
• 模型与数据并行：我们可以将模型和数据在各核心间划分。注意，不同的核心处理不同的数据批次。
• 专家并行：专家被放置在不同的工作节点上。如果与数据并行结合使用，每个核心拥有不同的专家，而数据则被划分到所有核心上。

在专家并行中，专家被放置在不同的工作节点上，每个工作节点处理不同的训练样本批次。对于非 MoE 层，专家并行的行为与数据并行相同。对于 MoE 层，序列中的 token 会被发送到所需专家所在的工作节点。

容量因子与通信成本

增加容量因子（CF）可以提高质量，但会增加通信成本和激活值的内存占用。如果全对全通信速度较慢，使用较小的容量因子会更好。一个好的起点是使用 top-2 路由、1.25 的容量因子，并且每个核心放置一个专家。在评估期间，可以更改容量因子以减少计算量。

服务技术

你可以将 mistralai/Mixtral-8x7B-Instruct-v0.1 部署到 Inference Endpoints。

MoE 的一个主要缺点是参数量庞大。对于本地使用场景，人们可能希望使用更小的模型。让我们快速讨论几种有助于服务部署的技术：

• Switch Transformers 的作者进行了早期的蒸馏实验。通过将 MoE 蒸馏回其稠密对应模型，他们能够保留 30-40% 的稀疏性收益。因此，蒸馏提供了更快的预训练和生产中使用更小模型的好处。
• 最近的方法修改了路由机制，将完整句子或任务路由到某个专家，从而允许提取子网络用于服务部署。
• 专家聚合（MoE）：这种技术合并专家的权重，从而减少推理时的参数量。

更多关于高效训练的内容

FasterMoE（2022 年 3 月）分析了 MoE 在高效分布式系统中的性能，分析了不同并行策略的理论极限，以及倾斜专家流行度的技术、降低延迟的细粒度通信调度，以及一种调整后的拓扑感知门控机制，该机制根据最低延迟选择专家，从而实现了 17 倍的加速。

Megablocks（2022 年 11 月）通过提供能够处理 MoE 中动态性的新 GPU 内核，探索了高效的稀疏预训练。他们的方案从不丢弃 token，并能高效映射到现代硬件上，从而带来显著的加速。诀窍是什么？传统的 MoE 使用批量矩阵乘法，这假设所有专家具有相同的形状和相同数量的 token。相比之下，Megablocks 将 MoE 层表示为块稀疏操作，可以适应不均衡的分配。

开源 MoE

如今有多个用于训练 MoE 的开源项目：

• Megablocks：https://github.com/stanford-futuredata/megablocks
• Fairseq：https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm
• OpenMoE：https://github.com/XueFuzhao/OpenMoE

在已发布的开放获取 MoE 中，你可以查看：

• Switch Transformers（Google）：基于 T5 的 MoE 集合，专家数量从 8 到 2048 不等。最大的模型有 1.6 万亿参数。
• NLLB MoE（Meta）：NLLB 翻译模型的 MoE 变体。
• OpenMoE：一个社区项目，发布了基于 Llama 的 MoE。
• Mixtral 8x7B（Mistral）：一个高质量的 MoE，性能超过 Llama 2 70B，且推理速度快得多。还发布了指令微调模型。更多信息请阅读公告博客文章。

令人兴奋的工作方向

进一步实验将稀疏 MoE 蒸馏回参数更少但质量相近的稠密模型。

另一个领域将是 MoE 的量化。QMoE（2023 年 10 月）是这一方向上的良好进展，它将 MoE 量化到每个参数不到 1 比特，从而将使用 3.2TB 加速器的 1.6T Switch Transformer 压缩到仅 160GB。

所以，TL;DR，一些值得探索的有趣领域：

• 将 Mixtral 蒸馏为稠密模型
• 探索专家的模型合并技术及其对推理时间的影响
• 对 Mixtral 执行极端量化技术

一些资源

• Adaptive Mixture of Local Experts（1991）
• Learning Factored Representations in a Deep Mixture of Experts（2013）
• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer（2017）
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding（2020 年 6 月）
• GLaM: Efficient Scaling of Language Models with Mixture-of-Experts（2021 年 12 月）
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity（2022 年 1 月）
• ST-MoE: Designing Stable and Transferable Sparse Expert Models（2022 年 2 月）
• FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models（2022 年 4 月）
• MegaBlocks: Efficient Sparse Training with Mixture-of-Experts（2022 年 11 月）
• Mixture-of-Experts Meets Instruction Tuning:A Winning Combination for Large Language Models（2023 年 5 月）
• Mixtral-8x7B-v0.1、Mixtral-8x7B-Instruct-v0.1。

引用

@misc {sanseviero2023moe,
    author       = { Omar Sanseviero and
                     Lewis Tunstall and
                     Philipp Schmid and
                     Sourab Mangrulkar and
                     Younes Belkada and
                     Pedro Cuenca
                   },
    title        = { Mixture of Experts Explained },
    year         = 2023,
    url          = { https://huggingface.co/blog/moe },
    publisher    = { Hugging Face Blog }
}

Sanseviero, et al., "Mixture of Experts Explained", Hugging Face Blog, 2023.