这里带来模块化MoE将成为视觉多任务学习基础模型
UMass Amherst 淦创团队提出了 Mod-Squad 模型，它可以从多任务大模型中提取针对单一任务的相同性能小模型，在 Taskonomy 大数据集和 PASCALContext 数据集上取得了最佳效果。
多任务学习（MTL）存在很多挑战，因为不同任务之间的梯度可能矛盾。为了利用任务之间的关联，作者引入了 Mod-Squad 模型，它是多个专家组成的模块化模型。模型可以灵活优化任务和专家的匹配，针对任务选择部分专家。模型让每一个专家只对应部分任务，每一个任务只对应部分专家，以此最大化利用任务之间的正向联系。Mod-Squad 整合了 Mixture of Expert (MoE) 层到 Vision Transformer 模型中，并引入了新的损失函数鼓励专家和任务之间的稀疏但强烈的依赖关系。此外，对于每个任务，模型都可以只保留小部分专家网络，并且性能与原来的大模型相同。模型在 13 个视觉任务的 Taskonomy 大数据集和 PASCALContext 数据集上取得了最佳效果。
GPT-4是8个2200亿MoE模型
GPT-4远不止1万亿，甚至，还是8个2200亿参数组成的混合专家模型（MoE）。
2023年6月，美国知名骇客George Hotz在接受采访时透露，GPT-4由8个220B模型组成。这么算来，8 x 220B = 1.76万亿。就连PyTorch的创建者Soumith Chintala对此也深信不疑。
论文地址：https://arxiv.org/abs/2212.08066
项目地址：https://vis-www.cs.umass.edu/mod-squad/
Github地址：https://github.com/UMass-Foundation-Model/Mod-Squad
多任务学习（MTL）的目的是建模任务之间的关系，并为多种任务构建统一的模型。如图 1 所示，Mod-Squad 的主要动机就是要让专家只被一些任务更新而不是所有任务，且每一个任务只更新部分专家。这样可以利用模型的全部容量的同时避免任务间的互相干扰。
MoE 应用于大模型，GPT-4并不是第一个。在2022年的时候，Google 就提出了MoE大模型Switch Transformer，模型大小是1571B，Switch Transformer在预训练任务上显示出比 T5-XXL（11B） 模型更高的样本效率。在相同的训练时间和计算资源下，Switch Transformer 能够达到更好的性能。
除了GPT-4和Switch Transformer，国内的团队DeepSeek 也开源了国内首个 MoE 大模型 DeepSeekMoE。
- DeepSeekMoE 2B可接近2B Dense，仅用了17.5%计算量。
- DeepSeekMoE 16B性能比肩 LLaMA2 7B 的同时，仅用了40%计算量。
- DeepSeekMoE 145B 优于Google 的MoE大模型GShard，而且仅用 28.5%计算量即可匹配 67B Dense 模型的性能。
一时间，国内大模型开始朝着MoE方向大步前进，估计在2024年，会有越来越多大模型选择MoE架构。
那么，究竟什么是MoE大模型？MoE大模型具备哪些优势？本文就带你一探究竟。
什么是MoE大模型？
MoE，全称为Mixed Expert Models，翻译过来就是混合专家模型。MoE并不是什么最新技术，早在1991年的时候，论文Adaptive Mixture of Local Experts就提出了MoE。
我们知道，模型规模是提升模型性能的关键因素之一，这也是为什么今天的大模型能取得成功。在有限的计算资源预算下，用更少的训练步数训练一个更大的模型，往往比用更多的步数训练一个较小的模型效果更佳。
MoE 的一个显著优势是它们能够在远少于 Dense 模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下，您可以显著扩大模型或数据集的规模。特别是在预训练阶段，与稠密模型相比，混合专家模型通常能够更快地达到相同的质量水平。
MoE基于Transformer架构，主要由两部分组成：
- 稀疏 MoE 层： 这些层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干“专家”(例如 8 个)，每个专家本身是一个独立的神经网络。在实际应用中，这些专家通常是前馈网络 (FFN)，但它们也可以是更复杂的网络结构。
- 门控网络或路由: 这个部分用于决定哪些 token 被发送到哪个专家。例如，在下图中，“More”这个 token 可能被发送到第二个专家，而“Parameters”这个 token 被发送到第一个专家。有时，一个 token 甚至可以被发送到多个专家。token 的路由方式是 MoE 使用中的一个关键点，因为路由器由学习的参数组成，并且与网络的其他部分一同进行预训练。
总结来说，在混合专家模型 (MoE) 中，我们将传统 Transformer 模型中的每个前馈网络 (FFN) 层替换为 MoE 层，其中 MoE 层由两个核心部分组成: 一个路由器（或者叫门控网络）和若干数量的专家。
MoE大模型具备哪些优势？
MoE的最大优势就是与Dense模型相比，在相同计算资源下，训练速度更快，而且可以训练更大的模型。比如Google的Switch Transformer，模型大小是T5-XXL的15倍，在相同计算资源下，Switch Transformer模型在达到固定困惑度 PPL 时，比T5-XXL模型快4倍。
相同计算资源下，Google的MoE大模型能够在相同计算资源下，以更快的速度达到相同的PPL，而且模型是T5的15倍；DeepSeek的16B MoE大模型，仅在40% 的计算量的情况下，性能和LLaMA 2 7B效果比肩。
总结MoE大模型优点，主要有以下3点：
- 训练速度更快，效果更好。
- 相同参数，推理成本低。
- 扩展性好，允许模型在保持计算成本不变的情况下增加参数数量，这使得它能够扩展到非常大的模型规模，如万亿参数模型。
- 多任务学习能力：MoE在多任务学习中具备很好的新能（比如Switch Transformer在所有101种语言上都显示出了性能提升，证明了其在多任务学习中的有效性）。
而MoE大模型的缺点，主要有以下4点：
- 训练稳定性：MoE在训练过程中可能会遇到稳定性问题。
- 通信成本：在分布式训练环境中，MoE的专家路由机制可能会增加通信成本，尤其是在模型规模较大时。
- 模型复杂性：MoE的设计相对复杂，可能需要更多的工程努力来实现和优化。
- 下游任务性能：MoE由于其稀疏性，使得在Fine-tuning过程中容易出现过拟合。
接下来，我们就介绍下MoE的主要原理。通过后面的介绍，我们主要需要回答以下3个问题：
- MoE为什么可以实现更大模型参数、更低训练成本？
- MoE如何解决训练稳定性问题？
- MoE如何解决Fine-Tuning过程中的过拟合问题？
一、Adaptive mixtures of local experts
Adaptive mixtures of local experts，这是大多数MoE论文都引用的最早的一篇文章，发表于1991年，作者中有两个大家熟知的大佬：Michael Jordan 和 Geoffrey Hinton。
论文介绍了一种新的监督学习过程，用于由多个独立网络组成的系统，每个网络处理训练集合的子集。这种新方法可以看作是多层监督网络的模块化版本，或者是竞争性学习的关联版本，因此提供了这两种看似不同的方法之间的新联系。
如果一个多层网络用来训练不同的子任务，通常会有强烈的干扰效应，这会导致学习过程变慢和泛化能力差。这种干扰效应的原因在于，当网络试图同时学习多个子任务时，不同任务的学习过程可能会相互干扰。例如，学习一个子任务时对权重的调整可能会影响其他子任务的学习效果，因为这些权重变化会改变其他子任务的loss。这种相互影响使得网络在处理每个子任务时都试图最小化所有其他子任务的loss。
为了解决这个问题，论文提出了使用多个模型（即专家，expert）去学习，使用一个门控网络（gating network）来决定每个数据应该被哪个模型去训练，这样就可以减轻不同类型样本之间的干扰。
二、Sparsely-Gated MoE
在 2010 至 2015 年间，两个独立的研究领域为混合专家模型 (MoE) 的后续发展做出了显著贡献：
- 组件专家：在传统的 MoE 设置中，整个系统由一个门控网络和多个专家组成。在支持向量机 (SVMs) 、高斯过程和其他方法的研究中，MoE 通常被视为整个模型的一部分。然而，Eigen、Ranzato 和 Ilya 的研究 探索了将 MoE 作为更深层网络的一个组件。这种方法允许将 MoE 嵌入到多层网络中的某一层，使得模型既大又高效。
- 条件计算（Conditional Computation）：传统的神经网络通过每一层处理所有输入数据。在这一时期，Yoshua Bengio 等研究人员开始探索基于输入 token 动态激活或停用网络组件的方法。
在 2017 年，Shazeer 等人(团队包括 Geoffrey Hinton 和 Jeff Dean，后者有时被戏称为“谷歌的 Chuck Norris”) 将这一概念应用于 137B 的 LSTM 。通过引入稀疏性，这项工作在保持极高规模的同时实现了快速的推理速度。在牺牲极少的计算效率的情况下，把模型规模提升1000多倍。
这项工作被发表在论文Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer中，和 1991 年Adaptive mixtures of local experts的工作对比，这里的 Sparsely-Gated MoE 主要有两个区别：
- Sparsely-Gated：不是所有expert都会起作用，而是极少数的expert会被使用来进行推理。这种稀疏性，也使得我们可以使用海量的experts来把模型容量做的超级大。
- token-level：前面那个文章，是 sample-level 的，即不同的样本，使用不同的专家，但是这篇则是 token-level 的，一个句子中不同的 token 使用不同的专家。
下图是Sparsely-Gated MoE的模型结构，被运用于RNN结构中：
图2-1 Sparsely-Gated MoE
如图 2-1 所示，每个 token，都会有一个 MoE Layer，每个 MoE layer 中包含了一堆的 experts，每个 expert都是一个小型的 FFN，还有一个 Gating Network 会根据当前 token，选择少数几个 expert 来进行计算。
2.1 门控网络（Gating Network）
门控网络（Gating Network）的设计和实现，这是Sparsely-Gated MoE 层的核心组成部分。门控网络负责为每个输入 token 选择一个稀疏的专家组合。
2.2 平衡专家利用率（Balancing Expert Utilization）
论文指出，门控网络倾向于收敛到一种状态，总是为相同的几个专家产生大的权重。这种不平衡是自我强化的，因为受到青睐的专家训练得更快，因此被门控网络更多地选择。这种不平衡可能导致训练效率低下，因为某些专家可能从未被使用过。
三、GShard
上面的两篇，都是 MoE 系列的基础工作，而且都没有在大模型上得到广泛应用。接下来介绍的工作，都是近几年在大模型上应用上比较出色的工作。
GShard 是谷歌 2021 年在论文 GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding 中提出的使用 GShard 实现 MoE 跨设备分片的方法。
按照文章的说法，GShard 是第一个将 MoE 的思想拓展到 Transformer 上的工作。具体的做法是，把Transformer 的 encoder 和 decoder 中，每隔一个的 FFN 层，替换成 MoE 层，使用的都是 Top-2 门控网络。
GShard 将在编码器和解码器中的每个 FFN 层替换为使用 Top-2 门控的 MoE 层。下图展示了编码器部分的结构。这种架构对于大规模计算非常有效：当扩展到多个设备时，MoE 层在不同设备间共享，而其他所有层则在每个设备上复制。
图中3-1中，分为标准Transformer（a）、MoE Transformer（b）以及GShard Transformer：
- 标准 Transformer（a）：是标准的Transformer编码器，其中每个 token 通过一个标准的 FFN。
- MoE Transformer（b）：将每隔一个的 FFN 层替换为 MoE 层。这意味着在编码器中，不再是每个 token 都通过相同的 FFN，而是通过一个由多个专家组成的 MoE 层。
- MoE跨设备分片（c）：它展示了 MoE 层是如何在多个设备上进行分片的。GShard MoE 层中的专家网络（experts）被分布在不同的设备上。每个专家网络负责处理一部分输入数据，并且每个 token 根据门控机制的输出被分配到一个或两个专家网络中。这样，整个 MoE 层的计算被分散到了多个设备上，每个设备负责处理一部分计算任务。
实现 MoE 跨设备分片的关键技术是模型并行化（model parallelism）和数据并行化（data parallelism）的结合。在模型并行化中，模型的不同部分（在这里是 MoE 层的专家网络）被分配到不同的设备上。在数据并行化中，输入数据（token）被分割成多个部分，每个部分被分配给不同的设备进行处理。
为了实现这种分片，论文中提到的 GShard 模块提供了一套 API 和编译器扩展，允许用户在模型代码中简单地注释关键张量，指定它们应该如何在设备集群上进行分片。这样，编译器就可以自动地将计算图（computation graph）转换为可以在多个设备上并行执行的程序，而不需要用户手动处理复杂的数据分片和通信逻辑。
由于专家被分配到不同设备，可以并行计算，因此大大提升了模型的计算效率，这也解释了为什么 MoE 可以实现更大模型参数、更低训练成本。
为了保持负载平衡和训练效率，GShard 的作者除了引入上节 Sparsely-Gated MoE 中的辅助 loss 外，还引入了一些关键变化：
- 随机路由： 在 Top-2 设置中，GShard 始终选择排名最高的专家，但第二个专家是根据其权重比例随机选择的。
- 专家容量： 我们可以设定一个阈值，定义一个专家能处理多少 token。如果两个专家的容量都达到上限，token 就会溢出，并通过残差连接传递到下一层，或在某些情况下被完全丢弃。专家容量是 MoE 中最重要的概念之一。为什么需要专家容量呢？因为所有张量的形状在编译时是静态确定的，我们无法提前知道多少 token 会分配给每个专家，因此需要一个固定的容量因子。
注意： 在推理过程中，只有部分专家被激活。同时，有些计算过程是共享的，例如自注意力 (self-attention) 机制，它适用于所有 token。这就解释了为什么我们可以使用相当于 12B Dense 模型的计算资源来运行一个包含 8 个专家的 47B 模型。如果我们采用 Top-2 门控，模型会使用高达 14B 的参数。但是，由于自注意力操作 (专家间共享) 的存在，实际上模型运行时使用的参数数量是 12B。
四、Switch Transformers
尽管 MoE 显示出了很大的潜力，但是由于复杂性、通信成本以及训练和微调过程的不稳定性，模型广泛采用仍需要优化。
而在2022年，Google 提出的 Switch Transformers 一定程度缓解了这些问题。Switch Transformers 是一项非常激动人心的工作，它深入研究了这些话题。作者在 Hugging Face 上发布了一个 1.6 万亿参数的 MoE，拥有 2048 个专家，你可以使用 transformers 库来运行它。Switch Transformers 实现了与 T5-XXL 相比 4 倍的预训练速度提升。
Switch Transformers 简化了 MoE 路由算法，设计了直观的改进模型，降低了通信和计算成本。Switch Transformers 的训练方法减轻了不稳定性，并且首次展示了用较低精度（bfloat16）格式训练大型稀疏模型的可能性。
和 T5 Base、T5 Large 相比，Switch Transformers 在相同计算资源情况下获得了高达 7 倍的预训练速度。在多语言实验中，Switch Transformers 在所有 101 种语言测试中都取得了提升。Switch Transformers 通过在爬虫语料库上预训练了一个高大万亿参数规模的模型，实现了与 T5-XXL 相比4倍的加速。
4.1 Switch Transformer 主要优化
Swith Transformer 在论文中提到其设计的指导原则是——尽可能地把 Transformer 模型的参数量做大！（同时以一种简单高效的实现方式）
和其他 MoE 模型的一个显著不同就是，Switch Transformer 的门控网络每次只路由到 1 个 expert，也就是每次只选取 top1 的专家，而其他的模型都是至少 2 个。这样就是最稀疏的 MoE 了，因此单单从 MoE layer 的计算效率上讲是最高的了。
与最初使用至少两个专家的想法相反，Switch Transformer 采用了简化的单专家策略，每次只选择一个专家。这种方法的效果包括：
- 减少了路由计算，一个 token 每次只路由到一个专家
- 每个专家的 batch size（专家容量、Expert Capacity） 至少可以减半
- 简化路由的实现，降低了 MoE 中的通信成本
4.2 Switch Routing
上面提到了 Switch Transformer 可以降低每个专家的专家容量，那么什么是专家容量？
专家容量（Expert Capacity） 是指每个专家在模型中处理的 token 数。专家容量的计算方式如下：
这里为什么要计算一个专家容量？这个专家容量又有什么作用？
在编译时，所有 tensor 的形状都是静态确定的。这意味着在编译阶段，模型的架构和数据布局已经被定义，包括模型的层数、每层的输入和输出维度等。
尽管 tensor 的形状是静态的，但在训练和推理过程中，模型的计算是动态的。这是因为模型中的路由器（门控网络）会根据输入数据动态地将 token 分配给不同的专家。这种动态性要求模型能够在运行时灵活地处理数据分布。
而这个专家容量的作用就是将 batch 中的总 token 数平均分配给所有专家。然后，为了应对 token 分布不均的情况，会通过一个容量因子（capacity factor）来扩展每个专家的容量。
容量因子是一个大于 1.0 的数，它的作用是为每个专家提供额外的缓冲空间，以容纳可能超出平均分配的 token。这样，即使某些专家接收到的 token 数量超过了平均值，也能够处理这些额外的 token，而不会因为容量不足而导致计算跳过。
下图是不同容量因子下的动态路由。
如图4-5所示，容量因子 Capacity Factor = 1.0的时候，输入6个 token，那么每个专家的专家容量等于 2（Expert Capacity = 6/3 * 1 = 2），Expert 1 被分配到了 3 个 token，超出了专家容量，这些超出的 token 被称为“溢出 token”（图4-5（左）中的虚线部分）。对于这些溢出的 token，模型会跳过计算，直接将 token 的表示通过残差连接传递到下一层。
而如果容量因子 Capacity Factor = 1.5，这时专家容量等于 3，每个专家就能处理 3 个 token（图4-5（右））。
虽然增加容量因子可以减少 token 溢出，但是它也有缺点。如果容量因子设置得过高，会导致计算资源和内存的浪费，因为模型会为可能永远不会用到的 token 分配额外的资源。在论文中，Switch Transformers 在低容量因子 (例如 1 至 1.25) 下表现出色。
下表是不同容量因子的效果对比。
表4-1：Switch Transformer 和 MoE 的效果对比
模型 | 容量因子 | 训练100k steps后的负对数困惑度（越大越好） | 模型到达指定负对数困惑度（-1.5）所需时间（单位小时） | 训练速度（每秒处理的样本数） |
T5-Base | -1.731 | 没有达到 | 1600 | |
T5-Large | -1.550 | 131.1 | 470 | |
MoE-Base | 2.0 | -1.547 | 68.7 | 840 |
Switch-Base | 2.0 | -1.554 | 72.8 | 860 |
MoE-Base | 1.25 | -1.559 | 72.8 | 790 |
Switch-Base | 1.25 | -1.553 | 65.0 | 910 |
MoE-Base | 1.0 | -1.572 | 80.1 | 860 |
Switch-Base | 1.0 | -1.561 | 62.8 | 1000 |
Switch-Base+ | 1.0 | -1.534 | 67.6 | 780 |
- 以上模型都是在相同的计算资源（32核）和硬件（TPUv3）上进行训练的。
- 所有 MoE 和 Switch Transformer 模型都使用 128 个专家。
- 为了达到负对数困惑度为-1.50，所有模型都需要进行超过 100k steps 的预训练。
- Switch-Base+：对于这个模型，作者增加了模型的大小，直到其训练速度与 MoE 模型相匹配。这通过增加模型的隐藏层大小（从768增加到896）和 head 的数量（从14增加到16）来实现。
- T5-Base 在训练的 100k 步内没有达到这个负对数困惑度：这表示在给定的训练步数内，T5-Base 模型没有达到设定的效果，这可能是由于其性能不如 Switch Transformer 或 MoE Transformer 模型。
Switch Transformer 的作者还重新审视并简化了前面章节中提到的负载均衡损失（公式（12））。通过合理设置负载均衡损失的系数，可以在训练过程中实现专家之间的良好负载分布。下面介绍下具体实现。
4.3 不同的负载均衡损失
在稀疏模型中，专家的数量通常分布在多个设备上，每个专家负责处理一部分输入数据。理想情况下，每个专家应该处理相同数量的数据，以实现资源的均匀利用。然而，在实际训练过程中，由于数据分布的不均匀性，某些专家可能会处理更多的数据，而其他专家可能会处理较少的数据。这种不均衡可能导致训练效率低下，因为某些专家可能会过载，而其他专家则可能闲置。为了解决这个问题，论文中引入了一种辅助损失函数，以促进专家之间的负载均衡。
4.4 稀疏路由和负载均衡loss的合并效果
前面介绍了 Switch Transformer 的主要优化：稀疏路由和负载均衡损失。下面介绍一下将这两项优化合并在一起的实验效果。
实验设置如下：
- 首先从 C4 数据集上进行预训练，使用 MLM（Masked Language Modeling） 作为预训练目标。在这个任务中，模型被训练来预测被 mask 的 token。
- 对比 Switch Transformer 与 MoE Transformer 以及 T5。Switch Transformer 在计算量（FLOPs）上与 T5-Base 匹配，即每个 token 应用的计算量相同。MoE Transformer 使用 top-2 路由。
- 所有模型都在相同的硬件（TPUv3）上进行了相同数量的步骤训练。
实验结论如下（可以参见前面的表4-1）：
- Switch Transformer 在速度和效果上都优 MoE Transformer。对于固定的计算量和时间，Switch Transformer 实现了最佳结果。
- Switch Transformer 的计算量小于同等参数的 MoE 模型。如果将 Switch Transformer 的规模增加到匹配MoE Transformer 的训练速度，那么它在每个步骤上都优于所有 MoE 模型。
- Switch Transformer 在较低的容量因子（1.0, 1.25）下表现更好。较低的专家容量表明在大模型中，模型内存非常稀缺，容量因子应尽可能小。
4.5 改进训练和Fine-Tuning技术
- 精度选择
作者还尝试了混合精度的方法，例如用 bfloat16 精度训练专家，同时对其余计算使用全精度进行。较低的精度可以减少处理器间的通信成本、计算成本以及存储 tensor 的内存。然而，在最初的实验中，当专家和门控网络都使用 bfloat16 精度训练时，出现了不稳定的训练现象。这种不稳定性主要是由路由计算引起的，因为路由涉及指数函数等操作，这些操作对精度要求较高。因此，为了保持计算的稳定性和精确性，保持更高的精度是重要的。为了减轻不稳定性，路由过程也使用了全精度。
下面的表 4-2 显示了混合精度训练的效果，将路由器输入转换为 float32，同时保持其他部分的精度为 bfloat16。这种策略允许模型在几乎与 bfloat16 精度相同的训练速度下，实现与 float32 训练相当的稳定性。
表4-2：不同精度效果对比
模型精度选择 | 效果（负对数困惑度） | 训练速度（每秒处理样本数） |
Switch-Base (float32) | -1.718 | 1160 |
Switch-Base (bfloat16) | -3.780 | 1390 |
Switch-Base (混合精度) | -1.716 | 1390 |
实验表明，使用混合精度的 Switch-Base 在固定步数的早期训练中，其效果（以负对数困惑度为衡量标准）与使用 float32 训练的模型相似，同时速度接近 bfloat16。
- 更小的参数初始化
在深度学习中，适当的权重初始化对于模型的成功训练至关重要。作者观察到，在 Switch Transformer 模型中，这一点尤其明显。
为了提高模型的稳定性，作者建议减少默认的 Transformer 初始化规模。在 Transformer 模型中，权重矩阵通常是从一个截断的正态分布，其均值为0，标准差由一个超参数 s 决定。作者建议将这个初始化超参数 s 从默认值1.0 减少 10 倍，即 s = 0.1。这种较小的初始化规模有助于提高模型效果和减少训练过程中的不稳定性。
表4-3：减小参数初始化规模可以提升训练稳定性
权重初始化规模 | 负对数困惑度 | 负对数困惑度标准差 |
0.1倍初始化 | -2.72 | 0.01 |
1.0倍初始化 | -3.60 | 0.68 |
表 4-3 中的数据表明，通过减少初始化规模，模型效果和稳定性得到了提升。这种改进对于大模型，如 Switch Transformer，尤其重要。
- Fine-Tuning 过程正则化
为了解决 Fine-Tuning 过程中的过拟合问题，作者提出了增加 dropout的策略，特别是在专家层（expert layers）中。他们称之为“expert dropout”，即在 Fine-Tuning 时只在专家层增加 dropout 率。
表 4-4 显示了在 Fine-Tuning Switch Transformer 时，不同 dropout 率的实验结果。这些模型是在 C4 数据集上预训练的，然后进行了 Fine-Tuning。
表4-4：Fine-Tuning 过程中正则化效果
模型（dropout） | GLUE | CNNDM | SQuAD | SuperGLUE |
T5-Base (d=0.1) | 82.9 | 19.6 | 83.5 | 72.4 |
Switch-Base (d=0.1) | 84.7 | 19.1 | 83.7 | 73.0 |
Switch-Base (d=0.2) | 84.4 | 19.2 | 83.9 | 73.2 |
Switch-Base (d=0.3) | 83.9 | 19.6 | 83.4 | 70.7 |
Switch-Base (d=0.1, expert d=0.4) | 85.2 | 19.6 | 83.7 | 73.0 |
通过这种 expert dropout 策略，有效地减少了过拟合的风险，同时保持了模型在下游任务上的性能。这种正则化方法对于处理具有大量参数的稀疏模型特别有用，因为它可以帮助模型更好地泛化到未见过的数据。
4.6 高效训练：数据、模型、专家并行
任意增加专家数量会导致收益递减（如图4-3所示）。这意味着在某个点之后，继续增加专家数量不会显著提高模型性能。但是可以通过增加模型的维度，如模型的隐藏层大小（dmodel）或前馈网络的维度（dff）来继续提升模型效果。但是这样又会导致显存和内存开销增加，这时候就可以通过并行技术，解决高效训练问题。
这里补充一下关于各种并行的方法的解释。标准的数据并行的定义是一个 batch 的数据在不同的 device 上并行处理，这时每一个 device 上都保存了模型的一份完整拷贝，前向计算完进行梯度汇总和更新。模型并行表示模型不同的参数（层、组件）分配到不同的 device 上，处理一个 batch 的数据。
- 数据并行（Data Parallelism）
第一列表示数据并行，模型权重拷贝 16 份，16 个同一种颜色矩阵分别表示一个完整的模型，图4-6（b）则是一个完整的矩阵，这里可以理解为 16 个模型计算完成后由于存在梯度汇总再更新的步骤，所以整体更新的是一个batch，因此这里 Data Parallelism 是一个唯一的矩阵。简单来说就是模型复制，数据并行。
2. 模型并行（Model Parallelism）
模型并行部分从模型侧看出来，16个 cores 维护的是一个整体的模型，但是每一个 core 只分配到其中部分模型参数（图4-6（a）），同一个 batch 数据在所有的 core 上计算（图4-6（b）），由于 1 个 core 中分布了不同的模型权重，每次计算完都需要和其他的 core 进行通信。
3. 模型和数据并行
总共有 NN 个 cores，其中 N=n\times mN=n\times m ， nn 代表数据并行维度上的分割因子， mm 代表模型并行维度上的分割因子。现在每个 core 处理的是 B/nB/n 个 token 以及
d_{ff}/md_{ff}/m 个权重。
4. 专家和数据并行
5. 专家、模型和数据并行
五、GLaM
除了 Switch Transformer，Google还推出另外一个 MoE 模型：GLaM (Generalist Language Model)。
GLaM 比 GPT-3 大三倍，但是由于使用了 Sparse MoE 的设计，训练成本却只有 GPT-3 的 1/3，而且在 29 个NLP 任务上超越了 GPT-3。
下面是 Google Blog 中 GLaM 的模型结构，非常形象。
从图5-1来看，和前面的 GShard 非常像，反正都是出自 Google，也不知道谁借鉴的谁。
表5-1：GLaM 实验
模型 | 模型类型 | 参数量 | 激活的参数量 |
BERT | Dense Encoder-only | 340M | 340M |
T5 | Dense Encoder-decoder | 13B | 13B |
GPT-3 | Dense Decoder-only | 175B | 175B |
Jurassic-1 | Dense Decoder-only | 178B | 178B |
Gopher | Dense Decoder-only | 280B | 280B |
Megatron-530B | Dense Decoder-only | 530B | 530B |
GShard-M4 | MoE Encoder-decoder | 600B | 1.5B |
Switch-C | MoE Encoder-decoder | 1.5T | 1.5B |
GLaM (64B/64E) | MoE Decoder-only | 1.2T | 96.6B |
上表展示了 GLaM 跟其他大模型的对比。可以看到，虽然 GLaM 的总参数量有 1.2T，但是在计算中实际激活的参数量只有 96B，所以在 inference 的时候，比 GPT-3 等 dense model 要快得多。
六、ST-MOE
之前讨论的负载均衡损失可能会导致稳定性问题。我们可以使用许多方法来稳定稀疏模型的训练，但这可能会牺牲模型质量。例如，引入 dropout 可以提高稳定性，但会导致模型质量下降。
6.1 用 Router z-loss 稳定模型训练
在论文 ST-MOE: Designing Stable and Transferable Sparse Expert Models 中，作者提出了一种新的辅助损失函数，称为 Router z-loss，用于提高稀疏模型的训练稳定性，同时保持或稍微提高模型质量。这个损失函数是针对稀疏专家模型中的路由器（router）部分设计的，路由器负责将输入的 token 路由到最合适的专家（expert）层。
在 MoE 模型中，每个输入 token 可能被路由到多个专家，但通常只有一个专家层会被激活。为了确保路由器能够稳定地工作并产生高质量的输出，作者引入了 Router z-loss。这个损失函数的目标是鼓励路由器产生较小的logits 值，因为较大的 logits 值在 softmax 激活函数中会导致较大的梯度，这可能会引起训练不稳定。
6.2 专家如何学习？
ST-MoE 的研究者们发现，encorder 中不同的专家倾向于专注于特定类型的 token 或浅层概念。例如，某些专家可能专门处理标点符号，而其他专家则专注于专有名词等。与此相反，decorder 中的专家通常具有较低的专业化程度。此外，研究者们还对这一模型进行了多语言训练。尽管人们可能会预期每个专家处理一种特定语言，但实际上并非如此。由于 token 路由和负载均衡的机制，没有任何专家被特定配置以专门处理某一特定语言。
6.3 专家的数量对预训练有何影响？
增加更多专家可以提升处理样本的效率和加速模型的运算速度，但这些优势随着专家数量的增加而递减 (尤其是当专家数量达到 256 或 512 之后更为明显)。同时，这也意味着在推理过程中，需要更多的显存来加载整个模型。值得注意的是，Switch Transformers 的研究表明，其在大规模模型中的特性在小规模模型下也同样适用，即便是每层仅包含 2、4 或 8 个专家。
6.4 Fine-Tuning MoE 模型
稠密模型和稀疏模型在过拟合的动态表现上存在显著差异。稀疏模型更易于出现过拟合现象，因此在处理这些模型时，尝试更强的内部正则化措施是有益的，比如使用更高比例的 dropout。例如，我们可以为稠密层设定一个较低的 dropout 率，而为稀疏层设置一个更高的 dropout 率，以此来优化模型性能。
在 Fine-Tuning 过程中是否使用辅助损失是一个需要决策的问题。ST-MoE 的作者尝试关闭辅助损失，发现即使高达 11% 的 token 被丢弃，模型的质量也没有显著受到影响。token 丢弃可能是一种正则化形式，有助于防止过拟合。
实验观察到，在相同的预训练 PPL 下，稀疏模型在下游任务中的表现不如对应的稠密模型，特别是在理解任务 (如 SuperGLUE) 上。另一方面，对于知识密集型任务 (如 TriviaQA)，稀疏模型的表现异常出色。作者还观察到，在Fine-Tuning 过程中，较少的专家的数量有助于改善性能。另一个关于泛化问题确认的发现是，模型在小型任务上表现较差，但在大型任务上表现良好。
一种可行的 Fine-Tuning 策略是尝试冻结所有非专家层的权重。实践中，这会导致性能大幅下降，我们可以尝试相反的方法：仅冻结 MoE 层的参数。实验结果显示，这种方法几乎与更新所有参数的效果相当。这种做法可以加速 Fine-Tuning 过程，并降低显存需求。
在 Fine-Tuning MoE 时还需要考虑的一个问题是，它们有需要特殊设置的超参数，例如，稀疏模型往往更适合使用较小的 batch size 和较高的学习率，这样可以获得更好的训练效果。
七、开源 MoE 模型
目前已经有一些开源的 MoE 大模型。国内的 MoE 大模型则是最近 DeepSeek 团队开源的 DeepSeekMoE，模型、代码、论文均已同步发布。
- 模型下载：https://huggingface.co/deepseek-ai
- 微调代码：https://github.com/deepseek-ai/DeepSeek-MoE
- 技术报告：https://github.com/deepseek-ai/DeepSeek-MoE/blob/main/DeepSeekMoE.pdf
此外还有一些国外的开源 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
下面是开源了模型，但是没有开源代码：
- Switch Transformers (Google)：基于 T5 的 MoE，专家数量从 8 到 2048。最大的模型有 1.6 万亿个参数。
- NLLB MoE (Meta)：NLLB 翻译模型的一个 MoE 变体。
- OpenMoE：社区对基于 Llama 的模型的 MoE 尝试。
- Mixtral 8x7B (Mistral)：一个性能超越了 Llama 2 70B 的高质量 MoE，并且具有更快的推理速度。此外，还发布了一个经过指令微调的模型。
图 1.Mod-Squad: 专家和任务互相选择。MoE ViT: 所有专家都被所有任务使用。
下面简单介绍下该文章。
模型结构
图 2.Mod-Squad: 将专家组 (mixture-of-expert) 插入到 Vision Transformer.
如图 2 所示， Mod-Squad 的结构就是将 Mixture-of-expert (MoE) 引入 Vision Transformer (ViT)。MoE 是一种机器学习模型，其中多个专家组成了一个混合模型。每个专家都是一个独立的模型，并且每个模型对于不同的输入有不同的贡献。最后，所有专家的贡献被加权并组合在一起以得到最终的输出。这种方法的优势在于它可以根据输入图像的内容动态地选择最佳的专家并且控制计算量。
之前的 MoE 模型收敛后，可以根据不同图片使用不同的专家，但是针对某个任务，模型会收敛到倾向于使用全部专家。Mod-Squad 可以做到让模型针对图片来使用不同的专家，并且模型可以在收敛后，达到一个任务只使用一部分专家的状态。接下来介绍这是怎么实现的。
最大化专家和任务之间的 mutual information
本文提出了一个任务和专家的联合概率模型来优化专家 E 和任务 T 之间的分配。这个概率模型会用来计算专家和任务之间的 mutual information，并作为额外的损失函数来优化 MoE 里的权重网络。Mutual information 公式如下，E 和 T 的概率可以由 MoE 里的权重网络得到，具体可以参见论文。
最大化任务和专家之间的 mutual information 之后，模型就可以让专家和任务拥有稀疏且非常强的依赖关系，如图 3 所示。最左边的就是 Mod-Squad 的任务使用专家频率。可以看出，Mod-Squad 的任务和专家之间有着更稀疏但尖锐的频率。
图 3. 任务使用不同专家的频率图对比。横轴是不同的专家，纵轴是不同的 task，颜色深代表更高的使用频率。Mod-Squad 的频率图更加稀疏且尖锐。
这个任务和专家之间稀疏且非常强依赖关系的好处就是：
1. 相近的任务倾向于使用同一个专家；
2. 专家倾向于被一组正相关的任务使用；
3. 模型的容量被全部使用，但每个任务只使用部分容量，可以根据任务调整使用容量；
4. 可以针对特定任务从多任务大模型中提取出单任务小模型，并具有和大模型一样的性能。这个特性能用于从超大多任务模型中提取出单任务小模型。
根据任务之间分享专家的频率，模型还可以算出任务之间的相似性，如下图所示。可以看出，偏 3D 的任务之间更倾向于使用相同专家，因此更加相似。
实验部分
Mod-Squad 可以在不损失精度的情况下针对单一任务进行剪枝，下图纵轴是性能，横轴是参数量。
在大数据集 Taskonomy 上也有很大的提升，可以看到，Mod-Squad 比单纯的 MTL 平均高了 2.8 个点，并且在剪枝以后保持着一样的性能。
在 PASCAL-Context 上跟其他方法的对比，Mod-Squad 比其他 MoE 方法平均高出了接近两个点。
总结
本文系统性地介绍了混合专家模型（MoE），主要介绍了针对 MoE 的高效训练方法，以及如何提升训练和 Fine-Tuning 的效果。现在我们回答下开篇提出的三个问题。
第一个问题：MoE 为什么能够实现在低成本下训练更大的模型。
这主要是因为稀疏路由的原因，每个 token 只会选择 top-k 个专家进行计算。同时可以使用模型并行、专家并行和数据并行，优化 MoE 的训练效率。而负载均衡损失可提升每个 device 的利用率。
第二个问题：MoE 如何解决训练稳定性问题？
可以通过混合精度训练、更小的参数初始化，以及 Router z-loss 提升训练的稳定性。
第三个问题：MoE 如何解决 Fine-Tuning 过程中的过拟合问题？
可以通过更大的 dropout （主要针对 expert）、更大的学习率、更小的 batch size。目前看到的主要是预训练的优化，针对 Fine-Tuning 的优化主要是一些常规的手段。
## 原理
在阅读DeepSpeed-Megatron / Megatron / Fairscale / Tutel等开源框架代码后总结而成，最终选择以DeepSpeed实现原理为主线，其余框架实现原理为辅线进行讲解。
关于MoE并行，我体感最深的一点是，它的定义不像tp/dp/pp这类典型的并行方式那样清晰明确，表现在：
- 对MoE并行相关的术语定义，各个论文/代码鱼龙混杂，有时根本不给定义直接甩出名词，有时相同的名词甚至代表不同的含义，给人造成极大困扰（举个最简单的例子，说说什么叫EP？）
- 对MoE并行的代码实现，不同框架间存在较多差异，有些框架还做了前提假设，导致在阅读源码过程中难以做知识迁移。
总结起来，MoE相关的开源资料，都透露着一股浓浓的禅意，这就好像你和老板的日常对话：你：“老板，这事急吗？” 老板：“如急，懂否？” 你：“如懂。” 所以参悟的最好办法，就是去看源码实践了。
本文是在阅读DeepSpeed-Megatron / Megatron / Fairscale / Tutel等开源框架代码后总结而成，最终选择以DeepSpeed实现原理为主线，其余框架实现原理为辅线进行讲解，全文架构如下：
- 第一部分，介绍选择DeepSpeed实现方式的原因
- 第二部分，介绍以Gshard为代表的MoE模型架构。如果你不想了解MoE分布式训练，只想知道MoE模型长什么样，是如何运作的，可以只看这部分
- 第三部分，介绍MoE并行训练中的分布式初始化。阅读本章需要对Megatron混合并行原理和Megatron源码架构有了解。
- 第四部分，源码解读。这个放在下篇中做讲解
一、为什么选择DeepSpeed-Megatron
引入MoE并行训练的框架有很多，为什么这篇文章选择用DeepSpeed来做讲解呢？主要原因如下：
- 使用DeepSpeed做讲解时，可将其MoE实现作为主线，其余框架MoE实现作为分支，这样方便在讲解主线的同时，引入分支进行对比。之所以DeepSpeed能成为主线，是因为它的MoE代码是一个“大杂汇”，例如，它的MoE初始化设置借鉴了Megatron，它的MoE-Layer架构借鉴了Fairscale，它的MoE优化方式借鉴了Tutel等，在这些借鉴之上，它引入了自己的一些改进。所以使用DeepSpeed，更方便我们做对比讲解。
- DeepSpeed的MoE模型架构是Gshard。作为最早将MoE应用在Transformer上的模型，Gshard提出的框架和思想一直影响至今。后续我们看到的很多LLM MoE的架构改进，其实都是在Gshard的这一套逻辑上做的迭代，比如loss改造、topKexpert的选择，稀疏矩阵计算优化等等。所以从Gshard入手，更利于我们对基础的把握。
二、Gshard架构
2.1 直觉上理解MoE设计
从架构图中我们可以发现，MoE其实就是将Transformer中的FFN层替换成了MoE-layer，其中每个MoE-Layer由一个gate和若干个experts组成。这里gate和每个expert都可以理解成是nn.linear形式的神经网络。
这样设计的直觉是什么呢？
expert：术业有专攻。假设我的输入数据是“我爱吃炸鸡”，在原始的Transformer中，我们把这5个token送去一个FFN层做处理。但是现在我们发现这句话从结构上可以拆成“主语-我”，“谓语-爱吃”，“宾语-炸鸡”，秉持着术业有专攻的原则，我把原来的1个FFN拆分成若干个expert，分别用来单独解析“主语”，“谓语”，“宾语”，这样可能会达到更好的效果。gate：那么我怎么知道要把哪个token送去哪个expert呢？很简单，我再训练一个gate神经网络，让它帮我判断就好了。
当然，这里并不是说expert就是用来解析主谓宾，只是举一个例子说明：不同token代表的含义不一样，因此我们可以用不同expert来对它们做解析。除了训练上也许能达到更好的效果外，MoE还能帮助我们在扩大模型规模的同时保证计算量是非线性增加的（因为每个token只用过topK个expert，不用过全量expert），这也是我们说MoE-layer是稀疏层的原因。
最后需要注意的是，在之前的表述中，我们说expert是从FFN层转变而来的，这很容易让人错理解成expert就是对FFN的平均切分，实际上你可以任意指定每个expert的大小，每个expert甚至可以>=原来单个FFN层，这并不会改变MoE的核心思想：token只发去部分expert时的计算量会小于它发去所有expert的计算量。
接下来，我们来看上图中MoE的各部分细节，下文中所有的符号遵从Gshard论文表示。
2.2 输入数据
首先，所有tokens正常过Attention层得到MoE-layer的输入，我们记输入数据的尺寸为(S, M)，其中:
- S : 输入batch中的token数量，例如图中S=8
- M: token_embedding维度
需要注意的是，我们一般是以batch的形式组织输入数据的（图中batch_size = 1），假设Attention层输入数据的维度是(batch_size, seq_len, M)，那么有S = batch_size * seq_len
2.3 Gate
接下来，我们就要使用线形层Gate帮助我们判断token应该送去哪个expert了。在别的MoE架构中，Gate有时也被称为Router（路由）。Gate的尺寸大小为(M, E)，其中E表示expert的数量。
输入数据(S, M)过Gate(M, E)后，得到prob数据(S, E)，它的含义是：每个token去向每个expert的概率。
由于在Gshard中我们使用的是top2Expert，因此对每个token，我们只关心它概率最大的两个expert。在图中，我们用深色表示最大概率，浅色表示次大概率。例如对token0来说，它被送去expert0的概率最大，被送去expert1的概率次大。
好，现在既然知道每个token的top2Expert了，是不是就可以直接发送了呢？别急，我们先来看看Expert的架构。
2.4 Expert与溢出处理
我们知道，token发去expert的概率不是我们能控制的，在实际操作中，可能某些expert接收到了好多token，而某些expert接收的token寥寥无几，我们管这种现象叫expert负载不均。这种情况不仅不符合我们MoE的设计初衷（术业有专攻），还影响计算效率（例如引起分布式训练中各卡通讯时的负载不均），所以我们急需想办法缓解这种问题，Gshard就提出了以下几种解决办法：
（1） capacity和capacity factor
在图中，你会看到一个叫Expert buffer的东西，这是什么呢？
在上文中我们提到，有些expert可能接收到非常多的token，为了缓解这个问题，我们可以给每个expert设置一个容量值（capacity），如果当前这个expert接收到的token数已经超过了容量，那么它就不再接收token了，此时我们称这个多出来的token为溢出(overflow)。
那么容量应该怎么设置呢？在我们的例子中，一共有8个token和4个expert，在理想的负载均衡的情况下，每个expert应该接收8/4 = 2个token，考虑到这里采用的是top2Expert，因此最终每个expert接收的token上限最好是(8/4)*2 = 4，这也是我们图中expert buffer的长度。
回到图中的例子上来，我们发现t0和t1都正常发去top2Expert上了。但是对于t6，它的2nd expert已经装满了；对于t7，它的1st和2nd expert都满了。所以t6和t7都发生了溢出。那么我们要怎么处理溢出的情况？别着急，我们马上来看。
（2） Random Routing
对于一个token，我们一定要把它发到top2Expert上吗？
从直觉上对于每个token，我们最好将其以100%的概率发去1st expert；但是对于它的2nd expert，我们可以不100%发送，而是以一定的概率（例如从uniform(0,1)中随机抽取一个数p，将其作为概率）发送，这样不就能节省对expert capacity的消耗，从而更好利用每个expert吗？这也是Gshard论文中提出的方法。
而我们下文要讲的DeepSpeed代码，在这一块又做了稍微不同的处理：以图中t0为例，1st expert它是肯定要发去的。但是在选择2nd expert时，它做了一些加噪处理：对产出的每个概率（更确切地说是logit），它从某种分布中采样4个噪声，加在这4个logit上，然后mask掉1st expert位置的logit，再从剩下3个logit中找到最大的作为其2nd Expert。
现在我们已经选出最终的top2Expert，我们再回到没有加噪时的4个概率上，取出相应位置的概率，做normalize计算：
回到上面的问题，token发生溢出时，要怎么办呢？
- 如果只有单个expert溢出，那么就把另一个expert的权重值为1，然后正常参与加权计算（如图中t6）
- 如果2个expert都溢出，那么该token就不经过任何expert，直接通过残差连接的方式，原样发去下一层的Attention上（如图中t7）
（3） Auxiliary Loss
除了capacity和random routing外，Gshard还通过增加一项辅助损失函数(Auxiliary Loss)来尽量保证Expert的负载均衡，其定义如下：
2.5 Zero Padding和Drop tokens
写到这里，我们稍微总结下：
- 首先，我们有一串过Attention层后的token序列
- 我们通过Gate，计算每个token去往每个expert的概率
- 我们希望不同expert处理的token数尽量均衡，所以我们同时采取三方面优化：
- Capacity: 为每个expert设置capacity（expert buffer），限制它能处理的最大token数量，多出来的token算为溢出，在top2Expert都溢出的情况下，该token会被直接发去下一层attention。
- Random Routing: 每个token一定会被发去1st Expert，在此基础上我们通过random routing加噪的方式，重新选出2nd expert。在做完capacity + random routing后，我们最终确认了每个token要发去的top2expert和其对应的权重，通过加权计算的方式，确认Moe-Layer最终的输出结果。
- Auxiliary Loss：添加辅助损失函数，对expert负载不均的情况做进一步惩罚。
到这里，Gshard MoE的核心架构内容我们就说完了，最后再提2点：
（1）Zero padding
我们上述的优化方法，只能“缓解”负载不均，而不能保证解决负载不均。也就是说，存在一些Expert，它的Expert buffer没有填满，这可怎么办呢？
最直接的方法，就是在没有buffer中空出来的位置，用0向量填充，我们称为Zero padding。更具体地说，最终每个expert上的输入数据维度为(E, C, M)，其中C表示capacity。0填充的好处是，我们保证每个expert上要处理的输入数据维度是一样的，这有利于硬件层面的后续处理（例如多卡通讯间的负载均衡等）。
（2）Drop tokens
我们知道，当发生溢出情况时，不是所有token都会被expert正常处理的，我们称这种对溢出的操作为drop tokens。如果被drop掉的tokens数量太多，也是一种信息损失（它们都没经过任何expert解析），我们当然可以通过调整capacity来缓解这个问题，但过大的capacity会引起更严重的zero padding问题（影响到矩阵的稀疏程度），所以这也是后续一些MoE模型架构侧重的优化。
2.6 伪代码
现在我们将整个过程以伪代码的形式写出（大家注意看注释细节）。这里在Gshard论文提供的伪代码上，按照deepspeed的实现方式做了些修正。
# -------------------------------------------------------------------------------------
# 1.通过gate，计算每个token去到每个expert的概率
#
# 【input】：Attention层输出的一整个batch的token，其尺寸为(seq_len, batch_size, M),
# 其中M表示token_embedding
# 【reshaped_input】：由input做reshape而来，尺寸为(S, M), 其中S = seq_len * batch_size
# 【Wg】: gate的权重，尺寸为(M, E)，其中E表示MoE-layer层的专家总数
# 【gates】: 每个token去到每个expert的概率，尺寸为(S, E)
# -------------------------------------------------------------------------------------
M = input.shape[-1]
reshape_input = input.reshape(-1, M)
gates = softmax(enisum("SM, ME -> SE"), reshape_input, Wg)
# -------------------------------------------------------------------------------------
# 2. 确定每个token最终要去的top2Expert，并返回对应的weight和mask
#
# 【combine_weights】：尺寸为(S, E, C)，其中C表示capacity（Expert buffer）。
# 表示对每个token（S）而言，它对每个专家（E）的weight，而这个weight按照
# 该token在buffer中的位置（C）存放，不是目标位置的地方则用0填充
# 例如图中token1，它将被发送到expert0和expert2，且它在expert0的buffer中排在
# 1号位置，在expert2中排在0号位置，那么token1的combine_weights就是：
# [[0., p0, 0., 0.],
# [0. , 0., 0., 0.],
# [p2, 0., 0., 0.],
# [0., 0., 0., 0.]]
# 最后再复习一下weight和gates所表示的prob的区别：前者是在后者基础上，
# 做了random + normalize，确定最终的top2Expert后返回的对应结果
#
# 【dispatch_mask】： 尺寸为（S，E，C），它等于combine_weights.bool(), 也就是对combine_weights
# 为0的地方设为False，为1的地方设为True。
# dispatch_mask后续将被用在zero padding上
# -------------------------------------------------------------------------------------
# (S, E, C) (S, E, C)
combine_weights, dispatch_mask = Top2Gating(gates)
# -------------------------------------------------------------------------------------
# 3. 将输入数据按照expert的顺序排好，为下一步送去expert计算做准备（很重要）
#
# 【dispatch_mask】：尺寸为（S, E, C），定义参见2
# 【reshape_input】：尺寸为(S, M)，定义参见1
# 【dispatched_expert_input】：本步的输出结果，表示按专家排序好的输入数据，尺寸为(E, C, M)
# 这个结果表示，每个专家（E）的buffer（C）下要处理的token_embedding（M），
# 例如dispatched_expert_input[0]就表示expert0 buffer中的情况
# 注意：
# （1）每个专家buffer中的token是按顺序排列好的，
# 回到图中的例子，expert0 buffer下0号位置排的是token0，
# 3号位置排的是token6，以此类推。dispatch_mask就起到了维护这种顺序的作用
# （2）当对应专家接收的token数不足buffer长度C时，不足的地方用0向量填充。
# -------------------------------------------------------------------------------------
dispatched_expert_input = einsum("SEC, SM -> ECM", dispatched_mask, reshape_input)
# -------------------------------------------------------------------------------------
# 4. 将排序好的input送入expert进行计算。
# 同正常的FFN层一样，每个expert也由2个线形层Wi, Wo组成
# 【dispatched_expert_input】：按专家顺序和专家buffer中的token顺序排好的输入数据，
# 尺寸为(E, C, M)，具体定义见3
# 【Wi】：experts的Wi层，尺寸为(E，M, H)，
# 【Wo】：experts的Wo层，尺寸为(E, H, M)
# 【expert_outputs】：experts的输出结果，不含加权处理，尺寸为（E, C, M）
# -------------------------------------------------------------------------------------
h = enisum("ECM, EMH -> ECH", dispatched_expert_input, Wi)
h = relu(h)
expert_outputs = enisum("ECH, EHM -> ECM", h, Wo)
# -------------------------------------------------------------------------------------
# 5. 最后，进行加权计算，得到最终MoE-layer层的输出
# -------------------------------------------------------------------------------------
outputs = enisum("SEC, ECM -> SM", combine_weights, expert_outputs)
outputs_reshape = outputs.reshape(input.shape) # 从(S, M)变成(seq_len, batch_size, M)再特别说明几点
（1）enisum的作用
- enisum在这里泛指我们自定义的某几种矩阵计算方式，enisum中诸如"SEC, SM -> ECM"只是用来表示输入数据和输出数据的维度，并不表示两个输入矩阵就一定是按照SEC和SM这样的尺寸直接相乘（我们肯定要对输入数据做些例如unsqeeze()，reshape()之类的操作才能把它们正确乘起来，得到想要的结果）
- enisum使得我们在矩阵计算的同时，能维持token和expert的顺序。你可能在阅读伪代码的过程中已经感受到，维持“顺序”是一件很重要的事，例如token在专家buffer中的顺序，各个专家间的排序等。为什么维持顺序很重要呢？因为一个batch里有很多token，我们将其发往不同的expert做计算后，输出结果的顺序肯定是打乱的，所以需要通过一种方式追踪顺序，把token permute回正常的位置再输入下一层Attention。在这里我们通过自定义的矩阵计算方式，巧妙维护住这种顺序，这样我们就不需要额外建索引表之类的来查找了。
在后文对deepspeed的源码解读中，我们会看到enisum的具体定义。不过这块不是源码解读的讲述重点（毕竟也只是矩阵计算而已）。对这块有兴趣的朋友，可以自己攥一些数据跑跑代码，研究它的运作原理。不感兴趣的朋友，只要记住输入输出的尺寸及各自含义即可。
（2）将输入数据按照expert的顺序排好
大家可以特别关注下伪代码步骤3中的操作，这个操作将有利于后续专家并行组之间的通讯。
三、MoE并行训练
正如前文所说，阅读本章需要了解Megatron混合并行原理，并掌握其代码中“分布式初始化”部分的相关知识。
当我刚开始研究MoE时，总会看到类似EP + DP，EP + TP + DP这样并行方式的缩写，例如DeepSpeed官方文档中所描述的。最开始我对这个符号的理解是：非MoE层的部分采取DP或DP+TP的方式；而MoE层的部分采取一种叫EP的新方式。然而当我把这样的理解代入代码中时，却发现有些部分难以解释。
摸索了一段时间后，我才发现不管是EP + DP，EP + TP +DP等等，它们都在特指MoE层的并行方式；而对non-MoE层，你采取什么样的并行方式，是不在这些并行符号的表示范围中的。
我们以EP + DP，EP + TP + DP这两种方式为例，来看看如何对MoE模型做分布式初始化。
3.1 EP + DP
如上图，我们先来看一个例子。在本例中，我们共有16块gpu：
- 对于non-moe的部分，采取tp + dp并行
- 对于moe部分，采取ep + dp并行。
（1）Non-MoE: tp + dp
non-moe的部分指模型中Attention、word embedding层等。tp + pp 的并行方式我们也很熟悉了，根据图中所示，我们有：
tp_world_size = 2
tp_groups = [
[g0, g1],[g2, g3], [g4, g5], [g6, g7],
[g8, g9],[g10, g11],[g12, g13],[g14, g15]
]
dp_world_size = 8
dp_groups = [
[g0, g2, g4, g6, g8, g10, g12, g14],
[g1, g3, g5, g7, g9, g11, g13, g15]
]
pp_world_size = 1（2）MoE: ep + dp
当我们安顿好non-moe的部分后，我们就可以开始考虑要怎么安排MoE层了，这里我们先给出划分方法，然后再对其进行详细解释：
ep_world_size = 4
ep_groups = [
[g0, g1, g2, g3 ],
[g4, g5, g6, g7],
[g8, g9, g10, g11],
[g12, g13, g14, g15]
]
ep_dp_world_size = 4
ep_dp_groups = [
[g0, g4, g8, g12],
[g1, g5, g9, g13],
[g2, g6, g10, g14],
[g3, g7, g11, g15]
]
ep_tp_world_size = 1还记得前面我们说MoE层采用的是EP + DP并行吗？那么这里的EP和DP的定义到底是什么呢？
假设我们每个MoE层有若干个专家（我们统称其为一套专家），现在我们想把这一套专家分布排列到gpu上，最直觉的做法就是：我们先定好要用几块GPU装下一套专家（EP），进而我们就能确认全局上共有多少套专家副本在跑（DP）。通过这种简单的方式，我们就做好了EP + DP形式的MoE层初始化。
回到我们的例子中，一共16块GPU：
- ep_world_size = 4：表示我们希望用4块GPU装下一套完整的专家。确定这个数值后，我们就能确认ep_groups
- local_expert_num：expert_num / ep_world_size，其中expert_num表示每层专家的总数。假设每层专家数量是4，那么1块gpu上就放一个专家；假设每层专家数量是8，那么1块gpu上就放2个专家。所以图中的e0等符号并不绝对表示这里只有1个专家，只是对local_expert的统称。
- ep_dp_world_size：类比于non-MoE层，MoE层同样也有数据并行的概念。例如图中[g0, g4, g8, g12]上都维护着e0，所以它们构成一个ep_dp_group。这个group的作用是当我们在计算bwd时，它们之间是需要做梯度的allreduce通讯的，我们会在下文详细图解这一点。另外需要注意的是，构成ep_dp_group的条件不仅是e相同，还需要每个e吃的batch的数据不同（类比于一个普通的dp_group，组内的每张卡吃的是不同的小batch）。现在你可能无法具象化感受这点，我们在后文将ep+tp+dp并行的时候再细说。
- ep_tp_world_size：类比于non-MoE层，MoE层同样也有张量并行的概念，即一个专家可以纵向切割成若干份，本例中我们不对专家做tp操作，在后文我们会详细来看做了tp操作的相关例子。
额外再说明两点：
- 你可能发现上面诸如
ep_dp_world_size这样的符号有点陌生，因为你并没有在相关论文/代码中看到过它。这是因为如本文开篇所说，不同框架对MoE并行的相关概念定义鱼龙混杂，这也是最令我痛苦的点。所以这里我自定义了一套符号，不管是什么框架，我都会把它的定义映射到这套符号上来。 - 以图中的e0来举例，我们再强调两点：首先，如上文所说，它不绝对表示1个专家，只是对local_expert的统称。其次，它不绝对表示1个MoE层的专家，它表示所有MoE层放在这块卡上的专家统称。
相信现在你对EP+DP的分布式设置有了初步认识了（这里我们特意举了non-MoE是tp+dp，而不是单纯dp的例子，来说明ep+dp这个并行定义是专门针对MoE部分的），但你可能对那些并行group的作用还不能有具象体会。现在让我们来给模型喂点数据，看看在1个FWD和BWD过程中，这些group都做了什么通讯吧！
（3）FWD与BWD过程
如图，三角形表示1个batch的数据，这里我们共喂给模型8个batch。每个tp组内的输入数据一致，所以它们对应的三角形颜色也相同。
好，让我们牢记分布式并行的使命：分布式训练效果应与单卡（假设这个单卡能装下一个完整的模型）的训练效果一致。放到我们的例子里，16卡吃8个小batch做完FWD+BWD后的结果，应该与单卡吃下由这8个小batch组成的大batch的结果一致。
现在开始做FWD与BWD，过程如下图：
- 在FWD中，数据先过non-MoE（Attention）层，由于一个tp组内每块卡的输出也是一致的，因此三角形颜色的分布没有改变。我们把三角形移动到对应的non-MoE分块下，表示在整个FWD中对应的non-MoE分块见过的batch。
- 继续做FWD，现在数据来到了MoE层，我们前面说过，每块卡上数据的维度是(E, C, M)，即我们已经计算好token和专家的对应关系，我们只需在ep_group内做all2all通讯，将token发送去对应的专家即可，这就是ep_group的作用。all2all通讯的细节我们放在后面说，这里只需记住在all2all通讯后，ep_group内每个专家见过的batch有了改变，例如对e0，现在它见过了蓝色和橘色两个batch的数据。每个专家计算完自己的结果后，再通过all2all的方式，将对应的token计算结果还给ep_group内的各gpu，然后继续mon-MoE->MoE的步骤，知道FWD完毕。
- 做完了FWD，进入BWD。我们首先来到MoE部分，以e0为例，根据分布式训练使命，我们应该allreduce 8个batch的梯度结果，用来更新e0。欸那这8个batch在哪里呢？当然是在图中的ep_dp_group内！所以在BWD过程中，我们对ep_dp_group中e0的梯度做allreduce，用来更新e0。现在，你是不是更好理解ep_group的作用了！
- 继续做BWD，数据来到了non-MoE部分，这块对梯度的通讯我们在Megatron解析中已经讲了很多，这里就不再说明了。
总结一下针对MoE部分的重点：
- 在FWD中，ep_group进行all2all通讯，将token发去对应的专家做计算，并将计算结果取回。
- 在BWD中，ep_dp_group进行AllReduce通讯梯度，用于更新对应的专家的参数。
对于这种在non-MoE部分采用tp，在MoE部分不采用tp的设计，在代码实现上有几个点要注意。举例来说，对non-MoE来说，[g0, g1]是一个tp组，我们要对这两块卡的输出做AllReduce。但是对MoE部分而言，[g0, g1]的输出是不需要做AllReduce的。看过Megatron代码的朋友，应该能想起这块相关的操作在RowParallelLinear/ColumnParallelLinear模块下，所以在deepspeed中，通过传入一个enable_expert_tensor_parallelism=True/False的参数来做相关调整，这点我们放在源码解读篇中说。
在一些代码框架（例如Megatron）中，为了多复用已有的并行方式，少做修改，一般都会做些强硬限制：例如MoE的mp（tp与pp）层面的并行设置须与non-MoE的mp设置保持一致，即如果non-MoE做了tp切分，MoE也必须以同样的方式做tp切分，在此基础上再去安排MoE的ep/ep_data等等并行。在这样的限制下，如果non-MoE采用dp，那么MoE只能用ep+dp；如果non-MoE采用tp+dp，那么MoE只能采用ep+tp+dp；欸发现了没有！这是不是和你我对ep+tp+dp这个符号表示的初印象很像？即tp+dp是non-MoE的并行方式，ep是MoE的并行方式。所以这样的理解，在某些代码框架上是通的，但是到别的更为灵活的代码实现上，就产生矛盾了。这也为什么我在本章开头说明，最好统一把这个符号理解成是对MoE部分并行方式的描述。
对于non-MoE只采用dp，MoE采用ep+dp的设计，比较简单，这里我们就不多说了，大家可以自己画画。
3.2 All2All通讯
在3.1中，我们说过每张卡进MoE前的输入数据尺寸为(E, C, M)，其中E表示expert_num，C表示capacity，M表示token_embedding。在每个ep_group内，我们通过all2all通讯将token发去指定的expert做计算，再通过all2all通讯将计算结果返回。现在我们来介绍all2all的细节。
（1）基础All2All
我们先来看基础All2All是怎么做的，再来看deepspeed改进后的All2All是怎么做的。
图中的MP表示的就是TP（在deepspeed的语系中，MP=TP），图中相关的分布式group为：
- tp_group: [[g0, g1], [g2,g3]]
- ep_group: [[g0,g1,g2,g3]]，也就意味四张卡上分别存着e0, e1, e2, e3
我们先来看最左侧的图，它描绘了数据刚过完non-MoE(Attention)层后的结果。因为tp组间做了AllReduce，所以g0和g1上存的数据完全一致（ABCD），g2和g3上存的数据完全一致(EFGH)。我们以[g0,g1]为例，因为有4个专家，所以图中将数据分为ABCD四份，每一份的维度为(C, M)，四份总维度为(E, C, M)。也就是说A的数据要发去e0，B的数据要发去e1，以此类推。
我们再来看中间的图，它描绘了ep_group内首次做all2all的过程，这个过程的目的是将token发去对应的expert上进行计算。你是否已经发现，all2all就相当于做了一次矩阵转置（对比一下左侧图和中间图的数据块排布）？因此通过All2All，我们就让数据块去到了它对应的位置：AE去e0，BF去e1，以此类推。而为了实现这种转置，我们必须提前对non-MoE做分块排序，让它按照要去的专家位置排好，现在你是不是能感受到排序的意义了？
最后来看右侧的图，它描绘了ep_group内第二次做all2all的过程，这个过程的目的是将MoE算完的token再返回给各卡，原理和上述一致。
一切都进行地很顺利，但我们能不能再做些优化呢？例如，属于一个tp组的g0和g1上存着一模一样的数据，在all2all的过程中是会被重复发送和计算的，我们能不能减少这种重复？
（2）改进All2All，理论版
为了避免tp组内的数据重复发送的问题，deepspeed在论文中提出了一种改进版的All2All算法，但值得一提的是，deepspeed在代码实现中可不是完全按照这种改进算法来的（手动狗头，虽然实现上还是借鉴了一些理论上的思路）。所以本节我们先来看deepspeed在理论上的改进，然后再来看它的实操。
咱们来看看这张图，确实是维持了deepspeed团队一贯的禅意风格：以培养读者悟性为宗旨，能不点破就不点破。所以我们也不要辜负他们的期望，努力地悟一悟吧！
deepspeed改进版all2all的核心宗旨是：既然你说tp组间的数据重复，那么我就在tp组间砍一刀，让tp组内的每块卡都维护不同的数据，不就不行了么？于是我把g0上的CD砍掉，g1上的AB砍掉，以此类推。这下完美了吧，每块卡上只有2块数据且不重复了！
但这样一来，all2all要怎么做呢？你现在每张卡上只有2个数据块，但是all2all group内一共4块卡，如果你想正常做all2all，必须卡上的数据块和卡数相同才行，否则就会出现问题。
所以，我们再想另一个办法：如果现在卡上只有两块数据了，那我如果把all2all group也一分为二，每个group内2张卡，一共两个all2all group，不就能解决这个问题了吗？那这个新的all2all group要怎么设呢？最简单的方法就是，tp rank相同的卡组成一个新的all2all group。例如对g0和g2，它们的tp_rank都是0，所以它们组成新的all2all组，g1和g3也是同理。所以现在，我们只需对[g0, g2] all2all，[g1, g3] all2all就可以了。这也是为什么在图中把CB和FG交换位置的原因。
那如果我想[g0,g3] all2all，[g1,g2] all2all，然后把CD, GH交换位置，那可以么？理论上是没问题的，但是代码上这样写就不太优雅了，比不上同一个tp rank间all2all来得简便。
知道了这个流程，我们再来看上面的图，是不是一目了然了？图中local transform为改进版all2all作准备，交换了数据；两个all2all就是token发送和接受的过程；最后再加一个all-gather，因为tp组内每块卡的输出要保持一致，这样才能进入接下来的non-MoE层继续计算。
（3）改进All2All，实操版
好，理论我们已经知道了。但你从改进的过程中肯定也发现了：原来是4卡all2all，现在是2卡all2all，这不就把我辛苦设置好的ep_group破坏掉了吗？如果真的这么操作，那对代码的改动肯定是比较大的。
所以，有没有办法还是4卡all2all，但是也避免发送重复数据呢？
当然有，大家想想我们每张卡的输出数据维度(E, C, M)，同个tp组内的所有卡这个(E, C, M)是一致的，此时你的脑子是不是灵光一现：如果我沿着C把数据切成C/tp_world_size份，每张卡只保留其中的一分，那么不就既能做到tp组内卡上的数据不重复的情况下复用原来的ep_group做all2all吗？在两次all2all后，我依然通过一个all-gather操作还原tp组内的完整数据，这不就行了吗？这就是deepspeed在代码实操中使用的方法。
deepspeed在代码中管这样的操作叫drop_tokens，大家注意和MoE理论部分所说的drop token区分开。
（4）改进版All2All使用注意事项
在上面改进版All2All的讲解中，你可能已发现一个重要的点：虽然non-MoE采取了tp，但是MoE却没有用tp。
答案是否定的，因为对于g0和g1来说，MoE层也被切开了，因此按照tp的方式，g0和g1上的输入各自过切开的MoE层后AllReduce的结果，才是最终结果。这时g0和g1上重复的输入数据是有实质意义的，我们不能drop。
deepspeed禅师，你看我们这样算悟了吗？
3.2 EP + DP + TP
（1）Non-MoE与MoE
现在我们再回到ep并行设置的例子上来，我们考虑对专家也做tp切分。实现这一点最简单的逻辑是让专家的tp切分方式和非专家的tp切分方式一致，这样省得我们再去对MoE层多写一套tp组的分布式设置。
non-MoE组的并行设置和3.1一致，这里不再介绍，我们来看看MoE组的并行设置：
ep_world_size = 4，保持和3.1一致
ep_groups = [
[g0, g2, g4, g6],
[g1, g3, g5, g7],
[g8, g10, g12, g14],
[g9, g11, g13, g15]
]
ep_dp_world_size = 2
ep_dp_groups = [
[g0, g8],
[g1, g9],
[g2, g10],
[g3, g11],
[g4, g12],
[g5, g13],
[g6, g14],
[g7, g15]
# 复用non-MoE tp相关的并行设置
ep_tp_world_size = 2
ep_tp_groups = [...]- ep_groups：在前面我们提过，每个ep_groups中的每个ep_group装下一套“完整的”专家，但我们也留了个坑，说明“完整”的含义需要留后讨论。现在我们就来看看这个坑。
- 在deepspeed中，虽然每个ep_group内的专家都是被纵向切开的，但只要它涉及到所有的专家，就认为它是“完整的”
- 在Megatron中，“完整”的含义就是参数完整，即g0~g7才被认为是一个ep_group。
定义不同，组内的通讯方式自然也不同。在deepspeed的定义下，ep_group内做的是all2all；在megatron的定义下，做的是ReduceScatter和AllGather（事实上Megatron的MoE实现就没有用到all2all）。这一点我们在源码讲解篇会来做比较。这里我先抛出我的结论：我认为deepspeed的ep_group设计是较好的，不仅操作上更符合直觉认知，还能避免重复数据发送（Megatron的实现方法会发送大量重复数据）。
（2）FWD与BWD
同样，我们通过描述1次FWD与BWD的过程，来说明各个group间是怎么运作的。
- 首先做FWD，数据过non-MoE层，因为采用了tp，所以相同tp组内的各卡输出一致，因此我们在图中将相同颜色的三角形放到对应的non-MoE参数块下。
- 继续做FWD，数据来到了MoE层。在例子中：
- ep_group [g0, g2, g4, g6]和ep_group[g1, g3, g5, g7]内都各做1次all2all，将token发给对应的expert进行计算。
- 计算完毕后，因为MoE层也是tp并行的，因此[g0, g1], [g2, g3], [g4, g5], [g6, g7]这几个tp组各自通过AllReduce取得完整的输出结果。
- 然后ep_group [g0, g2, g4, g6]和ep_group[g1, g3, g5, g7]再做一次all2all，把计算完毕的数据发送回去，以便进入下一个non-MoE->MoE操作。我们在图中把这一步骤里各卡维护的expert所见过的batch数据画在对应的expert下面。
- 开始做BWD，数据来到MoE层。同一个ep_dp_group内的参数需要对梯度做AllReduce，例如图中[g0, g8]，这个group维护了相同的e，每个e都各自吃过4个batch的数据，联合起来刚好构成全局上的8个batch（牢记前文分布式训练的使命）.
- 继续做BWD，数据来到non-MoE层，来到大家熟悉的领域了，tp+dp模式下的BWD，就不需要多说了吧。
3.3 PP去哪里了
在上面的例子中，我们见过了ep+tp+dp的混合，你可能想问，pp也是一种常见的并行方式，它去哪里了呢？
在MoE中，PP这个维度比较特殊，大部分论文和开源代码实践中，一般不考虑/不讨论再对MoE部分做pp切分。deepspeed中更是强制把PP维度设为1，并在代码注释里表示不支持PP（我猜这可能和deepspeed的zero实现有关，体感上可能会加剧通讯量，但我没做仔细研究，给不出确切答案）。我个人认为，如果tp+dp已经能满足显存限制的话，就不需要再引入pp将模型切得更碎了。同时在MoE模型中，你会发现non-MoE的模型副本数和MoE的模型副本数是不一致的。例如3.2的例子中，non-MoE有8个模型副本，但是MoE只有两个模型副本（g0~g7, g8~g15），却也能实现8个完整的non-MoE + MoE模型副本的分布式训练效果。从这一点上看tp+dp形式的训练方式已经基本够用了。
但并不是说不能使用pp，Megatron中就支持pp的使用，我们来看下引入pp后的情况（图中标题给错了，应该是EP+TP+DP+PP，太懒了没有改，大家凑合看）：
- non-MoE层采用tp + dp +pp并行（就是我们在Megatron解读中举的例子，是我们熟悉的味道）
- MoE层采用ep + tp +pp + dp，其中tp_group和tp_group直接复用non-MoE的设置
相信通过之前的讲解，大家已经能轻松看懂这张图了。这里只强调三点：
- ep_dp_groups = [[g0], [g1], [g2],...]，也就是每张卡单独组成了一个ep_dp_group（如果你觉得难以理解，就用前面说的每个e见过的batch来分析看看）
- ep_group = [[g0, g1, g2, g3], [g4, g5, g6, g7],...]，这个不难理解，想想上文说的Megatron对“完整的一套专家“的定义。
- 需要满足dp_world_size % ep_world_size == 0，事实上这也是Megatron并行设置的前置条件（不管你有没有使用pp，因为Megatron强制MoE复用non-MoE的并行配制，在此基础上再引入和ep相关的并行）。一般而言，world_size = tp_world_size * pp_world_szie * dp_world_szie，如果你的MoE层复用了non-MoE的tp和pp，那么ep_world_size只能在dp_world_size上做切割了。
好，关于MoE并行训练的模型架构和分布式配置，我们就介绍到这了。在下一篇中，我们通过源码讲解，和大家一起学习MoE并行训练的更多细节。
一、DeepSpeed MoE
1.1 执行脚本
执行脚本：Megatron-DeepSpeed/examples_deepspeed/MoE/ds_pretrain_gpt_1.3B_MoE128.sh
在deepspeed中，一共实现了两类MoE架构：
- 普通MoE：每个MoE层的专家数相等。因此你传入的num_expert可以是一个int
- PR-MoE：金字塔-残差连接型MoE（Pyramid-Residual MoE），这是deepspeed自己提出的创新模型，每个MoE层的专家数可以不等，因此你传入的应该是一个空格连接的字符串，即每层num_expert个数用空格隔开
这里我们只关注普通MoE。
这份sh脚本主要用于做参数配置，例如模型结构参数、分布式训练参数等等。这里主要关注分布式配置相关参数，其余可自行阅读。
megatron_options=" \
--override-opt_param-scheduler \
--adam-beta1 0.9 \
--adam-beta2 0.95 \
--tensor-model-parallel-size ${MP_SIZE} \
--moe-expert-parallel-size ${EP_PARALLEL_SIZE} \
--num-experts ${EP_SIZE} \
--moe-loss-coeff ${MLC} \
--moe-train-capacity-factor ${MOE_TRAIN_CAP_FACTOR} \
--moe-eval-capacity-factor ${MOE_EVAL_CAP_FACTOR} \
--moe-min-capacity ${MOE_MIN_CAP} \
--......- tensor-model-parallel-size：tp_world_size。由于在deepspeed moe实现中，默认pp_world_size = 1，则对于non-moe部分有dp_world_size = world_size / tp_world_size。
- moe-expert-parallel-size：ep_world_size。
- num-experts：每个MoE层专家数量。在普通MoE中，这是一个int（如128）；在PR-MoE中，其形式如"64 64 64 64 64 64 64 64 128 128"，表示每层的专家数量。当num_expert = 1时，说明不采用MoE架构（我们管这种叫dense model，同理采用MoE则称为sparse model）
对这些术语定义有疑惑的朋友，可以先看原理篇。
1.2 入口函数
入口函数：Megatron-DeepSpeed/pretrain_gpt.pyds_pretrain_gpt_1.3B_MoE128.sh脚本启动的入口函数位于Megatron-DeepSpeed/pretrain_gpt.py中，从这一步开始，整体代码就基本复用我们之前介绍过的Megatron框架了，只是在分布式环境和模型架构设计上会做些deepspeed特有的更改。所以再次建议大家在阅读本文前先了解Megatron源码。
1.3 分布式环境初始化
相关脚本：Megatron-DeepSpeed/megatron/initialize.py
在Megatron中我们讲过相关的内容（TODO：插入链接），这里就不赘述了。这份代码的主要作用就是将gpu划分成我们原理篇中说的各类group。
这里只提一个重要的点：deepspeed在这份脚本中，只对tp/pp/dp等group做了设置，对ep相关的group（例如原理篇中说的ep_group，ep_dp_group等）都没做相关设置。与之相反的是，Megatron在此处一次性把所有group都设置好了。
那么deepspeed在哪里设置ep相关的group呢？在“按分布式设置切割模型，并将模型搬运到gpu上”这一环节（我们马上就在1.4中说明）。为什么deepspeed要这么设计呢？如果你使用过deepspeed，你一定对deepspeed.initialize()这个api非常熟悉，它的作用是使用deepspeed的方式对传入的model做包装，以便在训练中能对model做各类我们所熟悉的deepspeed优化操作（例如zero显存优化技术等）。deepspeed moe也算是deepspeed特有的操作的一种，所以放在deepspeed.initialize()中做设置。我认为在初始化上，deepspeed和megatron的代码实现没有什么优劣之分，只是一种设计习惯而已。大家在读不同代码时能读懂相关内容即可。
1.4 模型切割
相关脚本：Megatron-DeepSpeed/megatron/training.py
熟悉Megatron的朋友应该能记起，在做完分布式初始化后，我们就按照初始化设计好的方式，切割我们的模型，并将其搬运到GPU上，我们来看这一步的核心函数setup_model_and_optimizer（只摘取关键代码做讲解）
def setup_model_and_optimizer(model_provider_func,
model_type,
no_wd_decay_cond=None,
scale_lr_cond=None,
lr_mult=1.0,
teacher=False,
data_post_process=None,
build_train_valid_test_datasets_provider=None):
"""Setup model and optimizer."""
args = get_args()
# 根据当前rank，切割好的模型
model = get_model(model_provider_func, model_type)
...
if args.deepspeed:
....
....
# 使用deepspeed做初始化
model, optimizer, args.deepspeed_dataloader, opt_param_scheduler = deepspeed.initialize(
model=model[0],
optimizer=optimizer,
args=args,
lr_scheduler=opt_param_scheduler,
training_data=train_ds,
mpu=mpu if args.no_pipeline_parallel else None,
config=args.deepspeed_config_dict,
)get_model：定义当前进程所属的gpu所要维护的模型块。我们在Megatron源码解读中专门写过一篇来介绍它，和MoE层相关的定义也由它来完成，下文我们会详细解读。deepspeed.initialize()：使用deepspeed来初始化定义好的模型，以便模型在训练中能使用deepspeed相关的优化操作。我们在1.3中说过，deepspeed MoE实现中，对MoE层分布式设置的部分（即设置各种ep并行组）都由deepspeed.initialize()来完成，我们详细来看下它是怎么做的，具体代码在DeepSpeed/deepspeed/runtime/engine.py，注意这里换了一个仓库，从Megatron-DeepSpeed仓库换到DeepSpeed仓库下，核心代码如下：
# Set deepspeed parallelism spec. for the model including expert parallelism
for _, module in self.module.named_modules():
if hasattr(module, 'set_deepspeed_parallelism'):
module.set_deepspeed_parallelism(self._config.use_data_before_expert_parallel_)这段代码的意思是，如果当前你定义的模型（module）中含有属性set_deepspeed_parallelism，说明这个模型会用到deepspeed自定义的分布式设置方法，这时我们对模型执行set_deepspeed_parallelism()方法就可以完成相关初始化设置了。现在看不懂也没关系，后文我们在介绍MoE层切块模型架构定义的时候，会同时介绍这个方法详细的代码实现。
好，到这里我们稍微总结下deepspeed MoE的代码流程：
- 在分布式设置环节，deepspeed MoE基本复用了Megatron分布式设置逻辑，但并没有对ep相关的并行组做设置
- 在模型切割环节，deepspeed MoE自定义了MoE部分的模型架构，并通过deepspeed.initialize方法对模型设置ep相关的并行组。
1.5 MoELayer
（1）整体框架
相关脚本：Megatron-DeepSpeed/megatron/model/transformer.py
现在，我们可以来看deepspeed是怎么定义它的MoE层了。
先简单回顾一下之前在Megatron源码解读时，画过的模型架构部分的层级关系：
图中每个红虚线框表示一块gpu/一个进程定义的模型块。我们知道MoE层其实是对原来MLP层的替换，所以主要改动应该在ParallelMLP相关的代码下。我们以ParallelTransformerLayer为入口，来看一步步看下细节。
class ParallelTransformerLayer(MegatronModule):
"""A single transformer layer.
Transformer layer takes input with size [s, b, h] and returns an
output of the same size.
"""
def __init__(self, config,
layer_number, layer_type=LayerType.encoder,
self_attn_mask_type=AttnMaskType.padding,
drop_path_rate=0., num_experts=1):
......
# -------------------------------------------------------------------
# Attention (non-MoE)部分
# -------------------------------------------------------------------
self.self_attention = ParallelAttention(
config,
layer_number,
attention_type=AttnType.self_attn,
attn_mask_type=self_attn_mask_type)
......
# -------------------------------------------------------------------
# MLP（MoE）部分，提供三种MLP定义方式
# 1、SwitchMLP：Megatron设计的MoE架构
# 2、普通MLP：当num_expert=1时，说明该模型不是MoE，只是一个普通的dense MLP层
# 3、deepspeed MoE：deepspeed自定义的MoE架构
# -------------------------------------------------------------------
self.num_experts = num_experts
if args.num_experts_switch is not None:
self.mlp = SwitchMLP(config) # 1. Megatron-LM's MoE
else:
if self.num_experts <= 1: # 2. dense, not MoE
self.mlp = ParallelMLP(config)
else: # 3. DeepSpeed's MoE
# enable_expert_tensor_parallelism：表示是否要对专家做tp切分
enable_expert_tensor_parallelism = args.enable_expert_tensor_parallelism
self.mlp = MoE(args.hidden_size, # token_embedding
# 定义单个专家
ParallelMLP(config,
moe=True,
enable_expert_tensor_parallelism=enable_expert_tensor_parallelism),
num_experts=self.num_experts, # 每层专家总数
ep_size=args.moe_expert_parallel_size, # ep_world_size
k=args.topk, # topKEpert中的K，deepspeed使用Gshard模型，一般而言K=2，但也提供了K=1时方法
use_residual=(args.mlp_type == 'residual'), # deepspeed自创的PR-MoE架构中提出的方法
capacity_factor=args.moe_train_capacity_factor, # train过程中expert的capacity factor
eval_capacity_factor=args.moe_eval_capacity_factor, # eval过程中expert的capacity factor
min_capacity=args.moe_min_capacity, # expert最少需要达到的capacity
drop_tokens=args.moe_token_dropping, # 是否需要对tokens做溢出处理
use_tutel=args.use_tutel, # 是否需要使用tutel路由优化方法
enable_expert_tensor_parallelism=enable_expert_tensor_parallelism # 是否需要对专家层做tp切分
)
......我们重点关注由deepspeed实现的MoE（Megatron实现的MoE我们放在后文说）。详细的解释都在代码注释中了，这里再额外提几点：
ParallelMLP()：定义单个expert的模型架构。我们知道1个expert其实也是1个mlp模块，一层MoE由若干个这样的expert/mlp组成。如果不对expert做tp切分，那么这里定义的就是完整的expert架构；如果对expert做tp切分，那么这里定义的就是对原始expert纵向切成tp_world_size块后某一块expert的架构。
回到上图中，你将图中的ParallelMLP层替换成一个MoE层，再在其中装入若干块ParallelMLP，就能将其改装成MoE模型下的分布式架构。
use_residual：前文说过，deepspeed提出了一种叫PR-MoE的模型架构，这个函数就是指代其中的“R”。PR-MoE架构图见下，首先，Pyramid（P）允许为每层设置不同数量的专家。Residual（R）则表示一个固定的mlp模块（也可以理解成一个固定的专家），即所有的token都一定会经过这个mlp模块解析，除此以外再为每个token在剩下的专家中选出top1。deepspeed这样做的原因是：他们认为在top2Expert策略中，2nd expert是对1st expert的纠偏，也就是1st expert选得不一定正确，我们需要用2nd expert弥补一些信息回来。那既然总要做纠偏，为何我不干脆把1st expert固定下来，然后只对2nd expert做选择呢？这就是R的意义。
drop_tokens：是否需要做溢出处理。在原理篇中我们讲过溢出的定义和溢出的处理方式，这里就不多说了。注意，这需要与all2all通讯中在同一个tp组内对输入做drop tokens的操作区分开。use_tutel：是否需要用tutel的路由优化方法。这个说来话长，可以单开一篇文章来讲了。在本文中可以暂时不关心，感兴趣的朋友可以自行阅读。
下面，我们先来看单个expert(ParallelMLP)如何定义，再来看如何将其组装为一个MoE层。
（2）定义单个expert模型架构
单个expert模型架构的入口为ParallelMLP()，我们来看看它的代码。
class ParallelMLP(MegatronModule):
"""MLP.
MLP will take the input with h hidden state, project it to 4*h
hidden dimension, perform nonlinear transformation, and project the
state back into h hidden dimension.
"""
def __init__(self, config, moe=False, enable_expert_tensor_parallelism=False):
......
# --------------------------------------------------------------------
# self.dense_h_to_4h：Wi，尺寸大小(h, 4h/tp_world_size)
# --------------------------------------------------------------------
self.dense_h_to_4h = tensor_parallel.ColumnParallelLinear(
config.hidden_size,
ffn_hidden_size,
config=config,
init_method=config.init_method,
bias=self.add_bias,
gather_output=False,
skip_bias_add=True,
moe=moe,
enable_expert_tensor_parallelism=enable_expert_tensor_parallelism
)
......
# --------------------------------------------------------------------
# self.dense_4h_to_h, Wo, 尺寸大小为(4h/tp_world_size, h)
# --------------------------------------------------------------------
self.dense_4h_to_h = tensor_parallel.RowParallelLinear(
config.ffn_hidden_size,
config.hidden_size,
config=config,
init_method=config.output_layer_init_method,
bias=self.add_bias,
input_is_parallel=True,
moe=moe,
enable_expert_tensor_parallelism=enable_expert_tensor_parallelism
)
def forward(self, hidden_states):
# --------------------------------------------------------------------
# 输入数据过Wi层
# [s, b, 4h/tp_word_size]
# --------------------------------------------------------------------
intermediate_parallel, bias_parallel = self.dense_h_to_4h(hidden_states)
......
# --------------------------------------------------------------------
# Wi层输出数据过Wo层
# 在对expert采取tp切分的情况下，这里的输出需要在tp_group内做AllReduce
# [s, b, h]
# --------------------------------------------------------------------
output, output_bias = self.dense_4h_to_h(intermediate_parallel)
return output, output_bias在原理篇我们说过，和普通MLP层一样，每个expert也是由Wi和Wo两个nn.linear层组成。其中Wi采用ColumnParallelLinear，Wo采用RowParallelLinear。如果每个专家是按tp维度切开的，那么数据过Wo层后的输出需要在对应tp_group内做AllReduce。
乍看之下这段代码似乎和没有MoE时长得差不多，这是因为入参moe和enable_expert_tensor_parallelism都是在ColumnParallelLinear和RowParallelLinear中起作用的，现在我们马上来看这两块的细节。
（a）ColumnParallelLinear
相关脚本：Megatron-DeepSpeed/megatron/core/tensor_parallel/layers.py
这里把列切割的原理再写一遍，方便大家对照着看：
f和g是两个共轭算子，可理解为两个torch.autograd.Function类。在这个类下，我们可以根据需要重写forward和backward方法。f: forward中，直接copy输入；backward中，对梯度做AllReduce。在代码里定义为class _CopyToModelParallelRegion(torch.autograd.Function)g: forward中，all-gather输出；backward中，对梯度做split（每张卡经过all-gather已有完整的Y了，因此以Y为起点计算梯度后，沿着列做split就可得到Y1和Y2的梯度）。在代码里定义为class _GatherFromModelParallelRegion(torch.autograd.Function)
class ColumnParallelLinear(torch.nn.Module):
"""Linear layer with column parallelism.
......
"""
def __init__(self, input_size, output_size, *,
config: ModelParallelConfig,
init_method: Callable,
bias=True, gather_output=False, stride=1,
keep_master_weight_for_test=False,
skip_bias_add=False,
skip_weight_param_allocation: bool=False,
moe=False, enable_expert_tensor_parallelism=False):
torch.nn.Module.__init__(self)
# Keep input parameters
self.input_size = input_size
self.output_size = output_size
self.gather_output = gather_output
# ---------------------------------------------------------------
# 判断相关module是否表expert
# 1、如果是moe，但不希望对expert做tp处理，则强制设tp_world_size = 1
# 2、其余情况（非moe，或者是moe且希望对expert做tp处理），则复用non-Moe
# 的tp_world_size
# ---------------------------------------------------------------
if moe and (not enable_expert_tensor_parallelism):
world_size = 1
self.is_expert_without_slicing = True
else:
world_size = get_tensor_model_parallel_world_size()
self.is_expert_without_slicing = False
self.output_size_per_partition = divide(output_size, world_size)
self.skip_bias_add = skip_bias_add
self.config = config
......
def forward(self,
input_: torch.Tensor,
weight: Optional[torch.Tensor] = None):
"""Forward of ColumnParallelLinear
......
"""
......
# ------------------------------------------------------------------------
# 定义f算子
# 1、如果expert未采用tp切分，则对expert的forward和backward方法不需要做任何修改
# 2、如果expert采用tp切分，则需要对expert的forward和backward方法做修改
# 具体为forward时直接copy input，backward时allreduce梯度
# 可参考下放示意图
# ------------------------------------------------------------------------
if self.async_tensor_model_parallel_allreduce or \
self.sequence_parallel or \
self.is_expert_without_slicing: # non-expert only tensor parallelism
input_parallel = input_
else:
input_parallel = copy_to_tensor_model_parallel_region(input_)
# Matrix multiply.
output_parallel = linear_with_grad_accumulation_and_async_allreduce(
input=input_parallel,
weight=weight,
bias=bias,
gradient_accumulation_fusion=self.gradient_accumulation_fusion,
async_grad_allreduce=self.async_tensor_model_parallel_allreduce,
sequence_parallel=self.sequence_parallel
)
# ------------------------------------------------------------------------
# 定义g算子
# 同理根据expert是否做了tp切分，决定要不要更改forward和backward方法
# 需要注意，当你对单个expert做tp切分时，不管你的gather_output是True/False,
# 单个expert的输出结果一定会在同个tp组内强制做allReduce
# ------------------------------------------------------------------------
if self.gather_output and not self.is_expert_without_slicing:
# All-gather across the partitions.
assert not self.sequence_parallel
output = gather_from_tensor_model_parallel_region(output_parallel)
else:
output = output_parallel
output_bias = self.bias if self.skip_bias_add else None
return output, output_bias（b）RowParallelLinear
这块关于expert的相关操作和ColumnParallelLinear差不多，也是针对“单个expert是否做了tp”这个条件做分类处理，具体代码留给读者自行阅读。
好，到这里我们对单个expert的模型架构简单总结下：
- 单个expert模型架构复用ParallelMLP定义，同时我们可根据需要决定单个expert是否要做tp切割
- 单个expert不做tp切割，则不需要重写它的forward和backward方法
- 单个expert做tp切割，则需要重写它的forward和backward方法，这和之前Megatron中做列切割和行切割的相关操作一致。
- 单个expert做tp切割时，数据过单个expert后的输出结果一定会在同个tp_group内做AllReduce
关于总结的第4点，我们再额外提几点：
在之前原理篇的介绍中，我们提过FWD中数据过MoE层时各个group的通讯情况，具体如下：
数据来到了MoE层。在例子中：
- ep_group [g0, g2, g4, g6]和ep_group[g1, g3, g5, g7]内都各做1次all2all，将token发给对应的expert进行计算。
- 计算完毕后，因为MoE层也是tp并行的，因此[g0, g1], [g2, g3], [g4, g5], [g6, g7]这几个tp组各自通过AllReduce取得完整的输出结果。
- 然后ep_group [g0, g2, g4, g6]和ep_group[g1, g3, g5, g7]再做一次all2all，把计算完毕的数据发送回去，以便进入下一个non-MoE->MoE操作。我们在图中把这一步骤里各卡维护的expert所见过的batch数据画在对应的expert下面。
有朋友问：第2步和第3步可以换下顺序吗？即我先完成所有All2All的步骤，然后再对结果做AllReduce行不行呢？从逻辑上来说是没问题的，但是从我们代码实现上来说，由于单个expert也是一个ParallelMLP模块，因此在expert也采用tp的情况下，做完FWD后数据一定是AllReduce的，因为这是定义在ParallelMLP的forward方法中的。这就是上面总结中强调第4点的含义。
（3）定义MoE层整体架构
相关脚本：DeepSpeed/deepspeed/moe/layer.py
我们已经知道了单个expert的定义，现在我们来看如何用单个expert组装起完整的MoE层。
class MoE(torch.nn.Module):
"""Initialize an MoE layer.
Arguments:
- 【hidden_size】 (int): token_embedding
- 【expert】 (torch.nn.Module): 单个expert架构，属于ParallMLP类
- 【num_experts】: 每层expert总数
- 【ep_size】 (int, optional): default=1, ep_world_size
- 【k】 (int, optional): default=1, topKexpert中的K，只支持K=1或2
- 【capacity_factor】 (float, optional): default=1.0, train步骤的容量因子
- 【eval_capacity_factor】 (float, optional): default=1.0, eval步骤的容量因子
- 【min_capacity】 (int, optional): default=4, 每个专家最小的容量值
- 【use_residual】 (bool, optional): default=False, 用于表示该层是否是一个residual expert层 (https://arxiv.org/abs/2201.05596) layer.
- 【noisy_gate_policy】 (str, optional): default=None, noisy gate policy（加噪策略）, valid options are 'Jitter', 'RSample' or 'None'.
- 【drop_tokens】 (bool, optional): default=True, whether to drop tokens - (setting to False is equivalent to infinite capacity).
- 【use_rts】 (bool, optional): default=True, whether to use Random Token Selection.
- 【use_tutel】 (bool, optional): default=False, whether to use Tutel optimizations (if installed).
- 【enable_expert_tensor_parallelism】 (bool, optional): default=False, 是否对expert做tp切分
"""
def __init__(self,
hidden_size,
expert,
num_experts=1,
ep_size=1,
k=1,
capacity_factor=1.,
eval_capacity_factor=1.,
min_capacity=4,
use_residual=False,
noisy_gate_policy: typing.Optional[str] = None,
drop_tokens: bool = True,
use_rts=True,
use_tutel: bool = False,
enable_expert_tensor_parallelism: bool = False):
super(MoE, self).__init__()
self.use_residual = use_residual
self.enable_expert_tensor_parallelism = enable_expert_tensor_parallelism
assert num_experts % ep_size == 0, f"Number of experts ({num_experts}) should be divisible by expert parallel size ({ep_size})"
self.ep_size = ep_size
self.expert_group_name = f"ep_size_{self.ep_size}"
self.num_experts = num_experts
# 单块gpu上需要存放的expert数量
self.num_local_experts = num_experts // self.ep_size
log_dist(
f'Creating MoE layer with num_experts: {num_experts} | num_local_experts: {self.num_local_experts} | expert_parallel_size: {self.ep_size}',
[0])
assert noisy_gate_policy is None or noisy_gate_policy in ['None', 'Jitter', 'RSample'], \
'Unsupported noisy_gate_policy: ' + noisy_gate_policy
# -------------------------------------------------------------------------
# 定义一个MoE层上所有的expert。见下面Experts类定义（很重要）
# -------------------------------------------------------------------------
experts = Experts(expert, self.num_local_experts, self.expert_group_name)
# -------------------------------------------------------------------------
# 定义MoE层
# -------------------------------------------------------------------------
self.deepspeed_moe = MOELayer(TopKGate(hidden_size, num_experts, k, capacity_factor, eval_capacity_factor,
min_capacity, noisy_gate_policy, drop_tokens, use_rts),
experts,
self.expert_group_name,
self.ep_size,
self.num_local_experts,
use_tutel=use_tutel)
if self.use_residual:
self.mlp = expert
# coefficient is used for weighted sum of the output of expert and mlp
self.coefficient = torch.nn.Linear(hidden_size, 2)
def set_deepspeed_parallelism(self, use_data_before_expert_parallel_=False):
"""
ep相关分布式设置。
如前文所说，我们在deepspeed.initialize()中，如果检测到一个module拥有
set_deepspeed_parallelism属性，则我们就对它执行相关的分布式设置操作
"""
self._create_process_groups(use_data_before_expert_parallel_=use_data_before_expert_parallel_)
def _create_process_groups(self, use_data_before_expert_parallel_=False):
"""
ep相关分布式设置
"""
# ----------------------------------------------------------------------------
# 如果当前还未做ep相关分布式设置，那么就先做设置
# ----------------------------------------------------------------------------
if self.expert_group_name not in groups._get_expert_parallel_group_dict():
print(f"No existing process group found, creating a new group named: {self.expert_group_name}")
# ----------------------------------------------------------------------------
# 1、当你没使用Megatron分布式并行，或者你使用了Megatron但又不想对expert组tp切分
# 那么你就按EP + DP的方式设置ep相关group
# ----------------------------------------------------------------------------
if (groups.mpu is None) or (not self.enable_expert_tensor_parallelism):
# Condition 1 - no groups.mpu means no tensor parallelism
# Condition 2 - disabling expert tensor parallelism on purpose
groups._create_expert_and_data_parallel(
self.ep_size, use_data_before_expert_parallel_=use_data_before_expert_parallel_)
# ----------------------------------------------------------------------------
# 2、其余情况则使用EP + DP + TP方式
# ----------------------------------------------------------------------------
else:
# expert tensor parallelism is enabled
groups._create_expert_data_and_model_parallel(
self.ep_size, mpu=groups.mpu, use_data_before_expert_parallel_=use_data_before_expert_parallel_)
# ----------------------------------------------------------------------------
# 在做完ep相关分布式设置的情况下，为当前进程所属的MoE层显示设置ep_group
# 这样就可以在ep_group内做all2all通讯
# 如果不显示设置ep_group，则默认是对所有gpu卡（world_size）做all2all
# ----------------------------------------------------------------------------
self.deepspeed_moe._set_ep_group(groups._get_expert_parallel_group(self.expert_group_name))
def forward(self, hidden_states, used_token=None):
""" MoE forward
Arguments:
hidden_states (Tensor): input to the layer
used_token (Tensor, optional): default: None, mask only used tokens
Returns:
A tuple including output, gate loss, and expert count.
* output (Tensor): output of the model
* l_aux (Tensor): gate loss value
* exp_counts (int): expert count
"""
output = self.deepspeed_moe(hidden_states, used_token)
if self.use_residual:
# Residual MoE
output_mlp = self.mlp(hidden_states)
if type(output_mlp) is tuple:
output_mlp = output_mlp[0] # Ignore the bias term for now
coef = self.coefficient(hidden_states)
coef = torch.nn.functional.softmax(coef, dim=-1)
output = output * coef[..., 0:1] + output_mlp * coef[..., 1:]
return output, self.deepspeed_moe.l_aux, self.deepspeed_moe.exp_counts
class Experts(nn.Module):
"""
相关脚本：DeepSpeed/deepspeed/moe/experts.py
定义一个MoE层上所有的Expert
"""
def __init__(self, expert: nn.Module, num_local_experts: int = 1, expert_group_name: Optional[str] = None) -> None:
super(Experts, self).__init__()
# ----------------------------------------------------------------------------
# 每块gpu上共num_local_experts个expert
# ----------------------------------------------------------------------------
self.deepspeed_experts = nn.ModuleList([copy.deepcopy(expert) for _ in range(num_local_experts)])
self.num_local_experts = num_local_experts
# TODO: revisit allreduce for moe.gate...
for expert in self.deepspeed_experts:
# TODO: Create param groups to handle expert + data case (e.g. param.group = moe_group)
for param in expert.parameters():
param.allreduce = False
param.group_name = expert_group_name
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
"""
我们知道，在分发去experts前，每张卡上的输出结果为(E, C, M)，其中E=该MoE层专家总数，
C = capacity，M = token embedding。
设ep_world_size = G, num_local_experts = e, 则易知E = G * e
对于All2All通讯，你可以理解成对于ep_group内的每张卡，都将数据沿着E维度切成G块后，再进行通讯。
（保证每张卡上的数据块数量 = ep_world_size，这样All2All通讯才不会出错）
因此发送完毕后，每张卡上的数据可以又表示为(G*e, C, M)
进一步在正式把数据喂给这张卡上维护的experts前，我们可以把数据reshape成(G, e, C, M)的形式。
此时如果我们沿着e维度将数据切分为e个chunck，则一个chunk对应一个local_expert，再次实现了token
和local expert间一一对应的关系
"""
# -------------------------------------------------------------------
# 将input做切分后，将分块分别喂给该gpu上维护的若干个expert
# inputs尺寸：(G, e, C, M)，
# G = ep_world_size
# e = num_local_experts，满足G*e = E，E为该层expert总数
# C = expert capacity,
# M = token embedding
# chunk_input: 沿着e维度切分inputs，方便各块input喂给该gpu上对应的各个expert
# -------------------------------------------------------------------
chunks = inputs.chunk(self.num_local_experts, dim=1)
expert_outputs: List[torch.Tensor] = []
for chunk, expert in zip(chunks, self.deepspeed_experts):
# out尺寸：（G, C, M）
out = expert(chunk)
if isinstance(out, tuple):
out = out[0] # Ignore the bias term for now
expert_outputs += [out]
# concat后最终out尺寸: (G, e, C, M)
return torch.cat(expert_outputs, dim=1)细节都在注释中，我们不再赘述。这里只整理一下定义一个MoE层的整体流程
- 首先，我们定义好了单个expert模型架构（ParallelMLP）
- 然后，鉴于一张卡上可能不止维护1个expert（num_local_experts = num_experts // ep_world_size），我们需要定义这张卡上expert的集合Experts（nn.ModuleList，见代码细节）
- 最后，我们需要一个TopKGate策略，来帮助token选择expert
- 将以上内容组装成一个MOELayer
（4）MOELayer与TopKGate
MOELayer
相关脚本：DeepSpeed/deepspeed/moe/sharded_moe.py
阅读本节时可以配合原理篇2.6部分的伪代码进行阅读（TODO：插入链接），deepspeed在MOELayer的实现上基本完全照搬了fairscale的实现方式，细节都在注释中，这里不赘述。
class MOELayer(Base):
"""MOELayer module which implements MixtureOfExperts as described in Gshard_.
::
gate = TopKGate(model_dim, num_experts)
moe = MOELayer(gate, expert)
output = moe(input)
l_aux = moe.l_aux
.. Gshard_: https://arxiv.org/pdf/2006.16668.pdf
Args:
gate (torch.nn.Module):
gate network
expert (torch.nn.Module):
expert network
"""
def __init__(self,
gate: Module,
experts: Module,
ep_group_name,
ep_size,
num_local_experts: int,
use_tutel: bool = False) -> None:
super().__init__()
# -------------------------------------------------------------------------
# TopKGate类，用来决定token的分法策略（细节见下文代码解读）
# -------------------------------------------------------------------------
self.gate = gate
# -------------------------------------------------------------------------
# 当前进程所属的gpu上维护的所有experts，nn.ModuleList[ParallelMLP()]
# -------------------------------------------------------------------------
self.experts = experts
# -------------------------------------------------------------------------
# 当前进程所属的ep_group，为None时表示所有的gpu构成一个ep_group
# 当执行_set_ep_group方法时，可以自定义ep_group(参见MoE类下_create_process_groups方法)
# -------------------------------------------------------------------------
self.ep_group = None
# -------------------------------------------------------------------------
# 当前进程所属的ep_group的ep_world_size
# -------------------------------------------------------------------------
self.ep_size = ep_size
# -------------------------------------------------------------------------
# 当前进程所属的ep_group的名字
# -------------------------------------------------------------------------
self.ep_group_name = ep_group_name
# -------------------------------------------------------------------------
# 当前进程所属的gpu上所维护的experts数量，它即为self.experts中维护的experts数量
# -------------------------------------------------------------------------
self.num_local_experts = num_local_experts
# -------------------------------------------------------------------------
# 一些用于衡量MoE计算过程的时间技术器，可以忽略不看（后面代码中我都用省略号替换了）
# -------------------------------------------------------------------------
self.time_falltoall = 0.0
self.time_salltoall = 0.0
self.time_moe = 0.0
self.timers = SynchronizedWallClockTimer()
self.wall_clock_breakdown = False
# 是否使用tutel做路由优化（tutel的实现细节我们不在本文中分析）
self.use_tutel = use_tutel and TUTEL_INSTALLED and gate.k == 1
if self.use_tutel:
logger.info('Using Tutel optimizations.')
elif use_tutel and not TUTEL_INSTALLED:
logger.warning("Tutel optimization requested but not installed. "
"Proceeding without Tutel.")
elif use_tutel and TUTEL_INSTALLED and gate.k != 1:
logger.warning("To enable Tutel optimization, use top-1 instead of top-2 gate. "
"Proceeding without Tutel.")
def _set_ep_group(self, ep_group):
self.ep_group = ep_group
def forward(self, *input: Tensor, **kwargs: Any) -> Tensor:
......
# -------------------------------------------------------------------------
# 1. 对input做reshape
# 注意入参中input前面带*号，意味着传入的input是一个tuple，一般是一个二元组
# input[0]是我们真正要做计算的batch数据，其尺寸为(seq_len, batch_size, M)
# input[1]是掩码数据，其尺寸为(seq_len*batch_size)，有时在计算MoE结果时，我们相对
# 某些token做mask，使其不参与计算，就可以把mask数据装在这里。由于这一策略不常用，
# 因此在本文中我们不关注这块
#
# reshaped_input尺寸为(S, M)，其中S = seq_len * batch_size
# -------------------------------------------------------------------------
d_model = input[0].shape[-1]
# Initial implementation -> Reshape into S tokens by dropping sequence dimension.
# Reshape into G groups so that each group can distribute tokens equally
# group_size = kwargs['group_size'] if 'group_size' in kwargs.keys() else 1
reshaped_input = input[0].reshape(-1, d_model)
# -------------------------------------------------------------------------
# 是否使用Tutel做路由优化（不是本文讲解内容）
# -------------------------------------------------------------------------
if self.use_tutel:
self.l_aux, C, E, indices_, locations_, gates_, self.exp_counts = self.gate(reshaped_input, input[1], True)
S, M = reshaped_input.size(0), reshaped_input.size(1)
if not hasattr(self, '_tutel_dispatcher'):
self._tutel_dispatcher = tutel_moe.fast_dispatcher(E, C, M, dispatch_dtype=reshaped_input.dtype)
self._tutel_dispatcher.update(indices_, locations_, gates_, capacity=C)
dispatched_input = self._tutel_dispatcher.encode(reshaped_input)
# -------------------------------------------------------------------------
# 2. 使用自定义的Gshard gate，确定token的分法策略
# （对以下输出结果解释有疑惑的，都可以参考原理篇2.6节伪代码讲解，有详细的图例）
# gate：TopKGate类，后文做解读
# l_aux: 辅助损失函数值
# combine_weights: 尺寸为(S, E, C)，表示对每个token（S）而言，它对每个专家（E）的weight，
# 而这个weight按照该token在buffer中的位置（C）存放，不是目标位置的地方则用0填充
# dispatch_mask： 它等于combine_weights.bool(), 也就是对combine_weights
# 为0的地方设为False，为1的地方设为True。
# dispatch_mask后续将被用在zero padding上
# -------------------------------------------------------------------------
else:
# 确认token分法策略
self.l_aux, combine_weights, dispatch_mask, self.exp_counts = self.gate(reshaped_input, input[1])
# -------------------------------------------------------------------------
# 3. 将输入数据按照expert的顺序排好，并做zero padding,
# 为下一步送去expert计算做准备（很重要）
# dispatched_input: 尺寸为(E, C, M)，
# 表示每个专家（E）的buffer（C）下要处理的token_embedding（M），
# 当对应专家接收的token数不足buffer长度C时，不足的地方用0向量填充。
# -------------------------------------------------------------------------
dispatched_input = einsum("sec,sm->ecm", dispatch_mask.type_as(input[0]), reshaped_input)
......
# -------------------------------------------------------------------------
# 4. 当expert不采用tp切分，而non-MoE部分采用tp切分时，为避免数据重复发送，需要对
# 同一个tp组内的tokens做去重（见原理篇3.2节）
# -------------------------------------------------------------------------
if groups._get_expert_model_parallel_world_size() == 1:
# If the non-expert is tensor-parallel, it will create
# duplicate tokens on the tensor-parallel ranks.
# Since our experts are not tensor-parallel, these duplicates
# need to be dropped to ensure correctness.
# this also doubles up as a communication optimization as we are
# reducing the all-to-all communication volume.
dispatched_input = drop_tokens(dispatched_input, dim=1)
# -------------------------------------------------------------------------
# 5. 第一次All2All：将token发给对应的expert
# dispatched_input尺寸为(E, C, M)，又可以写成(G*e, C, M)，
# 其中G=ep_world_size, e = num_local_experts
# 在将它正式喂给expert前，把它reshape成(G, e, C, M)
# (参见1.5(3) Experts类注释)
# -------------------------------------------------------------------------
dispatched_input = _AllToAll.apply(self.ep_group, dispatched_input)
......
# Re-shape after all-to-all: ECM -> GeCM
dispatched_input = dispatched_input.reshape(self.ep_size, self.num_local_experts, -1, d_model)
# -------------------------------------------------------------------------
# 6. 将token喂给expert计算
# expert_output尺寸为(G, e, C, M)
# -------------------------------------------------------------------------
expert_output = self.experts(dispatched_input)
......
# -------------------------------------------------------------------------
# 7. 第二次All2All：将算好的token返回给产出它的gpu
# expert_output为(G, e, C, M)，即此时这张卡上维护的token过MoE的结果，
# 是由它从ep_group（G）内所有expert(e)的结果汇总而来的
# -------------------------------------------------------------------------
expert_output = _AllToAll.apply(self.ep_group, expert_output)
......
# Re-shape back: GeCM -> ECM
expert_output = expert_output.reshape(self.ep_size * self.num_local_experts, -1, d_model)
# -------------------------------------------------------------------------
# 8. 如果之前在tp组内做过数据去重处理，这里要把数据all-gather回来
# （参见原理篇3.2）
# -------------------------------------------------------------------------
if groups._get_expert_model_parallel_world_size() == 1:
# the dropped duplicate tokens need to be gathered on each
# tensor parallel rank again for the tensor-parallel
# non-expert of the next layer.
expert_output = gather_tokens(expert_output, dim=1)
# -------------------------------------------------------------------------
# 9. 使用combine_weights进行加权计算
# combined_output尺寸为(S, M），其中S = seq_len * batch_size
# -------------------------------------------------------------------------
if self.use_tutel:
combined_output = self._tutel_dispatcher.decode(expert_output.view(E * C, M))
else:
combined_output = einsum("sec,ecm->sm", combine_weights.type_as(input[0]), expert_output)
# 最终输出a尺寸为：(seq_len, batch_size, M)
a = combined_output.reshape(input[0].shape)
......
return aTopKGate
TopKGate可以理解成是token的分发策略，它决定每个token要发去哪些expert，以及每个token在这些expert上的weight。
TopKGate实际其实就是一堆矩阵计算（包括在原理篇中我们提过的一些enisum算法），没有太多内容可写，大家可以对照着原理篇第二部分Gshard的架构解读（TODO：插入链接）来阅读代码。这里有一个经验之谈：虽然TopKGate的主要思想不复杂，但是实现起来有些弯弯绕绕的，建议大家直接捏一些假数据，跑一遍相关代码，把一些关键变量打印出来，这样更有利于解读。
这里只额外提一点，是我认为deepspeed在实践中可能有点问题的地方，那就是对gate模块的初始化，我们先来看下相关代码实践：
class TopKGate(Module):
wg: torch.nn.Linear
def __init__(self,
model_dim: int,
num_experts: int,
k: int = 1,
capacity_factor: float = 1.0,
eval_capacity_factor: float = 1.0,
min_capacity: int = 8,
noisy_gate_policy: Optional[str] = None,
drop_tokens: bool = True,
use_rts: bool = True) -> None:
super().__init__()
# Only top-1 and top-2 are supported at the moment.
if k != 1 and k != 2:
raise ValueError('Only top-1 and top-2 gatings are supported.')
# 定义gate层架构
self.wg = torch.nn.Linear(model_dim, num_experts, bias=False).float()我们知道，gate的作用是计算每个token分发去各个expert上的prob，在deepspeed的设计中，每张卡上都会有一个gate模块。不难理解这个所有卡上这个gate模块应该是一模一样的，否则算法就会出现逻辑上的问题。
而为了保证每卡上的gate一模一样，它们应该采取同样的初始化方式（例如采用相同初始化方法和初始化种子），但是但看deepspeed的这行代码实现，似乎并没有保障这一点。由于我没有完整阅读过deepspeed初始化的全部代码，不确定deepspeed是否在后续的包装过程中保证了这一点，因此觉得这里“可能”有些问题。
好！到这里为止，deepspeed moe实现的核心代码部分我们就讲完了，建议大家按这个流程将源码串起来读一遍，加深理解。接下来我们讲解Megatron moe的核心实现。
二、Megatron MoE
2.1 分布式环境初始化
相关脚本：Megatron-LM/megatron/core/parallel_state.py
在1.3中我们说过，deepspeed把ep相关的初始化设置放到模型切割之后，megatron则在模型切割之前一次性把包括ep在内的初始化都做完。我们先来回顾下原理篇中给出的一个megatron分布式设置的方式：
- non-MoE层采用tp + dp +pp并行
- MoE层采用ep + tp +pp + dp，其中tp_group和tp_group直接复用non-MoE的设置
对照着这张图，我们来看代码:
......
# -------------------------------------------------------------------------------
# ep_group
# [[g0, g1, g2, g3], [g4, g5, g6, g7],[g8, g9, g10, g11], [g12, g13, g14, g15]]
# 假设当前进程为0，同样都是ep_world_size = 2的情况下：
# - 在Megatron中，当前进程对应的ep_group为[g0, g1, g2, g3]，
# ep_group内通讯为ReduceScatter/AllGather(下文会细说)
# - 在deepspeed中，当前进程对应的ep_group为[g0, g2],
# ep_group内通讯为All2All
# -------------------------------------------------------------------------------
_TENSOR_AND_EXPERT_PARALLEL_GROUP = None
# -------------------------------------------------------------------------------
# ep_dp_group
# [[g0], [g1], [g2],...], 每张卡单独组成1个ep_dp_group
# 对ep_dp_group定义有疑问的朋友，可先阅读原理篇第三部分
# deepspeed和megatron对ep_dp_group的定义是一致的
# -------------------------------------------------------------------------------
_DATA_MODULO_EXPERT_PARALLEL_GROUP = None # ep_dp_group
......
# -------------------------------------------------------------------------------
# 在megatron中，强制MoE部分复用non-MoE部分的tp和pp group，那么ep部分的划分就只能在
# dp维度上进行，所以需要强制data_parallel_size % expert_model_parallel_size == 0，
# 这里expert_model_parallel_size就是ep_world_size（注意在图例中ep_world_size是2不是4！）
# -------------------------------------------------------------------------------
if data_parallel_size % expert_model_parallel_size != 0:
raise RuntimeError(
f"data_parallel_size ({data_parallel_size}) is not divisible by expert_model_parallel_size "
)
# -------------------------------------------------------------------------------
# 具体划分_TENSOR_AND_EXPERT_PARALLEL_GROUP和_TENSOR_AND_EXPERT_PARALLEL_GROUP的代码
# 我们就不特别列出了，大家知道了这些group的定义后阅读起来应该不困难，
# 大家可以多尝试一些切分方式，跑一跑看看megatron会给出什么样的group，哪些组合是被megatron禁止的，
# 加深对代码理解
# -------------------------------------------------------------------------------
......2.2 Megatron SwitchMLP
相关脚本：Megatron-LM/megatron/model/transformer.py
Megatron将其MOELayer的实现命名为SwitchMLP，与deepspeed的MOELayer相比主要有以下不同：
- Megatron采取的是top1Expert策略
- Megatron是token dropless的，也就意味着所有的token都会被发送到其对应的1st expert上，不存在任何溢出问题
第1点我们比较好理解，但是对于第2点我们就有困惑了：在token dropless的情况下，megatron怎么解决expert间token负载不均的问题呢？这会影响它ep_group内的通讯吗？
为了解答这些困惑，我们来详细画出Megatron SwitchMLP的流程图：
这里我们只取pp切分后的某一layer来看，其余layer同理可推。
- 首先，non-moe层的输出结果尺寸为
(seq_len, batch_size, M)，我们将其reshape为(S, M)，其中S = seq_len * batch_size。同一个tp组内的输出因为经过了AllReduce所以完全一致。以g0这块卡上的数据来说，我们假设a0是要送往e0的tokens，a2是要送往e1的tokens，其余gpu上也是类推，这样我们就得到了图例中的第一行。 - 在ep_group内，对所有数据做AllGather，然后再把不是这块卡的expert维护的token给mask掉。从这一步开始你就能发现Megatron和deepspeed的不同了：
- deepspeed中，在对ep_group做All2All通讯时，token是定向发送到它对应的expert所在的卡上的。因为不同expert维护的token数量不同，因此通讯时肯定有负载不均问题。为了解决这个问题，deepspeed对expert引入了capacity机制，可以理解成token是装在一个固定容量的容器（buffer）中发送的，容器中没装满的地方就用0填充。通过这种方式deepspeed解决了通讯时负载不均的问题。
- Megatron中，对ep_group采用的是AllGather通讯，它才不管token对应的expert是哪个，它直接把数据全量发送到各卡上，然后再在各卡上mask掉不是这张卡的expert维护的token，这样就能保证ep_group在通讯时负载均衡，缺点就是会发送大量的重复数据。通过这一步，大家是不是能更好理解deepspeed和Megatron中ep_group定义不同的原因了？
- 在各卡上，当我们将不是这块卡expert维护的token mask掉后，我们就能把其余的token喂给expert了（见图中第三行）。和deepspeed一样，Megatron的单个expert也是ParallelMLP。但与deepspeed不同的是，在Megatron中，如果一个expert采用了tp切分，那么在同一个tp组内，我们不对expert的输出结果做AllReduce，这样做是因为我们马上就要在ep_group内做ReduceScatter操作，这个操作自然会把expert不同tp组的结果相加起来。
- 数据在MoE层计算完毕后，我们将数据每张卡上的数据划分为ep_world_size份（注意，这里ep_world_size是指Megatron ep_group内的卡数，其值为4。而我们分布式设置group组时这个ep_world_size为2。大家阅读代码时可特别注意这一点），然后将卡上的数据做ReduceScatter，就能得到最终的结果了（图例中最后一行）。大家可以发现现在同个tp组内的数据又完全一致了，就可以将其当作输入喂给下一层non-MoE层了。
不难发现：
- 在ep_group内，Megatron通过AllGather + ReduceScatter的方式，替换掉了deepspeed All2All + expert buffer的形式，保持了ep_group内各卡通讯时的负载均衡，但Megatron这样做的缺点就是发送了大量的重复数据
- 不管是deepspeed还是Megatron，它们只是解决了ep_group内通讯时的负载不均问题，但是实际计算时，每个expert拿到的有效token（即非zero padding的token）数量还是不一样的，这一块我们一般通过辅助损失函数（原理篇中讲过）等方式来解决。
相信有了上面图例讲解，大家阅读Megatron SwitchMLP的代码就不困难了。这里我们不再把代码细节贴出（其实除了通讯外，就是一堆矩阵计算），留给大家自行阅读。
恭喜你竟然读完了这篇文章！现在你是不是能感受到我曾经说的：MoE并行训练开源资料有一种鱼龙混杂，难成体系的感觉？在原理篇和源码篇中，我尽量统一各种符号表达，明确定义，并努力将这个脉络理出来，希望对大家学习MoE并行能有所帮助。这里还遗留了一个坑，就是对Tutel优化的分析。
#关于MoE模型的面经汇总
Find something that you can benefit from scaling. ——Ilya
LLM 时代流传着一个法则：Scaling Law，即通过某种维度的指数上升可以带来指标的线性提升。
如下图所示，在 Compute、Data、Parameter 三个维度上的指数上升可以带来在 test loss 上的线性下降。
MoE（Mixture of Experts，混合专家模型）从本质上来说就是一种高效的 scaling 技术，用较少的 compute 实现更大的模型规模，从而获得更好的性能。
据说目前 LLM 的天花板 GPT-4 也使用了 MoE 技术，Mistral 7B /w 8 experts 的 checkpoint 释出，彻底引爆了 AI 社区对 MoE 的热情。
模型规模是提升模型性能的关键因素之一。在有限的计算资源预算下，用更少的训练步数训练一个更大的模型，往往比用更多的步数训练一个较小的模型效果更佳。
MoE 的一个显著优势是它们能够在远少于稠密模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下，您可以显著扩大模型或数据集的规模。特别是在预训练阶段，与稠密模型相比，混合专家模型通常能够更快地达到相同的质量水平。
MoE结构和原理
作为一种基于 Transformer 架构的模型，MoE 主要由两个关键部分组成:
- 稀疏 MoE 层: 这些层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干“专家”(例如 8 个)，每个专家本身是一个独立的神经网络。在实际应用中，这些专家通常是前馈网络 (FFN)，但它们也可以是更复杂的网络结构，甚至可以是 MoE 层本身，从而形成层级式的 MoE 结构。
- 门控网络或路由: 这个部分用于决定哪些 token 被发送到哪个专家。一般 input Token 会被 router 分配到一个或多个 expert 上做处理。如果是多个 expert 处理同一个 input token，那么不同 expert 的输出会再做一次 aggregate，作为 input token 的最终 output。
例如，在下图中，“More”这个 token 被发送到第二个专家，而“Parameters”这个 token 被发送到第一个专家。
MoE的优点
- 任务特异性：采用混合专家方法可以有效地充分利用多个专家模型的优势，每个专家都可以专门处理不同的任务或数据的不同部分，在处理复杂任务时取得更卓越的性能。各个专家模型能够针对不同的数据分布和模式进行建模，从而显著提升模型的准确性和泛化能力，因此模型可以更好地适应任务的复杂性。
- 灵活性：混合专家方法展现出卓越的灵活性，能够根据任务的需求灵活选择并组合适宜的专家模型。模型的结构允许根据任务的需要动态选择激活的专家模型，实现对输入数据的灵活处理。这使得模型能够适应不同的输入分布和任务场景，提高了模型的灵活性。
- 高效性：由于只有少数专家模型被激活，大部分模型处于未激活状态，混合专家模型具有很高的稀疏性。这种稀疏性带来了计算效率的提升，因为只有特定的专家模型对当前输入进行处理，减少了计算的开销。
- 表现能力：每个专家模型可以被设计为更加专业化，能够更好地捕捉输入数据中的模式和关系。整体模型通过组合这些专家的输出，提高了对复杂数据结构的建模能力，从而增强了模型的性能。
- 可解释性：由于每个专家模型相对独立，因此模型的决策过程更易于解释和理解，为用户提供更高的可解释性，这对于一些对模型决策过程有强解释要求的应用场景非常重要。
- 适应大规模数据：混合专家方法是处理大规模数据集的理想选择，能够有效地应对数据量巨大和特征复杂的挑战，可以利用稀疏矩阵的高效计算，利用GPU的并行能力计算所有专家层，能够有效地应对海量数据和复杂特征的挑战。
MoE的问题
尽管混合专家模型 (MoE) 提供了若干显著优势，例如更高效的预训练和与稠密模型相比更快的推理速度，但它们也伴随着一些挑战:
- 训练复杂性: 虽然 MoE 能够实现更高效的计算预训练，但其训练相对复杂，尤其是涉及到门控网络的参数调整。为了正确地学习专家的权重和整体模型的参数，反而可能需要更多的训练时间。另外在微调阶段往往面临泛化能力不足的问题，长期以来易于引发过拟合现象。
- 超参数调整：选择适当的超参数，特别是与门控网络相关的参数，以达到最佳性能，是一个复杂的任务。这可能需要通过交叉验证等技术进行仔细调整。
- 专家模型设计：专家模型的设计对模型的性能影响显著。选择适当的专家模型结构，确保其在特定任务上有足够的表现力，是一个挑战。
- 稀疏性失真：在某些情况下，为了实现稀疏性，门控网络可能会过度地激活或不激活某些专家，导致模型性能下降。需要谨慎设计稀疏性调整策略，以平衡效率和性能。
- 动态性问题：在处理动态或快速变化的数据分布时，门控网络可能需要更加灵活的调整，以适应输入数据的变化。这需要额外的处理和设计。
- 对数据噪声的敏感性：混合专家模型对于数据中的噪声相对敏感，可能在一些情况下表现不如其他更简单的模型。
- 通信宽带瓶颈：在分布式计算环境下可能面临通信宽带瓶颈的问题。这主要涉及到混合专家模型的分布式部署，其中不同的专家模型或门控网络可能分布在不同的计算节点上。在这种情况下，模型参数的传输和同步可能导致通信开销过大，成为性能的一个瓶颈。
- 推理挑战: MoE 模型虽然可能拥有大量参数，但在推理过程中只使用其中的一部分，这使得它们的推理速度快于具有相同数量参数的稠密模型。然而，这种模型需要将所有参数加载到内存中，因此对内存的需求非常高。
对于推理挑战，以 Mixtral 8x7B 这样的 MoE 为例，需要足够的 VRAM 来容纳一个 47B 参数的稠密模型。之所以是 47B 而不是 8 x 7B = 56B，是因为在 MoE 模型中，只有 FFN 层被视为独立的专家，而模型的其他参数是共享的。此外，假设每个token只使用两个专家，那么推理速度 (以 FLOPs 计算) 类似于使用 12B 模型 (而不是 14B 模型)，因为虽然它进行了 2x7B 的矩阵乘法计算，但某些层是共享的。
示例代码
MoE 的 Pytorch 示例代码如下，大家可以自己学习并运行一下：
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
# 创建一些随机数据（替换为真实数据）
num_samples = 1000
num_features = 300 # 假设文本已经转换为固定大小的向量
num_classes = 10 # 假设有10个类别
# 随机生成数据和标签
X = np.random.randn(num_samples, num_features)
y = np.random.randint(0, num_classes, num_samples)
# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 定义 Dataset
class TextDataset(Dataset):
def __init__(self, features, labels):
self.features = features
self.labels = labels
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
return torch.tensor(self.features[idx], dtype=torch.float), torch.tensor(self.labels[idx], dtype=torch.long)
# 创建 DataLoader
train_dataset = TextDataset(X_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_dataset = TextDataset(X_test, y_test)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
###模型定义
class TopKGating(nn.Module):
def __init__(self, input_dim, num_experts, top_k=2):
super(TopKGating, self).__init__()
# 初始化线性层作为门控机制
self.gate = nn.Linear(input_dim, num_experts)
# 设置要选择的顶部专家数量
self.top_k = top_k
def forward(self, x):
# 计算每个专家的分数
gating_scores = self.gate(x)
# 选取分数最高的 top_k 个专家，并返回它们的索引和 softmax 权重
top_k_values, top_k_indices = torch.topk(F.softmax(gating_scores, dim=1), self.top_k)
return top_k_indices, top_k_values
class Expert(nn.Module):
def __init__(self, input_dim, output_dim):
super(Expert, self).__init__()
# 为每个专家定义一个简单的神经网络
self.net = nn.Sequential(
nn.Linear(input_dim, 128),
nn.ReLU(),
nn.Linear(128, output_dim)
)
def forward(self, x):
# 通过专家网络传递输入数据
return self.net(x)
class MoE(nn.Module):
def __init__(self, input_dim, num_classes, num_experts, top_k=2):
super(MoE, self).__init__()
# 设置专家数量
self.num_experts = num_experts
# 设置类别数量
self.num_classes = num_classes
# 初始化 TopK 门控层
self.gating = TopKGating(input_dim, num_experts, top_k)
# 创建专家网络的列表，每个专家是一个 Expert 实例
self.experts = nn.ModuleList([Expert(input_dim, num_classes) for _ in range(num_experts)])
def forward(self, x):
# 获取批量大小
batch_size = x.size(0)
# 通过门控层获得 top_k 专家的索引和门控权重
indices, gates = self.gating(x) # 形状 indices：[batch_size, top_k], gates：[batch_size, top_k]
# 准备收集选定专家的输出
expert_outputs = torch.zeros(batch_size, indices.size(1), self.num_classes).to(x.device)
# 遍历每个样本和其对应的 top_k 专家
for i in range(batch_size):
for j in range(indices.size(1)):
expert_idx = indices[i, j].item() # 获取专家的索引
expert_outputs[i, j, :] = self.experts[expert_idx](x[i].unsqueeze(0))
# 将门控权重扩展到与专家输出相同的维度
gates = gates.unsqueeze(-1).expand(-1, -1, self.num_classes) # 形状：[batch_size, top_k, num_classes]
# 计算加权的专家输出的和
output = (gates * expert_outputs).sum(1)
return output, gates.sum(0) # 返回模型输出和门控使用率以用于负载平衡损失计算
import torch.nn.functional as F
def moe_loss(output, target, gating_weights, lambda_balance=0.1):
# 标准损失（例如交叉熵损失）
# output 是模型的输出，target 是真实的标签
standard_loss = F.cross_entropy(output, target)
# 负载平衡损失
# gating_weights 是门控权重，表示每个专家的使用率
# 使用标准差来衡量各专家使用率的平衡程度
balance_loss = torch.std(gating_weights)
# 总损失
# 结合标准损失和负载平衡损失，lambda_balance 是一个超参数，用于控制负载平衡损失在总损失中的比重
total_loss = standard_loss + lambda_balance * balance_loss
return total_loss
# 初始化模型
model = MoE(input_dim=num_features, num_classes=num_classes, num_experts=4, top_k=2)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 训练循环
num_epochs = 1
for epoch in range(num_epochs):
model.train()
total_loss = 0
for features, labels in train_loader:
optimizer.zero_grad()
outputs, gating_weights = model(features)
loss = moe_loss(outputs, labels, gating_weights)
loss.backward()
optimizer.step()
total_loss += loss.item()
print(f'Epoch {epoch+1}, Loss: {total_loss/len(train_loader)}')
def evaluate(model, data_loader):
model.eval（)
predictions, true_labels = [], []
with torch.no_grad():
for features, labels in data_loader:
s = time.time()
outputs, _ = model(features)
e = time.time()
print(e-s)
predicted = torch.argmax(outputs, dim=1)
predictions.extend(predicted.tolist())
true_labels.extend(labels.tolist())
return accuracy_score(true_labels, predictions)
