MoE 调研

MoE模型具有以下特点¹：

相比 dense 模型，训练速度更快；
相比相同参数量的模型推理速度更快；
因为要把所有的 expert 加载，所以需要非常高的 VRAM；
在微调阶段存在一些困难，但是有些 MoE 进行 instruction-tuning 的研究²；

什么是 Mixture of Experts (MoE)

模型规模是提升模型性能的关键因素之一。在有限的计算资源预算下，用更少的训练步数训练一个更大的模型，往往比用更多的步数训练一个较小的模型效果更佳。

混合专家模型 (MoE) 的一个显著优势是它们能够在远少于稠密模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下，您可以显著扩大模型或数据集的规模。特别是在预训练阶段，与稠密模型相比，混合专家模型通常能够更快地达到相同的质量水平。

MoE 由两部分组成：

使用稀疏 MoE 层代替密集前馈网络 (FFN) 层:混合专家模型 (MoE) 的一个显著优势是它们能够在远少于稠密模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下，您可以显著扩大模型或数据集的规模。特别是在预训练阶段，与稠密模型相比，混合专家模型通常能够更快地达到相同的质量水平。
门控网络或路由器，决定将哪些 token 发送给哪个 expert:这个部分用于决定哪些令牌 (token) 被发送到哪个专家。例如，在下图中，“More”这个令牌可能被发送到第二个专家，而“Parameters”这个令牌被发送到第一个专家。有时，一个令牌甚至可以被发送到多个专家。令牌的路由方式是 MoE 使用中的一个关键点，因为路由器由学习的参数组成，并且与网络的其他部分一同进行预训练。

Switch Transformers 论文MoE层

优点和缺点

MoE 优点：

提高预训练的计算效率；
更快的推理速度：MoE 在推理过程中只使用其中的一些参数。与具有相同数量参数的密集模型相比，有更快的推理速度。

MoE缺点：

微调阶段难以泛化，导致过拟合；
推理阶段需要将所有参数都需要加载到RAM中，因此对内存的要求很高；

MoE 的历史

MoE 来源于1991年的论文 aptive Mixture of Local Experts 。这个想法类似于集成方法，是为由单独网络组成的系统提供一个监督程序，每个网络处理训练案例的不同子集。每个单独的网络或专家专门研究输入空间的不同区域。门控网络确定每个专家的权重。在训练过程中，专家和门控都会受到训练。

2010年至2015年间，两个不同的研究领域为教育部后来的进步做出了贡献：

专家作为组件：Eigen、Ranzato 和 Ilya的研究探索了 MoE 作为更深网络的组成部分。这允许将 MoE 作为多层网络中的层，从而使模型既大又高效。
条件计算：传统网络通过每一层处理所有输入数据。Yoshua Bengio 研究了根据输入 token 动态激活或停用组件的方法。

Shazeer 等人（包括 Geoffrey Hinton 和 Jeff Dean、Google 的 Chuck Norris）通过引入稀疏性将这个想法扩展到 137B LSTM（当时事实上的 NLP 架构，由 Schmidhuber 创建），从而保持非常快的速度即使在高规模下也能进行推理。这项工作以翻译为主，但面临着沟通成本高、训练不稳定等诸多挑战。

LSTM 中的 MoE 层

什么是稀疏性（Sparsity）？

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

让我们深入分析 Shazeer 对混合专家模型 (MoE) 在翻译应用中的贡献。条件计算的概念 (即仅在每个样本的基础上激活网络的不同部分) 使得在不增加额外计算负担的情况下扩展模型规模成为可能。这一策略在每个 MoE 层中实现了数以千计甚至更多的专家的有效利用。

这种稀疏性设置确实带来了一些挑战。例如，在混合专家模型 (MoE) 中，尽管较大的批量大小通常有利于提高性能，但当数据通过激活的专家时，实际的批量大小可能会减少。比如，假设我们的输入批量包含 10 个令牌， 可能会有五个令牌被路由到同一个专家，而剩下的五个令牌分别被路由到不同的专家。这导致了批量大小的不均匀分配和资源利用效率不高的问题。在接下来的部分中，将会讨论让 MoE 高效运行的其他挑战以及相应的解决方案。

MoE 会导致 batch size 不平衡和利用率不足。

用一个可训练的 gating network 分配expert： \(y = \sum^{n}_{i=1}{G(x)_i E_i(x)}\) 对所有的输入所有的 expert 都会运行（加权乘法）。最简单的 gating network 是 softmax 函数： \(G_{\sigma}(x) = Softmax(x \cdot W_g)\) Shazeer 的工作还探索了其他门控机制，例如 Noisy Top-k Gating。这种门控方法引入了一些（可调）噪声，然后保留最高的 k 值具体来说:

添加一些噪声
选择保留前 K 个值
应用 Softmax 函数

这种稀疏性引入了一些有趣的特性。通过使用较低的 k 值 (例如 1 或 2)，我们可以比激活多个专家时更快地进行训练和推理。为什么不仅选择最顶尖的专家呢？最初的假设是，需要将输入路由到不止一个专家，以便门控学会如何进行有效的路由选择，因此至少需要选择两个专家。Switch Transformers 就这点进行了更多的研究。

混合专家模型中令牌的负载均衡

正如之前讨论的，如果所有的令牌都被发送到只有少数几个受欢迎的专家，那么训练效率将会降低。在通常的混合专家模型 (MoE) 训练中，门控网络往往倾向于主要激活相同的几个专家。这种情况可能会自我加强，因为受欢迎的专家训练得更快，因此它们更容易被选择。为了缓解这个问题，引入了一个 辅助损失，旨在鼓励给予所有专家相同的重要性。这个损失确保所有专家接收到大致相等数量的训练样本，从而平衡了专家之间的选择。接下来的部分还将探讨专家容量的概念，它引入了一个关于专家可以处理多少令牌的阈值。在 transformers 库中，可以通过 aux_loss 参数来控制辅助损失。

MoE 和 Transformers

Transformer 类模型明确表明，增加参数数量可以提高性能，因此谷歌使用 GShard 尝试将 Transformer 模型的参数量扩展到超过 6000 亿并不令人惊讶。

GShard 将在编码器和解码器中的每个前馈网络 (FFN) 层中的替换为使用 Top-2 门控的混合专家模型 (MoE) 层。下图展示了编码器部分的结构。这种架构对于大规模计算非常有效: 当扩展到多个设备时，MoE 层在不同设备间共享，而其他所有层则在每个设备上复制。我们将在 “让 MoE 起飞”部分对这一点进行更详细的讨论。

GShard 论文 MoE Transformer enconder：

MoE Transformer Encoder

除了辅助损失之外，GShard 作者还引入了一些更改：

随机路由：在 top-2 设置中，我们总是选择顶级 expert ，但选择第二个 expert 的概率与其权重成正比。
专家容量：我们可以设置一个 expert 可以处理多少 token 的阈值。如果两位 expert 都已满员，则令牌被视为溢出，并通过剩余连接发送到下一层（或者完全丢弃）。这个概念将成为 MoE 最重要的概念之一。为什么需要专家能力？由于所有张量形状都是在编译时静态确定的，无法提前知道每个 expert 将获得多少 token ，因此我们需要 fix the capacity factor。

注意: 在推理过程中，只有部分专家被激活。同时，有些计算过程是共享的，例如自注意力 (self-attention) 机制，它适用于所有令牌。这就解释了为什么我们可以使用相当于 12B 稠密模型的计算资源来运行一个包含 8 个专家的 47B 模型。如果我们采用 Top-2 门控，模型会使用高达 14B 的参数。但是，由于自注意力操作 (专家间共享) 的存在，实际上模型运行时使用的参数数量是 12B。

Switch Transformers

switch transformer 包括 2048 experts 。预训练速度比 T5-XXL 快 4 倍。

Switch Transformer Layer

与最初至少两个 expert 的想法相反，switch transformer 使用 single-expert 策略。这个方案的效果是：

减少了路由器的计算量
每个 expert 的 batch size 至少可以减半
交流损耗降低
效果能力得以保留

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

上述建议的容量是将批次中的令牌数量均匀分配到各个专家。如果我们使用大于 1 的容量因子，我们为令牌分配不完全平衡时提供了一个缓冲。增加容量因子会导致更高的设备间通信成本，因此这是一个需要考虑的权衡。特别值得注意的是，Switch Transformers 在低容量因子 (例如 1 至 1.25) 下表现出色。

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

作者还尝试了混合精度的方法，例如用 bfloat16 精度训练专家，同时对其余计算使用全精度进行。较低的精度可以减少处理器间的通信成本、计算成本以及存储张量的内存。然而，在最初的实验中，当专家和门控网络都使用 bfloat16 精度训练时，出现了不稳定的训练现象。这种不稳定性特别是由路由计算引起的，因为路由涉及指数函数等操作，这些操作对精度要求较高。因此，为了保持计算的稳定性和精确性，保持更高的精度是重要的。为了减轻不稳定性，路由过程也使用了全精度。

Switch Transformers 采用了编码器 - 解码器的架构，实现了与 T5 类似的混合专家模型 (MoE) 版本。GLaM 这篇工作探索了如何使用仅为原来 1/3 的计算资源 (因为 MoE 模型在训练时需要的计算量较少，从而能够显著降低碳足迹) 来训练与 GPT-3 质量相匹配的模型来提高这些模型的规模。作者专注于仅解码器 (decoder-only) 的模型以及少样本和单样本评估，而不是微调。他们使用了 Top-2 路由和更大的容量因子。此外，他们探讨了将容量因子作为一个动态度量，根据训练和评估期间所使用的计算量进行调整。

使用 router Z-loss 稳定训练

ST-MoE 使用 Router z-loss通过惩罚进入门控网络的大 logits，显着提高了训练稳定性，而不会降低质量。由于这种损失会导致值的绝对量值变小，因此舍入误差会减少，这对于门控等指数函数非常有影响。

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

ST-MoE 引入的 Router z-loss 在保持了模型性能的同时显著提升了训练的稳定性。这种损失机制通过惩罚门控网络输入的较大 logits 来起作用，目的是促使数值的绝对大小保持较小，这样可以有效减少计算中的舍入误差。这一点对于那些依赖指数函数进行计算的门控网络尤其重要。

expert 模型可以学到什么？

ST-MoE 的研究者们发现，编码器中不同的专家倾向于专注于特定类型的令牌或浅层概念。例如，某些专家可能专门处理标点符号，而其他专家则专注于专有名词等。与此相反，解码器中的专家通常具有较低的专业化程度。此外，研究者们还对这一模型进行了多语言训练。尽管人们可能会预期每个专家处理一种特定语言，但实际上并非如此。由于令牌路由和负载均衡的机制，没有任何专家被特定配置以专门处理某一特定语言。

增加更多专家可以提升处理样本的效率和加速模型的运算速度，但这些优势随着专家数量的增加而递减 (尤其是当专家数量达到 256 或 512 之后更为明显)。同时，这也意味着在推理过程中，需要更多的显存来加载整个模型。值得注意的是，Switch Transformers 的研究表明，其在大规模模型中的特性在小规模模型下也同样适用，即便是每层仅包含 2、4 或 8 个专家。

微调 MoE

稠密模型和稀疏模型在过拟合的动态表现上存在显著差异。稀疏模型更易于出现过拟合现象，因此在处理这些模型时，尝试更强的内部正则化措施是有益的，比如使用更高比例的 dropout。例如，我们可以为稠密层设定一个较低的 dropout 率，而为稀疏层设置一个更高的 dropout 率，以此来优化模型性能。

在微调过程中是否使用辅助损失是一个需要决策的问题。ST-MoE 的作者尝试关闭辅助损失，发现即使高达 11% 的令牌被丢弃，模型的质量也没有显著受到影响。令牌丢弃可能是一种正则化形式，有助于防止过拟合。

Switch Transformers 的作者观察到，在相同的预训练困惑度下，稀疏模型在下游任务中的表现不如对应的稠密模型，特别是在重理解任务 (如 SuperGLUE) 上。另一方面，对于知识密集型任务 (如 TriviaQA)，稀疏模型的表现异常出色。作者还观察到，在微调过程中，较少的专家的数量有助于改善性能。另一个关于泛化问题确认的发现是，模型在小型任务上表现较差，但在大型任务上表现良好。

一种可行的微调策略是尝试冻结所有非专家层的权重。实践中，这会导致性能大幅下降，但这符合我们的预期，因为混合专家模型 (MoE) 层占据了网络的主要部分。我们可以尝试相反的方法: 仅冻结 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 在多任务学习中表现更佳。与之前关闭 辅助损失 函数的做法相反，实际上这种损失函数可以帮助防止过拟合。

稀疏vs稠密

稀疏混合专家模型 (MoE) 适用于拥有多台机器且要求高吞吐量的场景。在固定的预训练计算资源下，稀疏模型往往能够实现更优的效果。相反，在显存较少且吞吐量要求不高的场景，稠密模型则是更合适的选择。

注意: 直接比较稀疏模型和稠密模型的参数数量是不恰当的，因为这两类模型基于的概念和参数量的计算方法完全不同。

如何优化 MoE 模型

最初的混合专家模型 (MoE) 设计采用了分支结构，这导致了计算效率低下。这种低效主要是因为 GPU 并不是为处理这种结构而设计的，而且由于设备间需要传递数据，网络带宽常常成为性能瓶颈。在接下来的讨论中，我们会讨论一些现有的研究成果，旨在使这些模型在预训练和推理阶段更加高效和实用。我们来看看如何优化 MoE 模型，让 MoE 起飞。

并行计算

让我们简要回顾一下并行计算的几种形式:

数据并行: 相同的权重在所有节点上复制，数据在节点之间分割。
模型并行: 模型在节点之间分割，相同的数据在所有节点上复制。
模型和数据并行: 我们可以在节点之间同时分割模型和数据。注意，不同的节点处理不同批次的数据。
专家并行: 专家被放置在不同的节点上。如果与数据并行结合，每个节点拥有不同的专家，数据在所有节点之间分割。

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

容量因子和通信开销

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

部署技术

部署混合专家模型 (MoE) 的一个关键挑战是其庞大的参数规模。对于本地使用情况，我们可能希望使用更小的模型。为了使模型更适合部署，下面是几种有用的技术:

预先蒸馏实验: Switch Transformers 的研究者们进行了预先蒸馏的实验。他们通过将 MoE 模型蒸馏回其对应的稠密模型，成功保留了 30-40%的由稀疏性带来的性能提升。预先蒸馏不仅加快了预训练速度，还使得在推理中使用更小型的模型成为可能。
任务级别路由: 最新的方法中，路由器被修改为将整个句子或任务直接路由到一个专家。这样做可以提取出一个用于服务的子网络，有助于简化模型的结构。
专家网络聚合: 这项技术通过合并各个专家的权重，在推理时减少了所需的参数数量。这样可以在不显著牺牲性能的情况下降低模型的复杂度。

高效训练

FasterMoE (2022 年 3 月) 深入分析了 MoE 在不同并行策略下的理论性能极限，并且探索了一系列创新技术，包括用于专家权重调整的方法、减少延迟的细粒度通信调度技术，以及一个基于最低延迟进行专家选择的拓扑感知门控机制。这些技术的结合使得 MoE 运行速度提升高达 17 倍。

Megablocks (2022 年 11 月) 则专注于通过开发新的 GPU kernel 来处理 MoE 模型中的动态性，以实现更高效的稀疏预训练。其核心优势在于，它不会丢弃任何令牌，并能高效地适应现代硬件架构 (支持块稀疏矩阵乘)，从而达到显著的加速效果。Megablocks 的创新之处在于，它不像传统 MoE 那样使用批量矩阵乘法 (这通常假设所有专家形状相同且处理相同数量的令牌)，而是将 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): Collection of T5-based MoEs going from 8 to 2048 experts. The largest model has 1.6 trillion parameters.
NLLB MoE (Meta): A MoE variant of the NLLB translation model.
OpenMoE: A community effort that has released Llama-based MoEs.
Mixtral 8x7B (Mistral): A high-quality MoE that outperforms Llama 2 70B and has much faster inference. A instruct-tuned model is also released. Read more about it in the announcement blog post.

一些有趣的研究方向

首先是尝试将稀疏混合专家模型 (SMoE) 蒸馏回到具有更少实际参数但相似等价参数量的稠密模型。

MoE 的量化也是一个有趣的研究领域。例如，QMoE (2023 年 10 月) 通过将 MoE 量化到每个参数不到 1 位，将 1.6 万亿参数的 Switch Transformer 所需的存储从 3.2TB 压缩到仅 160GB。

简而言之，一些值得探索的有趣领域包括:

将 Mixtral 蒸馏成一个稠密模型。
探索合并专家模型的技术及其对推理时间的影响。
尝试对 Mixtral 进行极端量化的实验。