大商创源码_晚上网站推荐靠谱的2021_网络销售推广平台_百度排行榜风云榜

探秘Transformer系列之（21）— MoE

0x00 概要

在足够的训练数据下，我们可以通过增加参数和计算预算来扩大语言模型规模就可以得到更强大的模型。然而，与之相关的问题是极高的计算成本。而MoE（Mixture-of-Experts/混合专家）架构通过条件计算，就可以在保持计算成本适度的情况下实现参数扩展，提供增强的模型容量和计算效率。简单理解，MoE就是将多个专家模型混合起来形成一个新的模型。但MoE不是让一个单一的神经网络处理所有任务，而是将工作分配给多个专门的“专家”，由一个门控网络决定针对每不同输入激活哪些专家。

---

注：全部文章列表在这里，估计最终在35篇左右，后续每发一篇文章，会修改文章列表。
csdn 探秘Transformer系列之文章列表

---

0x01 前置知识

1.1 MoE出现的原因

MoE的出现有几个主要方面原因：神经网络的稀疏性、神经元的多语义性和计算资源的有限性。我们从中也可以看到FFN的部分劣势。

1.1.1 神经网络的稀疏性

稀疏性是指我们可以仅使用整个系统的某些特定部分执行计算。这意味着并非所有参数都会在处理每个输入时被激活或使用，而是根据输入的特定特征或需求，只有部分相关参数集合被调用和运行。

虽然Transformer构建了庞大的参数网络，但是它的某些层可能会非常稀疏，即某些神经元的激活频率会低于其他神经元。论文“MoEfication: Transformer Feed-forward Layers are Mixtures of Experts”就指出，使用ReLU等激活函数会导致大部分的激活值都是0，这导致FFNs的激活值非常稀疏。这样在每次预测过程中，对于用户当前的问题来说，FFNs 中实际只有一小部分神经元被激活并参与计算。而且模型的规模越大，其稀疏性也越强。大型模型在处理输入时激活的神经元占总体的比例更小。其实，人脑也具备类似的稀疏性。如果对于所有问题，人脑都会使用全部神经元，恐怕人脑中的“CPU”早就烧毁了。

1.1.2 神经网络的过载性

对每个输入，主流深度神经网络都会载入网络中的所有层和神经元，所有模型参数都会一同参与处理该输入数据。因为上面提到的稀疏性，我们可知这意味着在处理大量参数的过程中，需要进行大量的不必要的计算。因此，网络实际上对于它们所做的大多数预测来说都太大了。LLM 成为世界上最低效和最耗能的系统之一。

除了不受控制的消耗之外，针对每个预测运行整个模型也会对性能产生重要影响，参数数量的增加会导致训练和推理过程中计算复杂度和内存消耗的增加。在追求速度和可扩展性的实际应用中部署如此庞大的模型是一项艰巨的任务。

1.1.3 神经元的多语义性

随着应用场景的复杂化和细分化，垂直领域应用更加碎片化，人们对大模型提出了更高的要求，希望一个模型既能回答通识问题，又能解决专业领域问题。

但是，有研究人员发现，神经元具有多义性的特点。也就是说，它们不专注于一个单一的主题，而是专注于许多主题。而且重要的是，它们在语义上可能并不相关。举个例子来说，在神经网络数十亿个神经元中的一个神经元可能每次在输入主题涉及“苹果”被激活，而当输入主题涉及“手机”时，这个神经元也可能被激活。这不仅使神经网络难以解释，而且也不是一个理想的情况。因为单个神经元必须精通各种彼此几乎毫无关系的主题。想象一下，你必须同时成为神经科学和地质学的专家，这将是一项艰巨的任务。

而目前不仅仅是知识范围更加广泛，多模态带来的各自数据集都可能各自的数据特征完全不同，这导致神经元很难获取知识。更糟糕的是，学习曲线可能相互矛盾，学习一个主题的更多知识可能会影响神经元获取另一个主题知识的能力。

1.1.4 计算资源的有限性

模型规模是提升模型性能的关键因素之一。而通常来讲，模型规模的扩展会导致训练成本显著增加，因此，计算资源的限制成为了大规模密集模型训练的瓶颈。

因此，需要一种技术来拆分、消除或至少缓解这些问题。这就是MoE希望达到的目的。

1.2 MoE的核心理念

MoE 的基本思想是将模型的参数计数与其使用的计算量分离。而其背后的理念则是模型的不同组件（即"专家"）在处理数据的不同任务或特征时具有专门化的能力。这种设计灵感来源于人类社会中的专业分工。在现实生活中，如果有一个包括了多个领域知识的复杂问题，我们通常会召集一个专家团队共同解决复杂问题。每位专家都拥有独特的技能。我们先拆分这个大问题到各领域，把不同的任务先分离出来，这样才便于分发给不同领域的专家。然后让各个领域的专家先逐个解决小问题，最后再把大家集合到一起来汇总结论，攻克这个任务。

MoE正是基于上述的理念，它由两个主要部分组成：专家和门控路由机制（或者路由机制）。

• 术业有专攻。模型的不同专家（expert）拥有不同领域的专业知识，负责处理不同的计算任务或者数据。每个专家子网络专门处理输入数据的子集，共同完成一项任务。相较于深度学习网络， MoE更像是宽度学习网络。另外，MoE与集成技术的主要区别在于，对于MoE，通常只有一个或少数几个专家模型针对每个输入进行运算，是稀疏模型；而在集成技术中，所有模型都会对每个输入进行运算，然后通过某种方式来综合这些模型的输出，这是密集模型。
• 有条件的计算。既然不同专家负责不同领域，怎么知道要把哪个token送去哪个expert呢？因此我们就需要对神经网络实际运行的程度拥有某种“决策权”，使得针对特定输入，只有特定专家被激活并处理（在生成式大模型中，就是根据token token 来选择专家的）。这部分工作就由门控机制来完成。其实，MoE的稀疏性与dropout的原理有些类似，MoE是根据任务的具体情况选择激活一定数量的专家模型来完成这个任务，而dropout则是对神经网络中的神经元进行随机性失活。

通过这种范式，模型将计算与参数解耦，仅激活与特定输入相关的专家，既保持了大规模知识库的优势，又有效控制了计算成本。而且，MoE能够在远少于 Dense 模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下，我们可以显著扩大模型或数据集的规模。这种可扩展且灵活的创新有效遵循了扩展规律，实现了模型容量的增长而不会导致计算需求的剧增。

0x02 发展历史

2.1 重要节点

下图是MoE发展历史上的一些重要节点。

2.1.1 Adaptive mixtures of local experts

MoE的开山之作是1991年的论文“Adaptive Mixture of Local Experts"。这篇论文引入了将复杂问题分解为子问题并分配给多个专门模型的思想。这种分而治之的策略成为了 MoE 架构的核心。

因为面对多任务学习时，多层网络的各层之间通常会有强烈的干扰效应，这会导致学习过程变慢和泛化能力差。为了解决这个问题，论文提出了一种新的监督式学习方法：由多个独立子网络（专家）组成一个系统，每个子网络独立学习整个训练数据集中的一个子集。模型使用一个门控网络（gating network）来决定每个数据应该被哪个子网络去训练，这样就可以减轻不同类型样本之间的干扰。在推理时，模型将输入同时传递给不同的子网络和门控网络，每个子网络给出自己的处理结果，门控网络会依据每个子网络的权重来决定每个子网络对当前输入的影响程度，最终给出所需的输出。

如何训练这个系统？如何让损失函数整合专家和门控网络的输出？论文作者作者提出了两种思路：鼓励竞争和鼓励合作，具体参见下图。

2.1.2 sparsely-gated mixture-of-experts layer

2017年，论文“Outrageously large neural networks: The sparsely-gated mixture-of-experts layer”首次将MoE引入自然语言处理领域，并提出了Sparse MoE（稀疏MoE）概念。和论文"Adaptive mixtures of local experts"相比，本论文 MoE 的主要区别如下：

• Sparsely-Gated：不是所有expert都会起作用，而是选择TopK的专家进行计算，即只激活部分专家对特定输入进行处理。这是一种条件计算，意味着只有部分专家被激活处理特定的输入，从而可以大大降低计算量。而且，门控网络依然为每个输入同时选取多个专家，使网络能够权衡并整合各专家的贡献，从而提升性能。这种稀疏性就是MoE可以把模型容量扩大的原因。
• token-level：相比于sample-level，此论文使用了在 token 级别进行处理，一个句子中不同的token使用不同的专家。由于其为每个输入令牌（token）选择对应专家的特点，该方法被称为token选择门控（token-choice gating）。

本论文之前的MoE中，每个专家都用于每个输入，但是每个专家的贡献由一个门控函数加权，这是一个学习到的函数，它为每个专家计算一个权重或重要性，使得所有专家的权重之和为1。由于每个专家都用于每个输入，这种方法仍然导致一个密集激活的模型，因此没有解决增加计算复杂度的问题。这种路由选择算法也叫做软性选择路由算法（也称为连续混合专家）。本论文使用的路由算法是硬选择路由算法，运行只有一部分专家用于任何给定的输入，这标志着从密集激活到稀疏模型的转变。

另外，门控网络会倾向于收敛到不均衡的状态，总是为少数专家产生较大的权重（相应的参数更新也会很不均衡）。因此，作者设计了额外的损失函数来促使所有专家具有同等的重要性，也首创了具有辅助负载平衡损失的可微分启发式方法，通过选择概率对专家输出进行加权，使门控过程可微分，从而能够通过梯度优化门控函数。这种方法随后成为MoE领域的主流研究范式。

本论文也是业界第一个实施专家并行的方案。

2.1.3 GShard

2021年的论文”GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding“是第一个将MoE的思想拓展到Transformer上的工作。具体而言，GShard 将每间隔一层的FFN层替换为 MoE 结构，MoE 中的每个专家都是一个FFN（每个专家大小相同）。每个 Token 都会通过 Gating 选择不同的专家，默认为 top2。

因为难以控制将token发给专家的概率，所以在实际操作中，可能某些expert接收到了好多token，而某些expert接收的token寥寥无几，我们管这种现象叫expert负载不均。这种情况不仅不符合我们MoE的设计初衷（术业有专攻），还影响计算效率（例如引起分布式训练中各卡通讯时的负载不均）。为了缓解“赢者通吃”问题，尽可能让不同专家处理的token数尽量均衡，Gshard提出了以下几种解决办法：

• 专家容量负载：为了确保负载平衡，GShard强制每个专家处理的token数量低于某个统一阈值，论文将其定义为专家容量。当token选择的两个专家都已经超出其容量时，该token被视为溢出token，这些token或者通过残差连接传递到下一层，或被完全丢弃。
• Local group dispatching（本地组调度）：将训练批次中的所有token均匀划分为 G 组，即每组包含 S = N/G 个token。所有组均独立并行处理。通过这种方式，我们可以确保专家容量仍然得到执行，并且总体负载保持平衡。
• Auxiliary loss（辅助损失）：添加辅助损失函数，对expert负载不均的情况做进一步惩罚。
• 随机路由：在Top-2 gating的设计下，GShard 始终选择排名最高的专家，但第二个专家是根据其权重比例随机选择的。直觉上认为在输出是加权平均且次要权重通常较小的情况下，次要专家的贡献可以忽略不计。

Gshard也提出了MoE跨设备分片的方法。当扩展到多个设备时，MoE 层在不同设备间共享，而其他所有层则在每个设备上复制。这样，整个 MoE 层的计算被分散到了多个设备上，每个设备负责处理一部分计算任务。这种架构对于大规模计算非常有效。这也解释了为什么 MoE 可以实现更大模型参数、更低训练成本。

GShard 为后续所有的 MoE 研究铺好了路：它证明了稀疏专家是有价值的，还为后来的方案指明了“容量因子”这些关键概念的重要性。

2.1.4 Swith Transformer

论文”Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity“用稀疏激活的混合专家层替换了Transformer块中的FFN，同时使用简化的门控机制使训练更加稳定，从而使 MoEs 成为语言建模应用更现实、更实用的选择。

Switch Transformer 的指导设计原则是以简单且计算高效的方式最大化 Transformer 模型的参数数量。在此原则指导下，论文做了一些有效努力，包括：简化稀疏路由、使用高效稀疏路由和增强的训练和微调技巧。Switch Transformer 作者发现仅使用一个专家也能保证模型的质量。一个专家可以让路由计算更简单，通信量也更少；一个 Token 仅对应一个专家，计算量也更少；平均每个专家对应的 batch size 至少可以减半。因此，Switch Transformer 的门控网络每次只路由到 1 个 expert，也就是每次只选取 Top1 的专家，而其他的模型都是至少 2 个专家。

2.2 详细时间线

下图是2017年之后的若干代表性的MoE模型的时序概览。时间线主要根据模型的发布日期构建。位于箭头之上的MoE模型是开源的，而位于箭头之下的则是专有闭源模型。来自不同领域的MoE模型用不同颜色标记：自然语言处理（ NLP）用绿色，计算机视觉（CV）用黄色，多模态（multimodal）用粉色，推荐系统（RecSys）用青色。

0x03 模型结构

MoE包括以下核心组件：

• 专家。在MoE架构中，专家是专门针对特定任务的子模型。专家拥有不同领域的专业知识，负责处理不同的计算任务或者特定输入子空间。形式上，每个专家网络$f_i$（通常是一个linear-ReLU-linear网络）由参数W来参数化，接受输入x并生成输出$f_i(x;W_i)$。
• 门控函数（也称为路由函数或路由器）：门控函数负责协调专家计算，即判定哪个输入样本应该由哪些专家处理，哪些专家将被激活并参与到当前的计算中。形式上，门控函数G（通常由linear-ReLU-linear-softmax网络组成）由参数O来参数化，接受输入x并产生输出。
• 聚合层（Combining Layer）：聚合层负责整合专家网络的输出，以形成最终的输出结果。很多资料并没有把聚合层单独罗列出来。

整个结构可以用下图来表示。门控网络输出是一个稀疏的n维向量，$G(x{)}_i$是门控网络给出的第 i 个专家的权重。$E_i(x)$ 是第 i 个expert的输出。那么对于在当前的输入x，输出就是所有 experts 的加权和。

下图给出了MoE的总体处理流程，具体是如下五步。这里每个专家和token都有颜色编码，门控网络权重（W）有每个专家的表示（颜色匹配）。为了确定路由，路由器权重对每个token embedding（x）执行点积，以产生路由器得分（h(x)）。然后将这些分数归一化为1（p(x)）。G使用了softmax函数。

• 将输入token的embedding和门控网络权重进行点积，得到门控分数。在语言建模的上下文中，这里每一列将表示输入序列中的一个token。因此，每个token可以路由到不同的专家。
• 在门控分数上施加softmax将门控分数进行归一化，得到概率。此概率表示每个专家模型对该token的贡献程度，即在给定输入情境下每个专家被激活的概率。或者说此概率表明专家处理传入token的能力如何。
• 使用此概率分布作为权重来选择最佳匹配的专家。
• 专家对输入token进行处理。
• 专家处理之后，将每个路由器的输出与每个选定的专家相乘，并对结果求和。

3.1 门控函数

3.1.1 条件计算

稀疏激活是 MoE 模型的关键部分和优势之一。与所有专家或参数对输入都活跃的密集模型不同，稀疏激活确保只有一小部分专家根据输入数据被激活。这种方法在保持性能的同时减少了计算需求，因为任何时候只有最相关的专家是活跃的。

本质上，我们现在所谈论的MoE 大模型是使用条件计算来强制稀疏激活（Sparse Activation）。条件计算是探讨如何分离计算复杂性和计算量需求，并在其之间进行合理的权衡的理论。在此处的意思是动态开启/关闭神经网络的部分功能。MoE 模型条件计算的核心是学习一个计算成本低的映射函数，该函数确定网络的哪些部分——换句话说，哪些专家可以最有效地处理给定的输入。条件计算在大模型中通常使用路由网络或者门控网络来实现。它是判断选择使用哪个专家的关键，即在网络中根据输入数据有选择地激活部分单元。在当前的大语言模型中，这是通过对每个 token 进行条件判断来实现的。当模型输入一个token时，路由网络根据上下文和当前token，选择合适的专家网络来计算。这种选择性激活的直接效果是加快信息在网络中的传播速度，无论是在训练阶段还是推理阶段。通过条件计算或者稀疏性，大模型能够在增加模型规模的同时，降低计算成本，实现了一个合适的均衡。其实大模型推理中常见的早退出机制（Early-Exit）也是一种条件计算，它允许在网络的早期层级就做出决策并减少计算。

因为是条件计算，所以与具有相同参数数量的模型相比，MoE具有更快的推理速度。也因为是条件计算，所以MoE需要把所有专家系统完全加载到内存中，所以需要大量显存。

3.1.2 定义

门控网络（Gating Network）的设计和实现是Sparsely-Gated MoE 层的核心组成部分。门控网络负责为每个输入 token 选择一个稀疏的专家组合，这些专家将参与到当前的计算中。门控函数是一个可以执行一系列非线性变换的网络，该网络对概率分布进行建模，根据概率去做出相应的选择。门控网络由学习到的参数组成，并且与网络的其余部分同时进行预训练。一个典型的门控网络就是一个带有 softmax 函数的简单网络。

假定注意力层的输入数据形状是 (batch_size, seq_len, embedding_size)，则门控网络的大小是 (token_size, expert_num)，门控网络的输入形状是 ( batch_size * seq_len, embedding_size)， 输出是 ( batch_size * seq_len, expert_num)，即每个token去向每个expert的概率。比如：

gates ( batch_size * seq_len = 3, expert_num = 4)：
[[0.2, 0.4, 0.1, 0.3],  # Token A 被分配到不同专家的概率[0.1, 0.6, 0.2, 0.1],  # Token B 被分配到不同专家的概率[0.3, 0.1, 0.5, 0.1]   # Token C 被分配到不同专家的概率
]

门控网络会学习将输入发送给哪个expert，softmax的输出作为每个专家的最终使用权重。门控函数的处理流程如下：

• 计算专家分数。门控函数接收单个token的emebdding作为输入，基于输入数据的特征进行计算，然后输出一组分数。这些分数表示每个专家模型对该token的贡献程度，或者说是表明专家处理传入token的能力如何。
• 计算专家的概率分布。以下图为例，门控函数使用softmax对分数进行处理，得到在给定输入情境下每个专家被激活的概率分布。这个分布反映了输入数据与各个专家相关性的大小，概率越高，表示该专家对于当前输入的预测任务越重要。下图中选择了两个专家，门控函数输出的概率是0.1和0.9，说明专家1对该token贡献是10%，专家2对该token的贡献是90%。
• 激活专家。门控函数将每个 Token 作为输入，并在 expert 上生成一个概率分布，以确定每个 Token 被发送给哪个 expert。根据门控输出的概率分布，一部分专家将被选中并激活。在下图中，如果使用top2 策略来选择，因为其它专家的激活概率不到0.1，因此没有被选中。专家2和专家n-1因为具有较高的激活概率，将被选中参与到后续的计算中。这意味着，只有这两个专家的参数将被用于处理当前的输入数据。假设专家2输出结果值是0.4，专家n-1输出结果值是0.5，则最终MoE返回选定专家的输出乘以门值（选择概率）是0.49。通过同时咨询多个专家对给定输入的意见，网络能够有效地权衡并整合他们的贡献，从而提升性能。

注：在MoE模型中，虽然动态训练门控功能是标准做法，但一些研究探索了非可训练的token选择门控机制（Non-trainable Token-Choice Gating）。这种机制的主要优势在于不需要额外的门控网络参数，通过特定的门控机制即可实现全面的负载均衡。比如，Hash Layer采用基于随机固定门控的方式，通过对输入token进行哈希，无需训练门控网络即可工作。人们还探索了其他更复杂的哈希函数，例如通过对单独预训练的Transformer模型产生的token嵌入应用k-means聚类，或者根据训练数据中token频率预计算一个哈希表，将token ID映射到专家，从而确保token到专家的分配更加平衡。

3.1.3 特点

我们接下来看看门控函数的特点。

首先，门控函数不仅决定在推理过程中选择哪些专家，还决定在训练过程中选择哪些专家。.这是因为只有让每个专家在训练期间学习到不同的信息，在推理时，才能知道哪些专家与给定的任务最相关。

其次，门控函数逐层都会选择专家。在具有 MoE 的 LLM 的每个层级中，我们都会找到（某种程度上专业的）专家。具体参加下图。

其实，在最微观层面，每个神经元就是一个专家，激活函数就起到了门控函数的作用。MoE只是把很多神经元聚类成一个专家。

另外，"OpenMoE"和Mixtral都对门控机制进行了分析，其中提到几个特点如下：

• 上下文无关的专精化（Context-independent Specialization）。MoE 倾向于简单地根据 token 级语义对 token 进行路由，即无论上下文如何，某些关键词经常被分配给同一位专家。因为路由规则和文本的语义主题无关，这说明MoE 模型中的专家各自擅长处理不同的 token，但是专家实际可能并不专门精通某一领域的知识。
• 路由早期习得（Early Routing Learning）。路由在预训练的早期就已建立，并且基本保持不变，因此 token 在整个训练过程中始终由相同的专家处理。这或许能启发我们设计更高效的路由机制。
• 序列尾部丢弃（Drop-towards-the-End）现象显著。序列后部的token因专家达到容量上限而更容易被丢弃。具体而言，在 MoE 模型中，为了保证负载均衡，通常会为每个专家设置其容量上限。当某个专家的容量达到上限时，该专家将不再接受新的 token 而将其丢弃（Drop）。如果我们从前往后为序列中的 tokens 分配专家，那么序列尾部的 tokens 将有更大的概率被丢弃，这在指令调优数据集中更为严重。
• 位置的局部性。相邻的token通常被路由到同一位专家，这表明token在句子中的位置会影响路由选择，会带来“高重复率”现象。这有利于减少专家负载的突发波动，但也可能导致专家被局部数据“霸占”。

因此，我们需要了解数据集的“局部规律”，因为一旦数据分布换了（比如从新闻文本转到代码），它原先的路由模式就可能失效。要做大规模的 MoE，就得好好考虑数据特征和专家分配之间的关系。

3.1.4 优化

关键因素

现在MoE 大模型的整体架构非常固定，而路由选择则成了关键。路由算法可以从简单（在张量的平均值上进行均匀选择或分箱）到复杂。在决定特定路由算法对问题的适用性的许多因素中，以下几个经常被讨论。

• 模型精度。MoE 模型对舍入误差很敏感，比如 softmax 中的指数操作可能会产生舍入误差，导致训练不稳定。 然而，简单地裁剪（即应用硬阈值以删除大值）门控函数输出的logit 又可能会损害模型性能。
• 平衡负载。我们希望尽可能让不同专家处理的token数尽量均衡。目前我们知道最简单的解决方案是根据 softmax 概率分布选择前 k 个专家。然而，这种方法会导致训练负载不平衡：训练期间，大多数token都会被分发给少数专家，因此这少数专家积累了大量的输入token，而其它专家比较空闲，这减慢了训练速度。与此同时，许多其他专家根本没有接受过足够的训练。因此需要更好的门控函数，以便在所有专家之间更均匀地分配token。
• 高效。如果门控函数只能串行执行，则很难实现负载均衡。假设我们有 E 个专家和N个token，则仅门函数的计算成本就至少为 O(NE)。在实际工作中，N 和 E的数量级会很大，门控函数的低效执行将使大部分计算资源（专家）在大多数时间处于闲置状态。因此，我们需要让门控函数可以高效并行实现来利用众多设备。

为了达成这些目标，研究人员做了不谢的努力，下图展示了MoE模型中使用的不同门控函数。包括 (a)使用top-1门控（top-1 gating）的稀疏MoE；(b)BASE层（BASE layer），©组合领域映射与随机门控（combining domain mapping with stochastic gating），(d)专家选择门控（expert selection gating），(e)注意力路由器（attention router），以及(f)带专家合并的软MoE（soft MoE with expert merging）。这些函数可能通过各种形式的强化学习和反向传播进行训练，做出二元或稀疏且连续、随机或确定性的门控决策。

改进示例

我们以论文”Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer“为例，看看如何对简单的 softmax 门控函数做升级，进而满足需求。

softmax 函数的问题：softmax 函数会让所有专家都会对输入进行运算，再通过门控网络的输出进行加权求和，如果experts的数量太大，就会导致计算量非常大。因此，我们需要找到一种方法能使某些专家模型的门控网络的输出为0，这样就没有必要对这个专家进行相应的计算，就可以节省计算资源。

普通” top-k 路由策略可以满足这个需求，其会根据 softmax 概率分布选择前 k 个专家。即，在Softmax函数应用于专家权重之前，执行一个KeepTopK操作，将除前k个专家之外的所有专家的权重设置为-∞。这确保了只有前k个专家在应用Softmax后权重大于0。因此，这种MoE可以帮助我们在扩大模型规模的同时保证计算量是非线性增加的（因为每个token只用过topK个expert，不需要使用全量expert），这也是我们说MoE-layer是稀疏层的原因。

普通” top-k 路由策略有一个重要缺点是，门控网络可能收敛到只激活少数专家。这是一个自我强化的问题：如果一小部分专家在早期被不成比例地选中，那么这些专家将更快地被训练，这会导致训练负载不平衡。而且，相对其它训练不足的专家，这些更快训练的专家会输出更可靠的预测，它们将继续被更多地选中。这种不平衡的负载意味着其他专家最终会成为名副其实的累赘。

带噪声的 TopK 门控 (Noisy Top-K Gating)能够缓解这个问题，其在为每个专家预测的概率值中添加了一些高斯噪声。在MoE模型中加入噪声的原因主要有以下几点：

• 提高模型的鲁棒性和泛化能力。当模型在训练或推理阶段遇到不确定或嘈杂的数据时，鲁棒性较强的模型更能保持稳定的性能。
• 噪声会增加一部分的随机性，减少过拟合的风险。
• 加入噪声可以实现不同专家之间的负载均衡。

该方案还为专家选择添加了两个可训练的正则化项：最小化负载均衡损失会惩罚过度依赖任何一个专家，而最小化专家多样性损失会奖励对所有专家的平等利用。我们用下图展示下上述解决思路。

• 原始处理流程如下图标号1，其中门控函数是softmax。路由策略就是将输入乘以权重矩阵并应用 softmax。
• 但是，这种方法并不能保证专家的选择将是稀疏的。为了解决这个问题，我们首先对输入进行线性变换，然后再加上一个softmax，这样得到的是一个非稀疏的门控函数。对应下图标号2。
• 但是这样依然不够，因此我们在进行softmax之前，先使用一个topk函数，只保留最大的k个值，其他都设为-∞。这样对于非TopK的部分，由于值是负无穷，这样在经过Softmax之后就会变成 0，不会被选中，就得到了稀疏性。在这个基础上，我们在输入上再加上一个高斯噪声。此处对应下图标号3。

随机路由是另一种解决问题的方法，例如，在top-2 设置中的“最佳”专家是使用标准 softmax 函数选择的，而第二个专家则是半随机选择的（每个专家被选中的概率与其连接的权重成正比）。因此，排名第二的专家最有可能被选中，但不再保证其肯定被选中。

另一种解决该问题的方法是根据专家容量的阈值来做调整。这种方法会设置一个阈值，该阈值定义了任何一个专家可以处理的最大token数量。还是以 Top-2的专家选择为例，如果 top-2 中选择的任何一个专家都已达到容量，则选择下一个专家（Top3）来处理后续token。但是，这可能导致token溢出，即超出容量的token无法被指定专家处理。另外，也有研究提出了DSelect-k，这是一种平滑的top-k门控算法，其平滑特性优于传统的top-k方法。

3.2 专家

在 MoE 架构中，专家是指训练好的子网络（神经网络或层），它们专门处理特定的数据或任务。专家和门控机制都通过梯度下降与其他网络参数一起进行联合训练。MoE里的“专家”是一种拟人的形象化的说法，其实，专家在本质上是基于某种人类先验“知识”或“策略”的“跨范畴采样”。

3.2.1 特点

专家具备如下特点：

• 架构。在实际应用中，一般来说，MoE 中的每个专家都是具有相同架构的前馈神经网络。但是，我们也可以使用更复杂的体系结构。我们甚至可以通过将每个 Experts 实现为另一个 MoE 来创建“分层”MoE 模块。在某些情况下，并非所有 FFN 层都被 MoE 取代，例如Jamba模型具有多个 FFN和MoE 层。

• 参数子集：FFN层被分解为多个专家，每个专家实际上是FFN参数的一个子集。专家并不是对FFN的平均切分，实际上我们可以任意指定每个expert的大小，每个expert甚至可以大于原来单个FFN层，这并不会改变MoE的核心思想：对于一个token，部分专家的计算量要小于所有专家的计算量。

• 输入分割：不同的专家会专注于不同的主题。用更专业的术语来说，输入空间被“区域化”了（或者说更精细地划分知识空间）。假设某个 LLM 可能收到的请求是一个“完整的知识空间”，而MoE将输入数据根据任务类型分割成多个区域，并将每个区域的数据分配给一个或多个专家模型。

• 专注学习。每个专家模型可以专注于处理自己接受到的输入数据，学习数据中的一种特定模式或特征。由于这些专家从一开始就存在，在训练过程中，每个专家都会在某些主题上变得更加专业，而其他专家则会在其他主题上变得更加博学。例如，在图像分类任务中，一个专家可能专门识别纹理，而另一个专家可能识别边缘或形状。

• 灵活扩展和组合作战。在MoE范式下，只有相关的专家被激活以处理给定输入，由于只有相关的专家被激活，因此可以减少不必要的计算（帮助我们在扩大模型规模的同时保证计算量是非线性增加的），从而加快模型的推理速度并降低运算成本。而且，MoE可以在减少计算开销、未相应增加计算成本的情况下扩展模型的参数空间，从而受益于大量专业知识。用户不必聘请一位“无所不知”的专家，而是组建一个拥有特定专业领域的团队。这种分工有助于整个模型更高效地处理问题，因为每个专家只处理它最适合的数据类型。另外，MoE这也使得模型能够更加灵活地适应不同的任务，因为不同的任务可能需要不同专家的组合来达到最优的预测效果。

我们再用一个示例来看看专家学习到了什么。

从目前的研究成果来看，专家并不专攻“心理学”或“生物学”等特定领域。它最多只是在单词层面学习句法信息：更具体地说，它们擅长于在特定上下文中处理特定的 tokens。专家学习的信息比整个领域的信息更加细粒度。因此，有时将它们称为“专家”可能会产生误导。

Mixtral 8x7B 论文中可以找到一个很好的例子，论文作者测量了所选专家在the Pile验证数据集不同子集上的分布（token分布比例）。下图显示了第0、15和31层的结果（第0层和第31层分别是模型的第一层和最后一层）。论文在根据主题分配专家时没有观察到明显的模式。例如，在所有层次上，ArXiv论文、生物学和哲学文档的专家分配分布都非常相似。真有针对数学的专家的分布才略有不同。这种差异可能是数据集的合成性质及其对自然语言的有限覆盖的结果，在第一层和最后一层尤其明显，其中隐藏状态分别与输入和输出嵌入非常相关。这表明门控网络确实表现出一些结构化的句法行为。

因此，尽管专家似乎没有专业知识，但他们似乎确实被一致地用于某些类型的token。下图显示了来自不同领域（Python代码、数学和英语）的文本示例，其中每个标记都用与其所选专家对应的背景颜色突出显示。该图显示，Python中的“self”和英语中的“Question”等单词经常通过同一个专家传递，即使它们涉及多个标记。同样，在代码中，缩进标记总是分配给相同的专家，特别是在隐藏状态与模型的输入和输出更相关的第一层和最后一层。我们还从图中注意到，连续的token通常分配给相同的专家。事实上，论文作者确实在The Pile数据集中观察到了一定程度的位置局部性。

因此，也有研究认为，专家提升的是记忆效果而不提升推理能力。比如论文“Mixture of Parrots: Experts improve memorization more than reasoning”研究了Mixture-of-Experts（MoE）架构性能和推理上的理论局限性，探讨了与标准密集型Transformer在记忆和推理方面的性能差异。研究发现，随着专家数量的增加，MoE模型在记忆任务上的表现提升，而在推理任务上达到饱和。论文也通过实证证明了MoE在特定记忆密集型任务上的优越性。

另外，MoE模型实际上还提供了一种细粒度的方式来研究和理解模型内部的工作机制。通过观察哪些专家被激活以及它们如何随着时间变化，研究人员可以更深入地洞察模型是如何学习和泛化知识，以及它是如何处理不同的输入特征的。

3.2.2 种类

专家的网络类型通常有如下几种：

• 前馈网络（Feed-Forward Network）：因为FFN层比自注意力层更加稀疏，且展示出更多的领域特定性，所以目前大多数MoE都是FNN的替代品。比如，有研究人员发现，大多数输入仅激活FFN的少量神经元，突显了FFN的内在稀疏性。对于同样输入，FFN层仅激活20%的专家，而自注意力层激活了80%的专家。预训练Transformer中的模块化涌现现象（Emergent Modularity）也揭示了神经元激活与特定任务之间的显著关联，支持了MoE结构反映预训练Transformer模块化特性的观点。另外，从参数量的角度我们也可以看到选择的原因，因为随着模型规模增长，FFN的计算开销呈现急剧上升趋势。例如，在拥有5400亿参数的PaLM模型中，约90%的参数分布在FFN层。
• 注意力（Attention）：尽管MoE研究主要集中在Transformer架构的FFN层，也有研究人员提出了多头注意力专家混合（Mixture of Attention Heads, MoA），将多头注意力层与MoE结合，以提升性能并降低计算成本。MoA使用两组专家（查询投影和输出投影），通过共同的门控网络选择相同的专家。为降低计算复杂度，MoA在所有注意力专家间共享$W_K$和$W_V$投影权重，专家仅在各自的查询（$q_t{W}_t^q$）和输出投影权重（$o_{i,t}{W}_i^O$）上有所区别，从而实现键（$K{W}_K$）和值（$V{W}_V$）序列的预计算共享。
• 其他类型。有些研究人员还探索了使用卷积神经网络（CNN）作为专家，也有将参数高效微调（PEFT）技术与MoE结合的努力，例如采用低秩适应（LoRA）作为专家。

3.2.3 位置

我们接下来看看专家如何嵌入到Transformer架构中。下图给出了一些实例。

• (a)展示了MoE与注意力机制中的Key和Value模块的集成。
• (b)表示MoE在FFN中的应用。
• ©指的是MoE在Transformer块层级的集成，其中应用了两组不同的专家到注意力和FFN层，分别为每个层分配专家，并通过各自的门控机制进行调控。
• (d)展示了MoE在每一层的集成，其中每个Transformer层视为一个统一体，门控机制协调专家之间的交互。

3.3 分类

之前的学习中，我们大致了解到了MoE有不同种类，此处我们再从最本质的角度（门控函数）来看看如何把MoE分类。门控函数是MoE架构的核心，它负责协调专家网络的参与并整合其输出。根据输入处理方式的不同，门控机制可分为稀疏型、稠密型和软性三类。稀疏门控只激活部分专家，稠密门控激活所有专家，而软性门控包括全微分方法，如输入token合并和专家合并。三种门控函数的特点如下：

• 稀疏门控：仅激活部分专家，包括基于token选择的top-k门控策略，以及使用辅助损失函数来促进专家间token均匀分布。
• 稠密门控：激活所有专家，在LoRA-MoE微调中表现出色，因为它可以有效地将多个LoRA整合到各种下游任务中。
• 软性（soft）门控：通过token或专家合并的方式实现完全可微性，避免了离散专家选择的问题，例如SMEAR、Lory和Omni-SMoLA。

根据门控函数的设计，MoE层可以大致分为以下两类：稠密MoE和稀疏MoE。

• 稠密MoE层在每次迭代中激活所有专家网络$f_1...f_N$。稠密MoE能充分利用所有参数，捕获潜在的复杂的模式和关系，因此通常能获得更高的预测精度，但计算开销较大。因此这种方法在早期研究中被广泛采用，近期有研究（EvoMoE、MoLE、LoRAMoE和DSMoE）重新探讨了稠密MoE的应用。因为稀疏激活专家虽然在计算效率上有优势，但当总参数量相同时，往往会导致性能损失。而在LoRA-MoE微调中，由于LoRA专家的计算开销较小，稠密激活表现出色。这种方法能够有效地将多个LoRA整合到各种下游任务中，既保持了原始预训练模型的生成能力，又维持了每个任务特定LoRA的独特性
• 稀疏MoE层在每次前向传递中仅激活选定的专家子集。稀疏型MoE不是汇总所有专家的输出，而是通过仅计算前k个专家输出的加权和来实现稀疏性。稀疏激活实际上是计算需求与模型性能之间的一种权衡策略。

下图展示了两种MoE的特点。右侧是示意图，左侧是门控函数以及负载函数。

3.3.1 稠密 vs 稀疏

我们通过例子来对稠密和稀疏MoE进行比对。

以人类分工为例。稠密就是类似手工业时代的生产模式。在这种模式下,每个工人（即神经元）都需要参与处理所有类型的任务，就像手工业时代的工匠需要精通产品制作的各个环节，掌握所有的生产技能。这种方法虽然直观且易于实现,但在面对复杂多变的任务时,往往效率低下且难以扩展。稀疏MoE则是工业革命之后的分工模式：每个岗位（专家）只需要完成一部分生产任务。这种分工的方式大大提高了生产效率，推动了工业化的进程，开启了机器大工业的时代。

3.3.2 软性（soft）门控

论文”A Survey on Mixture of Experts“中提出软性（soft）门控这种类型，以突出其通过门控加权合并输入token或专家来缓解计算需求的特点。为了为每个输入token分配适当的专家，稀疏性MoE通常需要启发式辅助损失来确保专家参与的平衡。在涉及分布外数据的场景中，这些问题变得更加突出。与密集MoE类似，软MoE方法通过利用所有专家处理每个输入来保持完全的可微性，从而避免了离散专家选择所固有的问题。

• token合并（Token Merging）：Soft MoE摒弃了传统的稀疏和离散门控机制，采用软分配策略来合并token。该方法计算所有token的加权平均值，权重取决于token和专家的关系，然后用相应专家处理每个加权平均结果。然而，token合并使其难以应用于自回归解码器，因为在推理时无法获取用于加权平均的未来token。
• 专家合并（Expert Merging）：SMEAR框架通过加权平均合并所有专家参数来避免离散门控。SMEAR的作者认为，传统稀疏MoE模型难以匹敌参数量相当的稠密模型或使用非学习启发式门控函数的模型，可能是由于非可微、离散门控决策的训练模块中存在梯度估计偏差。SMEAR通过单个合并专家处理输入token，既不显著增加计算成本，又支持标准梯度训练。

3.4 比对

论文“A Uniffed View for Attention and MoE”把注意力机制和MoE做了比对。Attention 结构的作用其实就是使用加权和的形式来聚合不同 token 的信息。MoE 其实是通过线性投影来学习一个 Router, 基于 Router 来聚合不同 expert 的信息. 本质上来说, router 的输出经过 top-k+softmax 来学习一个加权求和的权重, 最终聚合不同 token 的信息。

两者的相似之处:

• 都使用了 softmax 来归一化权重，本质上也是希望学习稳定。
• 本质都是对新特征的加权和。
• 都是动态权重, 根据输入来自适应地聚合信息。

二者的区别在于

• Attention 是聚合不同 token 的信息，MoE 是聚合不同地 Expert。
• Attention 会用上所有 token (Softmax 结果为正数，故所有 token 都会用上)；MoE 是一种稀疏选择，只会选择部分结果。

Attention和MoE都是学习一个权重来聚合信息，可以总结到如下图标号3所示，f是权重, g 是学习到的新特征。具体如下图。

从计算机体系结构的视角来看也有启发。注意力机制（含FFN）类似把所有上下文都放在内存中。MoE则可以理解为页表系统，需要的时候才把专家对应的页表放到内存中。

0x04 计算流程

4.1 算法

下图是top-k门控函数算法的伪代码，具体流程如下。

• 给定输入x，使用门控函数计算出得分score = G(x)。
• 选出前k的分数对应的索引。
• 遍历索引，使用索引对应的专家来进行计算，得到推理结果。
• 使用分数作为权重对专家计算的结果值进行修正。
• 综合这k个专家的结果作为最终推理结果。

4.2 流程

MoE的整个计算过程如下图所示：

• Routing。Routing也叫“experts selection”或者稀疏性激活，本质上是对门控函数的使用，是MoE的核心理念。Routing是一个对输入进行多分类的鉴别过程，目的是确定最适合处理输入的专家模型。在语言模型的应用中，当输入token通过MoE层时，Token通过和Router的权重矩阵相乘得到一个Expert Indices（决策矩阵）和一个概率张量，即索引和概率：
  • Expert indices是expert-to-token映射，用于指示每个token被分配给了哪个expert。即张量中第i个值代表本token应该分配到第i个专家。
  • Probabilities张量是分配置信度的概率，其中第i个值代表这个专家对于该token最终结果的权重。
• Permutation（排列/置换）。根据路由决策（expert-to-token映射）将Token分配给对应的专家。其间可能会依据专家的容量对token进行丢弃操作。
• Computation。专家依据分配到的经过Permutation重排序的tokens进行计算。通过使每个专家专注于执行特定任务，这一方法实现了计算的高效性。这种方式允许模型对不同类型的输入数据进行个性化处理，提高了整体效率和性能。每个专家网络并行处理其分配到的token，计算输出。这一步涉及到块稀疏矩阵乘法，其中输入矩阵 𝑥″ 与专家网络的权重矩阵相乘。设 𝑊 为专家网络的权重矩阵，则专家网络的输出可以表示为： 𝑦=𝑥″×𝑊 ，其中 𝑦 为所有专家网络输出的汇总结果。
• Un-Permutation。收集专家的计算结果。这是Permutation的逆运算，目的是为了将从各个experts收集到的处理后的tokens组合成一个完整的序列，这个序列保持了原始tokens的顺序。即，将每个专家网络的输出根据原始的token顺序重新排列。设 𝑦 为所有专家网络输出的汇总结果，则反排列操作可以表示为： 𝑦′=scatter(𝑦,𝑖𝑛𝑑𝑖𝑐𝑒𝑠) ，其中 𝑦′ 为最终的模型输出， scatter 操作根据 𝑖𝑛𝑑𝑖𝑐𝑒𝑠 中的索引将 𝑦 中的输出重新排列以匹配原始的token顺序。接着使用Routing步骤生成的分配置信度概率对结果进行缩放（加权求和），以得到最终的模型输出，然后将这个结果继续向下游处理。

因为Permutation相对复杂，我们接下来仔细分析。

4.3 Permutation

Permutation的主要作用是：

• 分发token。Permutation会依据Expert Indices构建本地的置换Token位置的后的临时矩阵（将输入Token根据路由结果重新排列），这样可以把属于每个专家的token分别放在一起，然后把tokens送给对应的专家。使得每个专家网络可以并行处理其分配到的token，以确保模型可以充分利用GPU的并行计算能力。比如上图中，“the”和“jumped”应该分配给专家1，所以就把它们放在一起。“quick”和“fox"都应该被发送给专家2，所以把它们也放在一起。
• 维持token和expert的顺序。因为一个batch里有很多token，我们将这些token发往不同的expert做计算后，专家输出结果的顺序肯定是打乱的，所以需要通过一种方式追踪顺序，把token permute回正常的位置再输入下一层网络。通过构建的矩阵，Permutation在计算时，就可以维护住这种顺序。
• 负载均衡。Permutation可以实现输入数据在不同专家之间的合理分配，以平衡各个专家的计算负载。不同的输入样本可能对不同专家的计算资源需求不同。通过对输入样本进行置换，使得每个专家能够相对均匀地接收到需要处理的样本，避免某些专家过度使用而其他专家闲置的情况。
• 增加多样性。Permutation可以增加模型对输入数据处理的多样性。因为不同的置换顺序可能会导致不同专家组合对数据进行处理，从而挖掘出数据的不同特征。

从代数的角度来看，MoE计算实际上是对token进行一次置换群的操作：$Pconcat(Experts)P^{-1}$。P（permutation操作）为一个进行Token位置置换的稀疏矩阵，实际上也构成了代数上的一个置换群的结构，$P^{-1}$需要对Token进行还原，保证原有的Token顺序输出到下一层。MoE实现的本质问题是：基于Permutation矩阵后构建的稀疏矩阵乘法如何进行并行。

4.4 实现

4.4.1 Mistral Inference

我们首先用https://github.com/mistralai/mistral-src来学习，此代码相对简单易懂。其中，gate是门控函数，在使用时会用配置如下：gate=nn.Linear(dim, moe.num_experts, bias=False)。

假设我们定义了4个专家，路由取前2名专家，即num_experts=4， num_experts_per_tok=2，同时词嵌入大小为32。MoE接收注意力层的输出作为输入X，即将输入从（batch_size，sequence_length，input_dim）的形状[2, 4, 32]投影到对应于（batch_size，sequence_length，num_experts）的形状[2, 4, 4]，其中num_experts即expert=4。然后通过torch.topk 将张量转换为[2, 4, 2]。torch.topk返回的selected_experts可以理解为对于每个token选中的两个专家的序号索引。

具体代码如下。

import dataclasses
from typing import Listimport torch
import torch.nn.functional as F
from simple_parsing.helpers import Serializable
from torch import nn@dataclasses.dataclass
class MoeArgs(Serializable):num_experts: int # 专家数量num_experts_per_tok: int # 每一个token被分配给几个专家# gate=nn.Linear(dim, moe.num_experts, bias=False)
class MoeLayer(nn.Module):def __init__(self, experts: List[nn.Module], gate: nn.Module, moe_args: MoeArgs):super().__init__()assert len(experts) > 0self.experts = nn.ModuleList(experts)self.gate = gateself.args = moe_argsdef forward(self, inputs: torch.Tensor) -> torch.Tensor:gate_logits = self.gate(inputs) # 通过门控网络获得各个专家的logits# 取出topk（k=2）专家的权重以及专家索引weights, selected_experts = torch.topk(gate_logits, self.args.num_experts_per_tok)#使用softmax来归一化权重weights = F.softmax(weights, dim=1, dtype=torch.float).to(inputs.dtype)# 创建形状和x一致，初始值为0的矩阵，用来存储每个expert的输出results = torch.zeros_like(inputs)for i, expert in enumerate(self.experts): #遍历每一个专家# selected_experts == i 得到的是一个矩阵，行为token的idx，列为专家的idxbatch_idx, nth_expert = torch.where(selected_experts == i)# 每一个token的结果都是由2个专家的结果进行加权求和得到的  # 利用None来增加维度，新增维度大小为1，有几个None就会增加几个维度。results[batch_idx] += weights[batch_idx, nth_expert, None] * expert(inputs[batch_idx])return results

使用代码如下，可以看到，门控网络就是nn.Linear的实例，而专家网络是FeedForward的实例。

class FeedForward(nn.Module):def __init__(self, args: TransformerArgs):super().__init__()self.w1 = nn.Linear(args.dim, args.hidden_dim, bias=False)self.w2 = nn.Linear(args.hidden_dim, args.dim, bias=False)self.w3 = nn.Linear(args.dim, args.hidden_dim, bias=False)def forward(self, x) -> torch.Tensor:return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))class TransformerBlock(nn.Module):def __init__(self, args: TransformerArgs):super().__init__()self.n_heads = args.n_headsself.dim = args.dimself.attention = Attention(args)self.feed_forward = MoeLayer(experts=[FeedForward(args=args) for _ in range(args.moe.num_experts)],gate=nn.Linear(args.dim, args.moe.num_experts, bias=False),moe_args=args.moe,)self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)self.args = argsdef forward(self,x: torch.Tensor,freqs_cis: torch.Tensor,positions: torch.Tensor,mask: Optional[torch.Tensor],) -> torch.Tensor:r = self.attention.forward(self.attention_norm(x), freqs_cis, positions, mask)h = x + rr = self.feed_forward.forward(self.ffn_norm(h))out = h + rreturn out

4.4.2 Mixtral 8x7B

我们在用Mixtral 8x7B来学习，此实现更加复杂。代码来源：transformers/src/transformers/models/mixtral

下面是MoE层的代码。此处要重点介绍下代码中的掩码矩阵。因为在MoE中，并不是所有的输入都要经过所有的专家。通常只有一部分输入经过特定的专家进行处理。因此使用掩码矩阵来决定决定哪些输入token与哪些专家进行交互。掩码矩阵的使用有多种可能，比如对于一个输入batch，掩码矩阵可以：

• 指定每个输入样本是由专家 1 处理，还是专家 2 处理，或者专家 3 处理，或者是其中几个专家的组合来处理。
• 也可以当掩码矩阵中的某个元素为 1，表示对应的token和专家之间有连接（即该专家会处理这个输入），0则表示没有连接。
• 也可以结合专家容量用到填充上。

另外，在训练过程中，Mask 矩阵还用于正确地计算和更新专家的参数。因为只有与输入有连接的专家才应该对这个输入的损失产生贡献，通过 Mask 矩阵可以准确地计算每个专家的梯度，从而正确地更新专家的参数，避免错误地更新那些没有处理特定输入的专家的参数。

class MixtralSparseMoeBlock(nn.Module):"""This implementation isstrictly equivalent to standard MoE with full capacity (nodropped tokens). It's faster since it formulates MoE operationsin terms of block-sparse operations to accomodate imbalancedassignments of tokens to experts, whereas standard MoE either(1) drop tokens at the cost of reduced performance or (2) setcapacity factor to number of experts and thus waste computationand memory on padding."""def __init__(self, config):# 初始化了MoE层中使用到的各个部分参数super().__init__()self.hidden_dim = config.hidden_sizeself.ffn_dim = config.intermediate_size# 设置专家数量self.num_experts = config.num_local_experts# 设置要选择的topk专家数量self.top_k = config.num_experts_per_tok# gating# 初始化线性层作为门控机制self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)# 创建专家网络的列表，每个专家是一个 MixtralBlockSparseTop2MLP 实例self.experts = nn.ModuleList([MixtralBlockSparseTop2MLP(config) for _ in range(self.num_experts)])# Jitter parametersself.jitter_noise = config.router_jitter_noisedef forward(self, hidden_states: torch.Tensor) -> torch.Tensor:""" """# 将注意力模块的输出隐状态hidden_states作为专家模块的输入。        batch_size, sequence_length, hidden_dim = hidden_states.shapeif self.training and self.jitter_noise > 0:hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise)# 原有hidden_states形状为[batch_size, sequence_length, hidden_dim]，为方便计算，将batch_size和sequence_length合并，重构为一个形状为(batch * sequence_length, n_experts)的二维张量。而且，后续计算不是样本维度，而是token维度，这样更清晰hidden_states = hidden_states.view(-1, hidden_dim)# 将hidden_states导入门控网络中，输出一个路由逻辑用于后续专家分配# 计算每个专家的分数，router_logits的形状为[batch_size * sequence_length, n_experts]router_logits = self.gate(hidden_states)# 计算专家经过softmax之后的概率，最后选取top k个专家用于输出（在dim=1进行softmax处理，即对应路由逻辑中的n_experts以及hidden_states中的hidden_dim）routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)# 通过门控函数获得分数最高的 top_k 和门控权重和专家索引# selected_experts和router_weight的形状都是 (b * s, top_k)routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)# 最后对专家权重进行归一化，确保权重之和为1，并且将权重的类别和输入统一routing_weights /= routing_weights.sum(dim=-1, keepdim=True)# we cast back to the input dtyperouting_weights = routing_weights.to(hidden_states.dtype)# 将token导入对应的专家网络中进行前向传播得到输出        final_hidden_states = torch.zeros( # 初始化全零矩阵，后续叠加为最终结果(batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device)# One hot encode the selected experts to create an expert mask# this will be used to easily index which expert is going to be sollicitated# 生成专家掩码，具体是根据专家网络的总数量构建一个one_hot编码，根据输入token所分配的专家网络，将one_hot编码中对应的专家网络索引由0变为1，据此来索引给定的专家网络，掩码原始形状是是 (b * s, top_k, expert_number)，permute之后形状是(expert_number, top_k, b * s)expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)# Loop over all available experts in the model and perform the computation on each expert# 在循环中依次通过各个专家网络处理每个输入tokenfor expert_idx in range(self.num_experts):#依次取出专家网络expert_layer = self.experts[expert_idx]# expert_mask[expert_idx] shape 是 (top_k, b * s)# idx 和 top_x 都是一维张量# idx 的值是 0 或 1, 表示这个token认为当前专家是 top1 还是 top2# top_x 的值是 token 在展平之后的 batch*seq_len 个token中的位置索引            idx, top_x = torch.where(expert_mask[expert_idx])# Index the correct hidden states and compute the expert hidden state for# the current expert. We need to make sure to multiply the output hidden# states by `routing_weights` on the corresponding tokens (top-1 and top-2)# 找出当前专家网络应该处理的，对应第top_x个token所对应的隐向量hidden_states# hidden_states 的形状是 (b * s, hidden_dim)current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)# 将hidden_states传入到专家网络当中进行处理# router_weights的形状是 (b * s, top_k)current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]# However `index_add_` only support torch tensors for indexing so we'll use# the `top_x` tensor here.# 将输出存入之前定义的张量当中final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))# 将计算完成的输出形状还原为之前的形状 (b * s, expert_number)    final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)return final_hidden_states, router_logits

专家网络的代码如下，Mixtral 8x7B使用了一个三层的MLP来作为MoE中的专家网络。MLP中的三个全连接层均参与了前向传播，最后对全连接层的输出和专家网络的输出进行加权得到最终的输出。此处与Mixtral of Experts代码的不同在于把激活函数用配置项来表示，这样更灵活，比如hidden_act=“silu”。

class MixtralBlockSparseTop2MLP(nn.Module):def __init__(self, config: MixtralConfig):super().__init__()self.ffn_dim = config.intermediate_sizeself.hidden_dim = config.hidden_sizeself.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)self.act_fn = ACT2FN[config.hidden_act]def forward(self, hidden_states):current_hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3(hidden_states)current_hidden_states = self.w2(current_hidden_states)return current_hidden_states

4.5 参数量

MoE 之所以有趣，很大一部分原因在于其计算要求。由于在给定时间内只使用专家的一个子集，因此我们可以访问比正在使用的更多的参数。但是如何计算MoE的参数量以及内存占用？

当前MoE 大模型有两种“署名”方法，各有优劣。

• 类似“8x22B”这样的名称。这表示模型有8个专家，每个专家有22B 参数。优势是清晰的给出了专家个数和规模，劣势是容易让大家误认为模型一共只有8x22B这么多参数。
• 类似”57BA14B”这样的名称。表示总参数规模一共57B，每次推理激活参数14B。但不清楚里面有多少专家，每个专家的大小。

因此，有研究人员建议使用如下方式来对MoE模型进行标记：总参数量-激活参数量-普通专家数量-激活专家数量-共享专家数量，我们接下来解释下。

• 总参数量。实际上专家只是模型的一部分，还有注意力层、embedding层，LM Head、门控网络等模块。总参数量通常被称为稀疏参数量，可以理解为模型容量的衡量标准，也是在推理时需要多大的显存 VRAM 的参考数值，因为模型的所有参数都必须加载到显存/内存中。

• 激活参数量。实际上用于处理单个token的参数数量，因为该token只通过一些专家块但不通过其他专家块，即在推理时，我们只使用了部分专家，所以激活的参数比较少。该指标可以理解为模型推理时计算成本的衡量标准。换句话说，MoE 模型需要更多的 VRAM 来加载整个模型（包括所有专家），但因为推理时只用到了部分专家，因此MoE在推理过程中的激活的参数较少、运行速度更快。

• 普通专家数量-激活专家数量-共享专家数量。这个依据不同MoE的实现不同而不同。因为有的MoE没有共享专家，只有专家。假如每一个Transformer层有64个普通专家和一个共享专家，激活专家数量为8，则推理时会从64个普通专家中选择8个专家进行激活。该层推理时主要加载的参数为：激活专家数 * 专家大小 + 共享专家数 * 专家大小 + 注意力层大小 + 门控函数大小。

另外，即使每个输入只使用一部分参数，所有专家的完整参数集通常也需要加载到内存中，这可能会增加推理过程中的整体内存占用 。

4.6 计算量

在推理的实际计算时，每个 Token 都会经门控网络来选择 1 个或多个 Expert，然后执行类似 FFN 的计算。我们可以参考论文“Scaling Laws for Fine-Grained Mixture of Experts”给出的MoE计算量的公式，具体指标如下：

• G（Granularity） ：粒度，也就是将一个完整的 FFN 切分为多少个细粒度的专家，比如一个 FFN 可以切分为 8 个小专家，每个专家的参数量为原始 FFN 的 1/8。

• E（Expansion rate） ：膨胀率，即所有专家的参数量扩展为单个标准 FFN 参数量的多少倍。比如，G=8，总共 64 个细粒度专家，则这里的 E=64/8=8，相当于有 8 个同标准规模的 FFN。该指标可以理解为MoE参数总数是其激活参数的比例。

• D：token数目。

• cf：表示除 Router 之外的计算量与参数量的比值，论文中设置为 6。对于每个token，线性层的一个活跃参数的FLOPs是6。操作分解如下：

  • 在前向传播过程中，使用2个操作（单次乘法和单次加法）来计算输入和线性投影的矩阵乘法。
  • 在反向传播过程中，使用2个操作来计算输入的梯度
  • 在反向传播过程中，使用2个操作来计算线性层权重的梯度。

• cr：表示 Router 的计算量与参数量的比值，对于比较简单的 Router，通常为 6-20。具体拆解如下：

  • 在前向传播过程中，使用2个操作来根据输入和“路由线性层”计算专家logits。
  • 在反向传播过程中，使用2个操作来计算“路由线性层”关于输入的梯度。
  • 在反向传播过程中，使用2个操作来计算“路由线性层”关于线性层的梯度，
  • 在前向传播过程中，使用2个操作将输入token路由到选定的专家。
  • 在前向传播过程中，使用2个操作将专家的输出路由回选定的token，并将这些输出乘以路由分数。
  • 在反向传播过程中，使用2个操作将梯度从输出token路由回到专家。
  • 在反向传播过程中，使用2个操作将梯度从专家路由到输入token。

  与cf的FLOP计算类似，cr的FLOP也是成对的，因为每次乘法后都是加法（用于累加输出或梯度）。

由于 Router 的计算量通常很小，常常可以忽略。如果不考虑 Router 引入的计算量，由于上述每个 Expert 实际上与非 MoE 模型的 FFN 一样，因此也可以用同样的方式推导出训练时每个 Token 的计算量依然为 C=6ND，只不过这里的 N 不再是整个模型的参数量，而是每个 Token 激活的参数量。

0x05 并行计算

MOE Transformer layer的并行方式一般如下：

• 非专家部分（注意力机制）采用张量并行和数据并行；
• 专家部分采用专家并行和张量并行；

5.1 通讯需求

我们首先看看MoE计算中的通讯需求。

5.1.1 单token

下图是一个MOE结构的Transformer layer的计算过程，其中：

• w1/w2表示自注意力模块的输入；
• a1/a2表示自注意力模块的输出，FFN的输入；
• f1/f2表示FFN的输出；

具体分为五步：1) 整理，2）发送，3）计算，4）原路返还，5) 加权求和。其中第2，第4步存在通信需求。

5.1.2 多token

上述是针对每个token单独操作，所以看起来很简单。但要知道，我们通常是同时处理一个batch的所有token，所以上述操作，无论是传输还是计算，都要针对矩阵做改变。此外，我们也不知道对于单个token，它选择的k个专家到底在哪些卡上。因此，MoE的实现就复杂在整理和传输之上。

实际中的通信操作如下图所示，在专家并行模式下，专家层的前后会分别引入 All-to-All 通信操作。前一个 All-to-All 用于将每个 Worker 上的 Token 按照 Router 后对应的专家发送到专家所在的 GPU，也叫 All-to-All Dispatch；后一个 All-to-All 用于将专家计算后的 Token 重新按照原来的方式排列，也叫 All-to-All Combine。可以看出来这里有两步通信需求：

• 把token发给专家（对于下图标号1）。需要把token按照expert-to-token映射序列发给对应的GPU。
• 收集k个计算结果（对于下图标号2）。需要对token进行还原，保证按照原有的token顺序输出到下一层。

上面提到了两个概念：专家并行和All-to-All通信，我们分别来分析下。

5.2 专家并行

在面对庞大而复杂的模型时，如何高效利用计算资源、提升训练速度，一直是研究者们关注的焦点。传统的做法通常是通过数据并行和模型并行的方式来进行加速计算，但这种方式在处理MoE模型时不是很适合，因为MoE 的特点是工作负载的稀疏性和动态性。因此，研究人员提出了专家并行（expert parallelism），通过分发切分后token（dispatching partitioned local tokens）并限制专家容量的负载均衡，来实现并行门控和专家计算。目前，专家并行已成为促进MoE模型高效扩展的基本策略。

5.2.1 定义

专家并行在本质上是一种模型并行方法，但也可以看作是数据并行的扩展。专家并行的思路是将MoE层中不同的专家分配到不同的计算设备上，每个设备负责存储和计算部分专家，而所有非专家层则在设备间复制。专家并行的流程包括以下几个顺序操作：门控路由、输入编码、All-to-All dispatch、专家计算、All-to-All combine和输出解码。

• 每个 EP rank 上只包含一部分 expert，而每个 EP rank 上的 token（即 token 对应的hidden state） 会根据 门控路由的结果分发到其他 EP rank 上的 expert。
• 和数据并行方案类似，每一张卡的输入是一个完整的任务，需要在每张卡任务的Batch和序列维度进行掩码，拆解为不同的子任务。因此，专家并行采用输入编码（input encode）将发往给同一专家的输入令牌（token）聚合到一个连续的内存空间中，这一空间由门控路由token-expert映射决定，该操作就是permutation。
• 随后，每个设备会根据MoE模型的路由规则，使用All-to-All dispatch将输入token（各自的任务）发送到相应专家所在的设备上。
• 经过专家的本地计算后，逆过程（All-to-All combine和输出解码）会根据门控索引来恢复原始的数据布局，将结果返回到原设备，这样就在每张卡上还原为完整的任务。

因为专家们被精心地分散部署于各个节点之上，但每个节点所承担的任务却有所不同。每个节点都拥有独特的专家资源，并且数据也被巧妙地分割，确保在所有节点之间实现均衡分配，使得计算资源得到了更加充分的利用。

注：为了达到计算设备所需的最佳利用率和吞吐量，通用矩阵乘法（GEMM）的输入大小需要足够大。

5.2.2 历史

业界鼻祖

论文”Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer“第一个用到了专家并行方案，论文给出的描述是：”Mixing Data Parallelism and Model Parallelism“。

论文遇到的问题如下。在现代CPU和GPU上，需要使用大的batch size来提高计算效率，这样可以分摊参数加载和更新的开销。如果门控网络为每个示例从n个专家中选择k个，那么假设一个batch中有b个样本，每个专家都会收到$\frac{kb}{n}\ll b$个样本。如果专家数量增加，这会导致原始（native）MoE实现的效率很低，即平均到每个 Expert 上的样本会远小于 b，这种方式很不利于模型的扩展。解决这个batch收缩（shrinking batch）问题的办法是使原始batch尽可能大。然而，batch大小往往受到存储前向和方向传播激活所需内存的限制。

为了解决这个问题，论文采用以下技术来增加batch size。论文让不同的batch同步运行，这样它们可以组合起来为MoE使用。论文根据传统的数据并行方案那样来分配模型中其它常规层和门控网络，但是每个专家只保留一个共享的拷贝。MoE层的每个专家都会收到一个由所有数据并行输入batch中相关样本的组合。同一组设备既可以作为数据并行副本（其它常规层和门控网络），也可以作为模型并行分片（每个设备承载专家集合中的部分专家）。如果模型分布在d个设备上，每个设备处理一批大小为b的数据，则每个专家都会收到一批数量大约是$\frac{kbd}{n}$个样本。因此，我们实现了专家batch size的d倍增大。

Gshard

Gshard则完善了MoE跨设备分片的方法。当扩展到多个设备时，MoE 层在不同设备间共享，而其他所有层则在每个设备上复制。这样，整个 MoE 层的计算被分散到了多个设备上，每个设备负责处理一部分计算任务。这种架构对于大规模计算非常有效。Gshard让人们真正意识到，只要把层和 token 智能拆分并均衡地分配给各个专家，训练几百亿甚至上千亿参数的模型是可以做到的。

Switch Transformers

Switch Transformers 应该是第一个显式提出来专家并行概念的论文。论文作者试图平衡每个令牌的FLOPS和模型参数之间的关系。因为如果用原生MoE实现，当扩大专家数量时，我们虽然会增加参数的数量，并不会改变每个令牌的FLOP。为了增加FLOP，我们还必须增加$d_{ff}$的维度（这会增加参数，但会导致速度较慢）。这导致我们需要进行一系列的平衡：当增加$d_{ff}$时，每个core的内存将耗尽，这就需要增加m。但由于core的数量是固定为N，并且N=n×m，我们必须减少n，这又迫使使用较小的批处理大小（以便保持每个核心的令牌常数）。因此论文采用了专家并行。

注：假定N是core的总数，n是数据并行度，m是模型并行度，E是专家数，C是专家容量。

小结

专家并行首先将experts分散保存到不同的设备上，然后将走同一数据流路径（data path）的输入tokens组合在一起构成token group（这些tokens选择了同样的expert），随后让不同的token groups（不同的data path）同时做并行计算。

在非专家并行方案中，如果专家数目越多，则每个激活专家受到的batch size会越小。而专家并行让每个专家收到的batch size与专家并行数成正比。只要 batch size 大小（d*b，d个设备，每个设备处理一批大小为b的数据） 随着专家数的增加而增加，就可以保证每个专家收到的样本数为常数，对显存和带宽的需求也就基本是常数，这样就减少了访存开销（data access）并最大化利用了内存带宽。

5.2.3 协同

随着模型规模的增加，往往需要结合各种分布式并行策略。包括数据并行（Data Parallel，DP）、模型并行（Model Parallel，MP，或者称为 Tensor Parallel，TP）和专家并行（Expert Parallel，EP）。下图给出了专家并行和其它并行方式结合的样例。因为不同专家被分配到不同的设备上，所以很好的解决了显存开销问题。而且，使用 EP 不会减少数据并行的数量，因为每个 EP 处理不同的数据。另外，因为 EP 和 TP 对通信的要求都很高，所以一般不会让 EP 和 TP 跨机。

一般情况下，单个expert不会被拆分(EP only)，但有时为了达到更好的性能，我们会使用更多的设备，其数量甚至会比expert数量还多，这种情况下还可以做expert-slicing，也就是把expert本身（MLP）做tensor-slicing并在多个设备上实现张量并行，这就是EP + TP。如果单个expert做tp切割，数据过单个expert后的输出结果一定会在同个ep_tp_group内做AllReduce。

分布式并行策略的选择影响计算效率、通信开销、内存占用之间的复杂相互作用，这可能受到不同硬件配置的影响。因此，实际应用中的部署策略需要在多个方面进行精细的权衡，并为特定的使用场景量身定制设计。

5.2.4 如何切分

下图是从模型参数切分和数据切分的角度（只考虑 FFN 层）来比较几种并行策略。

每个4×4的虚线网格代表16个core，阴影方块是该core上包含的数据（模型权重或一批令牌）。第一行：说明模型权重如何在core之间分配。阴影方块的大小标识了FFN层中不同的权重矩阵大小，每种颜色都标识了一个唯一的权重矩阵。第二行：说明数据batch如何在core之间分割。不同颜色代表不同令牌。每种策略的模型权重和数据张量分割如下。

• 第一列：数据并行。上方表明所有设备（1-16）都有相同、全部的模型参数。下方表明每个设备只有一个数据分片，且不重复，共 16 个数据分片。
• 第二列：模型并行。上方表明所有设备（1-16）都只有模型参数的一部分，共 16 个分片。下方表明所有设备使用共同的一份数据。
• 第三列：模型并行+数据并行，设备分为 4 组（1-4,5-8,9-12,13-16）。上方表明每组（4 个设备）都有完整的模型参数副本，但是每组内的设备只有参数的一部分。下方表明数据分为 4 个切片，每组（4 个设备）对应一个数据切片。
• 第四列：专家并行+数据并行，设备分为 16 组（1-16）。上方表明每一个设备都有不同的专家，共 16 个专家。下方表明每个设备都有不同的数据分片（Token），共 16 个数据分片，一个专家对应一个分片。
• 第五列：专家并行+模型并行+数据并行，有 4 组设备（1-4,5-8,9-12,13-16）。上方表明有 4 个专家，每个专家分布在对应的 4 个设备上，比如绿色专家分布在 5,6,7,8 设备上。下方表明有 4 个数据分片，每组设备（每个专家）对应一个数据分片，一组里的 4 个设备共享一份数据分片。

5.2.5 优势

EP主要优势如下：

• 大大减小了模型参数冗余，节省出大量的显存空间，为大batchsize提供显存空间基础。
• 消除了畸形矩阵的运算，增加AI值，解决bandwidth bound问题；
• 单卡的通信量不受整个实例的batchsize的影响，为大batchsize提供高效通信基础；
• 通过节省出的显存带来的大batch，可以将各个节点打满，实现稀疏模型到稠密模型的转化；

我们再对第一点做下解析。因为虽然MoE一次推理的激活权值只较少，但是依然需要将所有参数都载入内存，这些参数没有参与计算但是占用了大量的显存，导致batchsize很难打高。EP方案种，每张卡至少一个Expert，在保证足够大的batchsize和负载均衡的条件下，每一张卡的所有expert都处于一个完全激活的状态，消除了memory bound的问题，或者说把一个稀疏模型转化为一个稠密模型，也可以发挥出GPU的算力。

当然EP也有存在的问题，最大的问题可能是部署实例减少，且单部署实例故障概率增加，加上多卡之间的通信同步会大大增加出错的概率，最终导致整个系统的稳定性和可用性面临比较大的挑战。

5.3 All-to-All通信

可能我们会直观的认为，哪个Token通过Router决定了那个Expert，那么就直接通过NVLINK内存拷贝或者RDMA发送过去就好了，为什么需要AlltoAll的通信范式呢？我们仔细分析下。

5.3.1 困境

假设我们有2张显卡，4个专家。专家平均分布在2个显卡上，每个显卡上2个专家。在数据并行的场景下，每个显卡有自己独立的一个batch，在MoE层以外的地方，各个设备独立计算。一旦到了MoE层，数据在显卡间需要分发。因为当前设备上的某个token，它所需要的专家，可能坐落在其它设备上。

在MoE层，每张显卡需要传输哪些token给其它显卡呢？一种方法是，每张显卡把自己所有的a个token都发给其它显卡。每个显卡的每个专家都获得了2a个token。然而，对于每个专家，它实际只需要处理这2a个token中的部分token。因此，需要有一个机制让每个专家从自己获得的2a的token中找到自己应该处理的a/2个token，计算完成后，还需要把这a/2个token原路返回，还原成原来的形状。

我们看看这个机制。假设每个专家最多可以处理 p 个token，每个显卡有两个专家，则每个显卡最多处理 2p 个token。在传输前，每张卡把这些token整合成一个形状为 (2,p) 的张量，因为有两张显卡（本显卡也要发给自己），所以实际上每个显卡要：

• 整理出来形状为 (2, 2, p) 的张量，这是每个显卡要传给其它每个显卡的数据。假如每个token维度是m，则张量实际形状是 (2, 2, p, m) 。
• 对于第 1 个显卡，它自己保留 (∗,0,∗,∗) ，然后把 (∗,1,∗,∗) 传输给第 2 个显卡。第 2 个显卡保留 (∗,1,∗,∗) ，把 (∗,0,∗,∗) 发给第 1 个显卡。因此，每个显卡拿到的张量形状还是 (2, 2, p, m) 。
• 每个卡会发给每个专家形状为 (2, p, m) 的张量。

实际上，在集合通信中已经有了这个通信机制，这就是All-to-All通信。

5.3.2 All-to-All

在此通讯模式中，每个进程向每个其他进程发消息的一部分，最后每个进程拥有各个进程消息的一部分。All-to-All的作用相当于分布式转置Transpose操作。具体如下图所示。可以看到，GPU0把自己收到的4个绿色的块分配给了全部4个GPU 。

我们结合MoE来看看。我们需要通过All-to-All通讯将token发去指定的expert做计算，再通过All-to-All通讯将计算结果返回。假设我们有一个4张卡的GPU集群。下图标号1描绘了首次做All-to-All（All-to-All Dispatch）的过程，这个过程的目的是将token发去对应的expert上进行计算。对比一下左侧图和中间图的数据块排布，你会发现All-to-All就相当于做了一次矩阵转置。因此通过All-to-All，我们就让数据块去到了它对应的位置：A0、B0、C0和D0去GPU0，A1、B1、C1和D1去GPU1，以此类推。而为了实现这种转置，我们必须提前token做分块排序，让它按照要去的专家位置排好。标号2描绘了第二次做All-to-All（All-to-All Combine）的过程，这个过程的目的是将MoE算完的token再返回给各卡，原理和上述一致。

在 All-to-All Dispatch 操作之前准备好 All-to-All 输入的过程叫输入编码，即需要对本GPU上的 local token 按照路由结果进行 permute/group，将发往同一个专家的 token 进行分组。随后，这些 token 通过 All-to-All Dispatch 通信发送到对应的 expert rank（每个 EP rank 上只包含一部分 expert）。

在 All-to-All Combine 操作之后需要解包 All-to-All 的输出，组织为原始的顺序，这叫输出解码。大多数流行的DL框架利用NCCL的点对点（P2P）API来实现线性 All-to-All 算法，参见下图左侧。下图右侧是用python实现的输入编码。

5.4 分布式计算过程

5.4.1 多种范式结合

有了All-to-All的支撑，我们看看TP、DP和EP互相结合的MoE分布式计算过程，其中并行划分为：Transformer layer在四张GPU卡上并行，TensorParallel=2，DataParallel=2，ExpertParallel=2。张量并行组（GPU0和GPU1）在处理token 1和token 2，而张量并行组（GPU 2和GPU 3）处理token 3和token 4。整体计算过程分为7步。

• 步骤1：每个GPU首先计算属于它们自己的的自注意力块的分区。{GPU0,GPU1}与{GPU2,GPU3}之间是数据并行。

• 步骤2：在每个自注意力块的张量并行组内，每个GPU都会执行All-Reduce操作（对张量并行的部分和进行规约）来聚合它们各自token的完整输出激活（a1、a2、a3和a4）。这一步是聚合自注意力块的张量并行输出。

• 步骤3：每个GPU对于自己的本地token执行MoE路由功能。

• 步骤4：根据路由结果，将token发送到对应的专家。我们假设路由函数将令牌1和3映射到专家1，且将令牌2和4映射到专家2。于是执行如下操作。在专家并行组中执行一个All-to-All通信操作，根据路由函数决定的映射来路由令牌。让我们看看GPU 0和2组成的专家并行组。在GPU 0上，令牌1已映射到专家1，令牌2已映射到专家2。因此，我们希望GPU 0保留a1，并将a2发送到容纳专家2的GPU2。类似地，在GPU 2上，我们希望保留a4并将a3发送到GPU 0。

• 步骤5：专家计算。在All-to-All通信操作完成之后，每个专家分块在自己所在的GPU进行计算。

• 步骤6：在FFN的张量并行组内执行All-reduce操作来聚合完整输出。这一步是聚合FFN的张量并行输出。

• 步骤7：All-to-All通信，回到数据并行；执行All-to-All通信操作（实际是第一个All-to-All的逆操作），将令牌带回它们的原始GPU。

5.4.2 通信复杂度

TP vs EP

如上图左边所示，TP的基本思路是将模型参数切分到多个 GPU 进行计算。面对参数量大幅增大但计算量不变的MoE架构，TP 方案暴露出两大核心问题：通信会成为瓶颈，内存也会逐渐成为瓶颈，

通讯会成为瓶颈

随着TP size的增大，通讯会逐渐成为瓶颈。假设每次推理一个batch里一共有 S 个token，hidden dimension是 D，那么对于TP每一个MoE层每个GPU需要发送 $2SD$ 大小的数据，通讯量并不会随着TP size的增大而降低。

TP的部署方式使得每个GPU上都需要AllReduce来聚合所有input tokens的activation，无论是Self-Attention还是MLP都是对Reduce维度的划分，而reduce维度的划分，无论划分多少份，都是不会改变结果矩阵的大小的。那么TP中的AllReduce的通信量会随着整个部署实例batch size的增大而增大，并且即使增大TP并行度，其通信量也不会变小。这是因为TP技术是源自稠密模型的设计范式。

内存会成为瓶颈

由于TP划分的是权重，对同一个实例里的每一张卡都有相同的输入，所以TP的整体通信量是和当前实例的batch-size成正比的，且不受TP划分粒度的增加而减少。整个实例增大batchsize，那么对于实例内的每一张卡都需要增加batchsize。这极大限制了推理的batch size。

有限的batch size使得每次推理时，每个专家分到的token数量极其有限，使得专家的计算从compute-bound变为了memory-bound，大大降低了GPU的利用率。同时由于专家负载的不均衡性，较小的batch size甚至可能会导致单次推理时只有部分专家被激活。然而，TP 方案要求所有 GPU 加载全部专家的参数，即使某些专家未参与计算，也会占用显存资源。

EP的作用

TP很难增大推理的batch size，使得专家部分计算遇到内存瓶颈（memory-bound），TP也几乎不可能进行大规模模型并行扩展，来增加单个部署实例的batchsize。因此，面对超大参数量 MoE 模型以及大规模推理场景，传统 TP 方案有着巨大的本质上的局限性。

EP 方案为大规模 MoE 推理提供了一种全新的并行思路，能够有效解决 TP 方案的两大核心问题。

在通信开销方面，EP 采用 All-to-all 原语进行数据交换。在EP size增大的情况下，EP能大幅降低计算相同数量token的情况下单个GPU的通讯开销。同样以一个batch一共包含 𝑆 个token为例，假设每个token需要选择top-𝑘 的专家，以及专家之间负载平衡的话，那么每个GPU在token分发（dispatch）和重组（combine）两个阶段各需要发送$\frac{K\cdot S}{M}\cdot D$大小的数据，其中M是EP size。考虑dispatch和combine两阶段的通讯，当$\frac{K}{M}\ll 1$ 时，EP的通讯开销会远低于TP。

EP同时使得每个GPU可以计算不同的input data，而不需要像TP一样在每个GPU上处理相同的token并聚合 activation，EP可以极大的扩展batch size，使得每个专家都能分到足够数量的token，解决memory access的bottleneck。

通讯复杂度对比

下图给出了各个并行范式的通讯复杂度对比。

在专家并行（expert parallelism）中，每个MoE层在前向和反向传播阶段中，一共需要进行四次 All-to-All 通信，这会产生显著的开销，甚至成为效率的主要制约因素。这种通信的效率取决于多个因素，包括通道带宽的异质性、网络拓扑结构和集体通信算法。此外，MoE固有的负载不均衡可能通过引发同步延迟来加剧这些低效。为了优化节点内高带宽和节点间低带宽的使用，研究人员做了很多努力，比如：

• 最小化网络流量并利用高带宽连接。比如引入分层 All-to-All 、拓扑感知的路由策略、利用专家亲和性来进行分配等。
• 考虑到通信和计算的并发性，把流水线并行和专家并行集成，以此协调 All-to-All 通信和专家计算的重叠。也有研究人员利用GPU的大规模并行性和GPU发起的通信，将计算与依赖的集合通信进行融合。或者将通信依赖关系进行解耦来通信与计算之间的重叠。

5.4.3 代码示例

不同框架对MoE并行的相关概念不同，这恐怕是从事人工智能或者机器学习工作的人最苦恼的地方：模糊且有争议的定义。所以我们选取两个框架来看看。

DeepSpeed-Megatron

此处参考了 图解大模型训练系列之：DeepSpeed-Megatron MoE并行训练（源码解读篇） 猛猿 的精彩文章。

首先，我们给出一些背景知识。world size代表将要参与训练的进程数（或者计算设备数），每个进程都会被分配一个rank，该rank是一个介于0和world size-1之间的数字，该数字在作业中是唯一的。它作为进程标识符，并用于代替地址，将张量发送到指定的rank（进程）。

DeepSeed-Megatron中，假设每个MoE层有若干个专家（统称其为一套专家），现在我们想把这一套专家分布排列到若干GPU上。我们可以先定好要用几块GPU装下一套专家（EP），进而我们就能确认全局上共有多少套专家副本在跑（DP）。假设一共8张GPU，则：

• ep_world_size = 4：表示我们希望用4块GPU装下一套完整的专家。ep_group = 8 / ep_world_size = 8 /4 = 2，即一共2个专家组。我们需要在每个专家组内做All-to-All通信，将token发送去对应的专家。
• ep_dp_world_size = 2：MoE层的数据并行的大小。例如下图中[g0, g8]上都维护着e0，所以它们构成一个ep_dp_group。这个group的作用是当我们在计算bwd时，它们之间是需要做梯度的allreduce通讯的，构成ep_dp_group的条件不仅是e相同，还需要每个e接受到的的batch数据不同。

即，在FWD中，ep_group进行all2all通讯，将token发去对应的专家做计算，并将计算结果取回。
在BWD中，ep_dp_group进行AllReduce通讯梯度，用于更新对应的专家的参数。

我们来看看源码中的示例函数。

def _get_expert_parallel_ranks(world_size,tensor_parallel_size_,expert_parallel_size_,pipeline_parallel_size_=1,use_data_before_expert_parallel_=False):"""Generate expert parallel and expert data parallel group ranks list.Example - E + M + D parallelworld_size = 16model_degree = 2expert_degree = 4 # number of experts in same groupmp_group = [0, 1], [2,3], [4,5] ...data_parallel_group =[0,2,4,6,8,10, 12,14],                 [1,3,5,7,9,11,13,15]expert_parallel_group = [0,2,4,6], [8,10,12,14]             [1,3,5,7], [9,11,13,15]expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14],    [1,9],[3,11],[5,13],[7,15]Args:world_size (int): Distributed world size.tensor_parallel_size_ (int): Tensor parallel group size.expert_parallel_size_ (int): Expert parallel group size.pipeline_parallel_size_ (int): Pipeline parallel group sizeuse_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topologyReturns:Expert parallel group ranks and Expert data parallel group ranks list."""_ensure_divisibility(world_size, tensor_parallel_size_ * pipeline_parallel_size_)# dp_world_size = world_size // (tensor_parallel_size_ * pipeline_parallel_size_)_ensure_divisibility(dp_world_size, expert_parallel_size_)# Generate data parallel groupsdata_parallel_groups = []dp_group_size = tensor_parallel_size_pp_stride = world_size // pipeline_parallel_size_if use_data_before_expert_parallel_:dp_stride = world_size // expert_parallel_size_ // tensor_parallel_size_ // pipeline_parallel_size_for pp_stage_start in range(0, world_size, pp_stride):pp_stage_next = pp_stage_start + pp_stridefor i in range(dp_group_size):data_parallel_groups.append(list())for ds in range(dp_stride):# [0, 4, 8, 12, 16, 20, 24, 28, 2, 6, 10, 14, 18, 22, 26, 30]# [1, 5, 9, 13, 17, 21, 25, 29, 3, 7, 11, 15, 19, 23, 27, 31]data_parallel_groups[-1].extend(list(range(pp_stage_start + i + ds * tensor_parallel_size_, pp_stage_next,dp_stride * tensor_parallel_size_)))else:for pp_stage_start in range(0, world_size, pp_stride):pp_stage_next = pp_stage_start + pp_stridefor i in range(dp_group_size):data_parallel_groups.append(list(range(pp_stage_start + i, pp_stage_next, dp_group_size)))expert_parallel_groups = []expert_data_parallel_groups = []for dp_ranks in data_parallel_groups:# partition of expert parallel groups, e.g. [0,2,4,6], [8,10,12,14]part_ep_groups = []for i in range(0, dp_world_size, expert_parallel_size_):part_ep_groups.append(dp_ranks[i:i + expert_parallel_size_])expert_parallel_groups.extend(part_ep_groups)# zip part_ep_groups get expert data parallel ranks, e.g [0,8],[2,10],[4,12],[6,14]for expert_dp_ranks in zip(*part_ep_groups):expert_data_parallel_groups.append(list(expert_dp_ranks))return expert_parallel_groups, expert_data_parallel_groups

下图展示了一个MoE层的整体架构。

• 首先，我们定义好了单个expert模型架构（ParallelMLP）
• 然后，鉴于一张卡上可能不止维护1个expert（num_local_experts = num_experts // ep_world_size），我们需要定义这张卡上expert的集合Experts（nn.ModuleList，见代码细节）
• 再次，我们需要一个TopKGate策略，来帮助token选择expert。
• 最后，将以上内容组装成一个MOELayer。

下面是ParallelTransformerLayer的定义。

class ParallelTransformerLayer(MegatronModule):"""A single transformer layer.Transformer layer takes input with size [s, b, h] and returns anoutput 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):# retriever=None):args = get_args()super(ParallelTransformerLayer, self).__init__()self.layer_number = layer_numberself.layer_type = layer_type# MLPself.num_experts = num_expertsif args.num_experts_switch is not None:self.mlp = SwitchMLP(config) # Megatron-LM's MoEelse:if self.num_experts <= 1: # dense, not MoEself.mlp = ParallelMLP(config)else: # DeepSpeed's MoEenable_expert_tensor_parallelism = args.enable_expert_tensor_parallelismself.mlp = MoE(args.hidden_size,# 定义单个专家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_sizek=args.topk,use_residual=(args.mlp_type == 'residual'),capacity_factor=args.moe_train_capacity_factor,eval_capacity_factor=args.moe_eval_capacity_factor,min_capacity=args.moe_min_capacity,drop_tokens=args.moe_token_dropping, # 是否需要做溢出处理use_tutel=args.use_tutel,enable_expert_tensor_parallelism=enable_expert_tensor_parallelism,top2_2nd_expert_sampling=args.moe_top2_2nd_expert_sampling)

MoE代码位于from deepspeed.moe.layer import MoE，具体如下。

class MoE(nn.Module):"""Initialize an MoE layer.Arguments:hidden_size (int): token embedding.expert (nn.Module): 专家 (e.g., MLP, torch.linear)，此处使用ParallMLPnum_experts (int, optional): 每层专家数ep_size (int, optional): default=1, 专家并行中的rank数，或者说ep_world_size，即用ep_size张卡容纳全部专家k (int, optional): default=1, top-k gating value, only supports k=1 or k=2.capacity_factor (float, optional): default=1.0, 训练时的容量因子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.cnoisy_gate_policy (str, optional): default=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, # 是否对专家进行TP切分top2_2nd_expert_sampling (bool, optional): default=True, whether to perform sampling for 2nd expert"""def __init__(self,hidden_size: int,expert: nn.Module,num_experts: int = 1,ep_size: int = 1,k: int = 1,capacity_factor: float = 1.0,eval_capacity_factor: float = 1.0,min_capacity: int = 4,use_residual: bool = False,noisy_gate_policy: Optional[str] = None,drop_tokens: bool = True,use_rts: bool = True,use_tutel: bool = False,enable_expert_tensor_parallelism: bool = False,top2_2nd_expert_sampling: bool = True) -> None:super(MoE, self).__init__()self.use_residual = use_residualself.enable_expert_tensor_parallelism = enable_expert_tensor_parallelismself.ep_size = ep_sizeself.expert_group_name = f"ep_size_{self.ep_size}"self.num_experts = num_expertsself.num_local_experts = num_experts // self.ep_size # 单块GPU上需存放的专家数量# 定义一个MoE层上所有的专家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, None,top2_2nd_expert_sampling), 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 mlpself.coefficient = nn.Linear(hidden_size, 2)def set_deepspeed_parallelism(self, use_data_before_expert_parallel_: bool = False) -> None:# 专家分布相关设置self._create_process_groups(use_data_before_expert_parallel_=use_data_before_expert_parallel_)def _create_process_groups(self, use_data_before_expert_parallel_: bool = False) -> None:# 专家分布相关设# Create process group for a layer if neededif self.expert_group_name not in groups._get_expert_parallel_group_dict():# 按EP + DP方式设置专家并行相关组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 purposegroups._create_expert_and_data_parallel(self.ep_size, use_data_before_expert_parallel_=use_data_before_expert_parallel_)else:# 使用EP + DP + TP方式设置专家并行相关组# expert tensor parallelism is enabledgroups._create_expert_data_and_model_parallel(self.ep_size, mpu=groups.mpu, use_data_before_expert_parallel_=use_data_before_expert_parallel_)# Set the group handle for the MOELayer (deepspeed_moe) object# 为当前进程所属的MoE层设置ep_group，样就可以在ep_group内做All-to-All通讯,如果不设置ep_group，默认对所有GPU卡（ep_world_size）做All-to-All通信self.deepspeed_moe._set_ep_group(groups._get_expert_parallel_group(self.expert_group_name))def forward(self,hidden_states: torch.Tensor,used_token: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:""" MoE forwardArguments:hidden_states (Tensor): input to the layerused_token (Tensor, optional): default: None, mask only used tokensReturns:A tuple including output, gate loss, and expert count.* output (Tensor): output of the model* l_aux (Tensor): gate loss value* exp_counts (Tensor): expert count"""output = self.deepspeed_moe(hidden_states, used_token)if self.use_residual:# Residual MoEoutput_mlp = self.mlp(hidden_states)if isinstance(output_mlp, tuple):output_mlp = output_mlp[0]  # Ignore the bias term for nowcoef = self.coefficient(hidden_states)coef = F.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           

deepspeed在MOELayer的实现如下。

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.pdfArgs:gate (torch.nn.Module):gate networkexpert (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__() self.gate = gate # TopKGate类，用来决定token的分发策略self.experts = experts # 当前进程所属的GPU上维护的所有专家，nn.ModuleList[ParallelMLP()]self.ep_group = None # 当前进程所属的ep_group，为None时表示所有GPU构成一个ep_groupself.ep_size = ep_size # 当前进程所属的ep_group的ep_world_size，即GPU卡数self.ep_group_name = ep_group_name # 当前进程所属的ep_group的名字self.num_local_experts = num_local_experts # 当前进程所属的GPU上所维护的专家数量，即为self.experts中维护的专家s数量self.use_tutel = use_tutel and TUTEL_INSTALLED and gate.k == 1 # 是否使用tutel做路由优化def _set_ep_group(self, ep_group):self.ep_group = ep_groupself.gate._set_ep_group(ep_group)def forward(self, *input: Tensor, **kwargs: Any) -> Tensor:"""*号：传入的input是一个tuple，一般是一个二元组input[0]是做计算的batch数据，其形状为(seq_len, batch_size, embedding_size)，embedding_size就是d_modelinput[1]是掩码数据，其尺寸为(seq_len * batch_size)，可以在计算MoE结果时，对某些token做掩码，使其不参与计算"""# Implement Algorithm 2 from GShard paper.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尺寸为(seq_len * batch_size, token_embedding_size)reshaped_input = input[0].reshape(-1, d_model)if self.use_tutel: # 使用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)else:# 使用自定义的Gshard gate来确定token的分发策略# gate：TopKGate类，l_aux: 辅助损失函数值# combine_weights: 尺寸为(seq_len * batch_size, expert_num, capacity)，表示对每个token（总共seq_len * capacity个）而言，它对每个专家（总共expert_num个）的weight，这个weight按照该token在专家buffer中的位置（总共capacity个token）存放，不是目标位置的地方则用0填充# dispatch_mask：相当于combine_weights.bool()，combine_weights为0的地方设为False，为1的地方设为True。dispatch_mask后续将被用在zero paddingself.l_aux, combine_weights, dispatch_mask, self.exp_counts = self.gate(reshaped_input, input[1])# 将输入数据按照专家的顺序排好，并做zero padding,# dispatched_input: 尺寸为(expert_num, capacity, token_embedding_size)，表示每个专家的buffer下要处理的token_embedding           dispatched_input = einsum("sec,sm->ecm", dispatch_mask.type_as(input[0]), reshaped_input)tensor_model_world_size = bwc_tensor_model_parallel_world_size(groups.mpu)# 当expert不采用tp切分，而non-MoE部分采用tp切分时，为避免数据重复发送，需要对同一个tp组内的tokens做去重if tensor_model_world_size > 1:# If the non-expert is tensor-parallel,# Whether expert is tensor-parallel or not , it will create# duplicate tokens on the tensor-parallel ranks.# drop duplicate tokens also doubles up as a communication# optimization as we are reducing the all-to-all communication volume.# 1: for not tensor-parallel expert,drop duplicate tokens to ensure# both correctness and reduce all-to-all communication.# 2: for tensor-parallel expert,drop duplicate tokens to reduce all-to-all# communication volume,before expert execution, it is necessary to perform# an allgather to ensure correctness,dispatched_input = drop_tokens(dispatched_input, dim=1)# 第一次All2All：将token发给对应的expert，dispatched_input尺寸为(expert_num, capacity, token_embedding_size)，又可以写成(ep_world_size * num_local_experts, capacity, token_embedding_size)。dispatched_input = _AllToAll.apply(self.ep_group, dispatched_input)if tensor_model_world_size > 1 and groups._get_expert_model_parallel_world_size() > 1:# if both expert and non-expert are tensor-parallel# the dropped duplicate tokens need to be gathered on each# tensor parallel rank again to ensure correctnessdispatched_input = gather_tokens(dispatched_input, dim=1)# Re-shape after all-to-all: ecm -> gecm（g是ep_world_size，e是num_local_experts）# 在将dispatched_input正式喂给expert前，把它reshape成(ep_world_size, num_local_experts, capacity, token_embedding_size)dispatched_input = dispatched_input.reshape(self.ep_size, self.num_local_experts, -1, d_model)# 将token喂给expert计算，expert_output尺寸为(ep_world_size, num_local_experts, capacity, token_embedding_size)expert_output = self.experts(dispatched_input)# Re-shape before drop_tokens: gecm -> ecm# expert_output的形状是(ep_world_size * num_local_experts, capacity, token_embedding_size)。expert_output = expert_output.reshape(self.ep_size * self.num_local_experts, -1, d_model)if tensor_model_world_size > 1 and groups._get_expert_model_parallel_world_size() > 1:# if both expert and non-expert are tensor-parallel# drop duplicate tokens to ensure both correctness# and reduce all-to-all communication.expert_output = drop_tokens(expert_output, dim=1)# 第二次All2All，将算好的token返回给产出它的GPU， expert_output为(ep_world_size * num_local_experts, C, M)，即此时这张卡上维护的token过MoE的结果，是由它从ep_group（ep_world_size）内所有expert(num_local_experts)的结果汇总而来expert_output = _AllToAll.apply(self.ep_group, expert_output)# 如果之前在tp组内做过数据去重处理，这里要把数据all-gather回来if tensor_model_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)# 使用combine_weights进行加权计算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, token_embedding_size)a = combined_output.reshape(input[0].shape)return a

门控函数定义如下。

class TopKGate(Module):"""Gate module which implements Top2Gating as described in Gshard_.::gate = TopKGate(model_dim, num_experts)l_aux, combine_weights, dispatch_mask = gate(input).. Gshard_: https://arxiv.org/pdf/2006.16668.pdfArgs:model_dim (int):size of model embedding dimensionnum_experts (int):number of experts in model"""wg: torch.nn.Lineardef __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,ep_group: Union[torch.distributed.ProcessGroup, None] = None,top2_2nd_expert_sampling: bool = True) -> None:super().__init__()self.wg = torch.nn.Linear(model_dim, num_experts, bias=False)self.ep_group = ep_groupself.k = kself.capacity_factor = capacity_factorself.eval_capacity_factor = eval_capacity_factorself.min_capacity = min_capacityself.noisy_gate_policy = noisy_gate_policyself.timers = SynchronizedWallClockTimer()self.wall_clock_breakdown = Falseself.gate_time = 0.0self.drop_tokens = drop_tokensself.use_rts = use_rtsself.top2_2nd_expert_sampling = top2_2nd_expert_samplingdef _set_ep_group(self, ep_group):self.ep_group = ep_groupdef forward(self,input: torch.Tensor,used_token: torch.Tensor = None,use_tutel: bool = False) -> Tuple[Tensor, Tensor, Tensor]:  input_fp32 = input.float()# input jitteringif self.noisy_gate_policy == 'Jitter' and self.training:input_fp32 = multiplicative_jitter(input_fp32, device=input.device)logits = torch.nn.functional.linear(input_fp32, weight=self.wg.weight.float(), bias=None)if self.k == 1:gate_output = top1gating(logits, self.capacity_factor if self.training else self.eval_capacity_factor, self.min_capacity, used_token, self.noisy_gate_policy if self.training else None, self.drop_tokens, self.use_rts, self.ep_group, use_tutel)elif self.k == 2:gate_output = top2gating(logits, self.capacity_factor if self.training else self.eval_capacity_factor, self.min_capacity, self.drop_tokens, self.ep_group, self.top2_2nd_expert_sampling)else:gate_output = topkgating(logits, self.k, self.capacity_factor if self.training else self.eval_capacity_factor, self.min_capacity, self.drop_tokens, self.ep_group)return gate_output

Experts的定义如下，其定义一个MoE层上所有的Expert。

class Experts(nn.Module):def __init__(self, expert: nn.Module, num_local_experts: int = 1, expert_group_name: Optional[str] = None) -> None:super(Experts, self).__init__()self.deepspeed_experts = nn.ModuleList([copy.deepcopy(expert) for _ in range(num_local_experts)])self.num_local_experts = num_local_experts # 每块GPU上共num_local_experts个expertfor expert in self.deepspeed_experts:           for param in expert.parameters():param.allreduce = Falseparam.group_name = expert_group_namedef forward(self, inputs: torch.Tensor) -> torch.Tensor:"""inputs尺寸：(ep_world_size, num_local_experts, capacity, token_embedding_size)在分发去experts前，每张卡上的输出结果为(ep_world_size * num_local_experts, capacity, token_embedding_size)对于All2All通讯可以理解为，对于ep_group内的每张卡，都将数据沿着ep_world_size * num_local_experts维度切成ep_world_size块后，再进行通讯。目的是保证每张卡上的数据块数量 = ep_world_size，这样All2All通讯才不会出错，因此发送完毕后，每张卡上的数据可以又表示为(ep_world_size * num_local_experts, capacity, token_embedding_size)进一步在正式把数据喂给这张卡上维护的experts前，我们可以把数据reshape成(ep_world_size, num_local_experts, capacity, token_embedding_size)的形式。即沿着num_local_experts维度将数据切分为num_local_experts个chunck，则一个chunk对应一个local_expert，再次实现了token 和local expert间一一对应的关系        """# chunks: 沿着num_local_expert维度切分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尺寸：(ep_world_size, capacity, token_embedding_size)out = expert(chunk)if isinstance(out, tuple):out = out[0]  # Ignore the bias term for nowexpert_outputs += [out]# concat后最终out尺寸: (ep_world_size, num_local_experts, capacity, token_embedding_size)return torch.cat(expert_outputs, dim=1)

专家定义如下。

class ParallelMLP(MegatronModule):"""MLP.MLP will take the input with h hidden state, project it to 4*hhidden dimension, perform nonlinear transformation, and project thestate back into h hidden dimension."""def __init__(self, config, moe=False, enable_expert_tensor_parallelism=False):super(ParallelMLP, self).__init__()args = get_args()self.add_bias = config.add_bias_linearffn_hidden_size = config.ffn_hidden_sizeif config.gated_linear_unit:ffn_hidden_size *= 2# Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf# 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.bias_gelu_fusion = Falseself.activation_func = Noneself.swiglu = args.swigluif args.openai_gelu:self.activation_func = openai_geluelif args.onnx_safe:self.activation_func = erf_geluelif args.swiglu:def swiglu(x):x = torch.chunk(x, 2, dim=-1)return F.silu(x[0]) * x[1]self.activation_func = swigluelif args.squared_relu:def squared_relu(x):return torch.pow(F.relu(x), 2)self.activation_func = squared_reluelse:self.bias_gelu_fusion = args.bias_gelu_fusionself.activation_func = F.gelu# Project back to h.# 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):# [s, b, 4hp]# 输入数据过Wi层，如果做TP切分，则尺寸为[s, b, 4h/tp_word_size]intermediate_parallel, bias_parallel = self.dense_h_to_4h(hidden_states)if self.bias_gelu_fusion:# DeepSpeed FLOPS profiler temporarily substitues functions like F.gelu to calculate the throughputintermediate_parallel = bias_gelu_impl(intermediate_parallel, bias_parallel)else:if bias_parallel is not None:intermediate_parallel = intermediate_parallel + bias_parallelintermediate_parallel = self.activation_func(intermediate_parallel)# [s, b, h]# Wi层输出数据过Wo层，如果对expert采取tp切分，这里的输出需要在tp_group内做AllReduceoutput, output_bias = self.dense_4h_to_h(intermediate_parallel)return output, output_bias

FastMoE

我们使用 FastMoE 作为示例，看看如何实现分布式操作。 这里的代码比较显式。

论文”FASTMOE: A FAST MIXTURE-OF-EXPERT TRAINING SYSTEM“提出了 FastMoE，这是一个基于 Pytorch 的分布式 MoE 训练系统。该系统支持将不同的专家放置在多个节点上的多个 GPU 中，从而实现专家数量和 GPU 数量线性增加。

FastMoE 支持将专家分布在多个节点的多个 Worker 上，并且将不同 Worker 之间的数据通信隐藏起来，模型开发人员不用考虑。此外，在分布式 MoE 系统中的一个主要挑战为：动态路由导致分配给不同专家的输入样本数可能存在很大的差异。作者的方案为：在 Worker 之间交换实际的数据之前，先在 Worker 之间交换大小信息，Worker 根据相应信息分配 Buffer，然后传输真实的数据。

FastMoE 将所有输入样本一起 Batching 后发给同一个专家。由于数据表示的限制，FastMoE 使用专门开发的 CUDA Kernel 进行内存移动，以减少开销。如下图所示，给定每个样本要进入的索引（Gating 输出），通过 Scatter 操作将所有样本按照对应顺序进行排布，执行完专家计算之后，再按照相反的 Gather 操作进行复原。

MoE代码如下，关键函数是_fmoe_general_global_forward()函数，该函数会完成MoE的关键计算步骤。

class FMoE(nn.Module):r"""A general moe implementation that supports an arbitrary module as theexpert.* `num_expert` stands for the number of experts on **each** worker.* `world_size` stands for the total number of workers that containsdifferent experts.* `slice_group` can be a torch's communication group, indicating thatspecific model parallel is applied across the group, and workers in thegroup hold the same copy of input feature, and requires the same copy ofthe output. For each worker, FMoE only computes the output of a certainslice of the input batch, and will all-gather the outputs aftercomputation.* `mp_group` is a deprecated alias of `slice_group`* `moe_group` stands for the group of process that performs expertparallelism. The default value `None` means all processes. See theparallelism document for more details of the groups.* `top_k` stands for the number of experts each token is going to.* `gate` is a gate class which can found in `fmoe.gates`.* `expert` can be specified as a module class, it is used to generate`num_expert` expert modules.* `gate_bias` is only valid for naive_gate and its subclasses, it meanswhether to add bias to the gate module."""def __init__(self,num_expert=32,d_model=1024,world_size=1,mp_group=None,  # being deprecatedslice_group=None,moe_group=None,top_k=2,gate=NaiveGate,expert=None,gate_hook=None,mask=None,mask_dict=None,gate_bias=True,):super().__init__()self.num_expert = num_expertself.d_model = d_modelself.world_size = world_sizeself.slice_group = slice_groupif mp_group is not None:self.slice_group = mp_groupif self.slice_group is None:self.slice_size = 1self.slice_rank = 0else:self.slice_size = self.slice_group.size()self.slice_rank = self.slice_group.rank()self.top_k = top_kif type(expert) is list:self.experts = nn.ModuleList([e(d_model) for e in expert])self.experts_fused = Falseself.num_expert = num_expert = len(expert)elif expert is not None:self.experts = nn.ModuleList([expert(d_model) for _ in range(num_expert)])self.experts_fused = Falseelse:self.experts_fused = Trueif issubclass(gate, NaiveGate):self.gate = gate(d_model, num_expert, world_size, top_k, gate_bias=gate_bias)else:self.gate = gate(d_model, num_expert, world_size, top_k)self.gate_hook = gate_hookself.mask = maskself.mask_dict = mask_dictself.moe_group = moe_groupdef expert_fn(self, inp, fwd_expert_count):r"""The default expert function which either calls the experts as a wholeor as separate experts."""if self.experts_fused:return self.experts(inp, fwd_expert_count)if isinstance(fwd_expert_count, torch.Tensor):fwd_expert_count_cpu = fwd_expert_count.cpu().numpy()outputs = []base_idx = 0for i in range(self.num_expert):batch_size = fwd_expert_count_cpu[i]inp_slice = inp[base_idx : base_idx + batch_size]outputs.append(self.experts[i](inp_slice, torch.tensor([fwd_expert_count[i]])))base_idx += batch_sizereturn torch.cat(outputs, dim=0)def expert_fn_single(self, inp, fwd_expert_count, idx):r"""forward single expert for smart scheduling."""output = self.experts[idx](inp, fwd_expert_count)return outputdef mark_parallel_comm(self, expert_dp_comm="none"):r"""Automatically mark the data parallel comms of the parameters within themodule. This can be typically called at the end of the __init__ functionin child classes."""if self.experts is not None:comm = expert_dp_commif isinstance(self.experts, list):for e in self.experts:mark_module_parallel_comm(e, comm)else:mark_module_parallel_comm(self.experts, comm)mark_module_parallel_comm(self.gate, "gate")def forward(self, moe_inp):r"""The FMoE module first computes gate output, and then conduct MoE forwardaccording to the gate.  The score of the selected gate given by theexpert is multiplied to the experts' output tensors as a weight."""moe_inp_batch_size = tree.flatten(tree.map_structure(lambda tensor: tensor.shape[0], moe_inp))if self.world_size > 1:def ensure_comm_func(tensor):ensure_comm(tensor, self.moe_group)tree.map_structure(ensure_comm_func, moe_inp)if self.slice_size > 1:def slice_func(tensor):return Slice.apply(tensor, self.slice_rank, self.slice_size, self.slice_group)moe_inp = tree.map_structure(slice_func, moe_inp)gate_top_k_idx, gate_score = self.gate(moe_inp)if self.gate_hook is not None:self.gate_hook(gate_top_k_idx, gate_score, None)# delete masked tensorsif self.mask is not None and self.mask_dict is not None:# TODO: to fixdef delete_mask_func(tensor):# to: (BxL') x d_modeltensor = tensor[mask == 0, :]return tensormask = self.mask.view(-1)moe_inp = tree.map_structure(delete_mask_func, moe_inp)gate_top_k_idx = gate_top_k_idx[mask == 0, :]fwd = _fmoe_general_global_forward(moe_inp, gate_top_k_idx, self.expert_fn_single if fmoe_faster_schedule else self.expert_fn,self.num_expert, self.world_size,experts=self.experts)# recover deleted tensorsif self.mask is not None and self.mask_dict is not None:def recover_func(tensor):# to: (BxL') x top_k x dimdim = tensor.shape[-1]tensor = tensor.view(-1, self.top_k, dim)# to: (BxL) x top_k x d_modelx = torch.zeros(mask.shape[0],self.top_k,dim,device=tensor.device,dtype=tensor.dtype,)# recoverx[mask == 0] = tensorfor k, v in self.mask_dict.items():x[mask == k] = vreturn xmoe_outp = tree.map_structure(recover_func, fwd)else:def view_func(tensor):dim = tensor.shape[-1]tensor = tensor.view(-1, self.top_k, dim)return tensormoe_outp = tree.map_structure(view_func, fwd)gate_score = gate_score.view(-1, 1, self.top_k)def bmm_func(tensor):dim = tensor.shape[-1]tensor = torch.bmm(gate_score, tensor).reshape(-1, dim)return tensormoe_outp = tree.map_structure(bmm_func, moe_outp)if self.slice_size > 1:def all_gather_func(tensor):return AllGather.apply(tensor, self.slice_rank, self.slice_size, self.slice_group)moe_outp = tree.map_structure(all_gather_func, moe_outp)moe_outp_batch_size = tree.flatten(tree.map_structure(lambda tensor: tensor.shape[0], moe_outp))return moe_outp

_fmoe_general_global_forward()函数的代码如下。

def prepare_forward(gate, num_expert, world_size):r"""Prepare necessary information from gate output for MoE computation.Args:gate: a 1-d Long Tensor representing the target expert of each inputsample.num_expert: number of experts on each worker.world_size: number of workers that hold different experts.comm: the communicator of all workers in the expert-parallel group."""pos, local_expert_count, global_expert_count = count_by_gate(gate, num_expert, world_size)with torch.no_grad():fwd_expert_count = global_expert_count.view(world_size,num_expert).sum(dim=0)fwd_batch_size = int(fwd_expert_count.sum().item())return (pos,local_expert_count.cpu(),global_expert_count.cpu(),fwd_expert_count.cpu(),fwd_batch_size,)def _fmoe_general_global_forward(inp, gate, expert_fn, num_expert, world_size, **kwargs):r"""A private function that performs the following steps to complete the MoEcomputation.* Count the number of tokens from each worker to each expert.* Send the features to their target position so that input features to eachexpert are contiguous in memory.* Perform the forward computation of the experts using `expert_fn`* Gather the output features of experts back, and reorder them as sentences.Intermediate results like expert counts are hidden from users by thisfunction."""(pos,local_expert_count,global_expert_count,fwd_expert_count,fwd_batch_size,) = prepare_forward(gate, num_expert, world_size) # 获得专家index信息topk = 1if len(gate.shape) == 2:topk = gate.shape[1]def scatter_func(tensor): #  All-to-All dispatchreturn MOEScatter.apply(tensor,torch.div(pos, topk, rounding_mode='floor'),local_expert_count,global_expert_count,fwd_batch_size,world_size,)x = tree.map_structure(scatter_func, inp)x = expert_fn(x, fwd_expert_count) # 专家处理out_batch_size = tree.flatten(inp)[0].shape[0]if len(gate.shape) == 2:out_batch_size *= gate.shape[1]def gather_func(tensor): # All-to-All combine，返回给对应的rankreturn MOEGather.apply(tensor,pos,local_expert_count,global_expert_count,out_batch_size,world_size,)outp = tree.map_structure(gather_func, x)return outp

训练时的步进函数，可以看到torch.distributed.all_reduce的使用。

def patch_forward_step(forward_step_func, Megatron_Version="v2.2"):r"""Patch model's forward_step_func to support balance loss"""from megatron.mpu import is_pipeline_last_stagefrom megatron.mpu import get_tensor_model_parallel_groupfrom megatron import get_argsif not get_args().balance_strategy:return forward_step_funcdef forward_step_with_balance_loss_v2_2(data_iterator, model, input_tensor):args = get_args()output = forward_step_func(data_iterator, model, input_tensor)if not is_pipeline_last_stage() or not args.balance_strategy:return outputwhile hasattr(model, 'module'):model = model.moduleloss_list = [l.mlp.gate.get_loss(clear=False).view(1)for l in model.language_model.transformer.layersif l.mlp.gate.has_loss]if len(loss_list) == 0:return outputloss_name = args.balance_strategy + "_loss"(loss, state_dict), bal_loss = (output,torch.cat(loss_list).mean() * args.balance_loss_weight)# avarage across moe groupmoe_group = get_tensor_model_parallel_group()world_size = torch.distributed.get_world_size(group=moe_group)averaged_bal_loss = bal_loss.clone().detach()torch.distributed.all_reduce(averaged_bal_loss, group=moe_group)averaged_bal_loss /= world_sizeloss += bal_lossstate_dict[loss_name] = averaged_bal_lossreturn loss, state_dict

0xEE 个人信息

★★★★★★关于生活和技术的思考★★★★★★

微信公众账号：罗西的思考

如果您想及时得到个人撰写文章的消息推送，或者想看看个人推荐的技术资料，敬请关注。

0xFF 参考

一个关于MoE的猜想 渣B zartbot

Mixtral of Experts

SwitchHead：使用专家混合模型注意力加速 Transformer

LLM MOE的进化之路，从普通简化 MOE，到 sparse moe，再到 deepseek 使用的 share_expert sparse moe chaofa用代码打点酱油

Mixture of Parrots: Experts improve memorization more than reasoning

Scaling Laws for Fine-Grained Mixture of Experts

https://arxiv.org/pdf/2402.07871

大规模分布式 AI 模型训练系列—专家并行 AI闲谈

DeepSeek-R1模型架构深度解读（三）弄懂DeepSeekMoE AI算法之道 [AI算法之道](javascript:void(0)😉

详细谈谈DeepSeek MoE相关的技术发展 渣B [zartbot](javascript:void(0)😉

DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale

混合专家（MoE）的路由（Deepspeed ） 杨歆

Megatron-LM中MOE并行分组策略 庭叶藏

SGLang的Expert Parallel特性解读 BBuf [GiantPandaCV](javascript:void(0)😉

图解大模型训练系列之：DeepSpeed-Megatron MoE并行训练（源码解读篇） 猛猿

图解大模型训练系列之：DeepSpeed-Megatron MoE并行训练（原理篇） 猛猿

简单理解DeepSpeed-MoE专家模型和all2all通讯 voodoo

重新思考 MoE 王庆法 [清熙](javascript:void(0)😉

Moe模型的对比：Mixtral, Qwen2-MoE, DeepSeek-v3 Alex [算法狗](javascript:void(0)😉

混合专家模型Mixtral-8x7B模型挖坑指北 孟繁续 [青稞AI](javascript:void(0)😉

A Uniffed View for Attention and MoE

统一视角看 Attention 与 MoE Taki

关于Deepseek采用EP推理方式的一些思考 杨鹏程

MOE介绍及其LLM方案整理 假如给我一只AI

首篇MoE工作-Adaptive mixtures of local experts(1991) uihcgniw

【IDPT论文解读】Adaptive Mixtures of Local Experts - 多系统融合 JaPay

Deepseek-MOE架构图解（V1->V2->V3） 假如给我一只AI

DeepEP Dispatch/Combine 图示 Marlene

MoE Inference On AnyScale MoE-On-AnyScale