混合专家模型MoE的全面指南（二）路由机制、负载均衡

最新推荐文章于 2025-04-06 21:24:53 发布

大模型面试

最新推荐文章于 2025-04-06 21:24:53 发布

阅读量2.6k

点赞数 30

文章标签：人工智能 AI大模型语言模型 ai LLM 大模型

本文链接：https://blog.csdn.net/Code1994/article/details/145209710

版权

三、路由机制

既然我们有了这么一帮专家，那模型咋知道该用哪个专家呢？在专家前面加了个门控网络，它会学着挑出给定词用哪个专家。

1、路由器

路由器（门控网络）也是前馈神经网络，根据特定输入来选专家。它会输出概率，用这些概率来挑出最匹配的专家：

专家层会返回选中的专家的输出乘以门控值（选择概率）。

路由器加上专家（只选几个）就组成了专家混合（MoE）层：

一个MoE层有两种尺寸，要么是稀疏的，要么是密集的专家混合。

两者都用路由器选专家，但稀疏MoE只挑几个，密集MoE虽然全选，但可能按不同分布来选。

比如，给定一串词，MoE会把词分给所有专家，而稀疏MoE就只挑几个专家。

在目前大型语言模型（LLM）的状况下，看到‘MoE’一般指稀疏MoE，因为它能用专家的一个子集。这样计算成本低，对大型语言模型来说很重要。

2、专家的选择

门控网络可以说是专家混合（MoE）里最重要的部分，因为它不仅决定了推理时用哪些专家，训练时也一样。

最基础的形式下，我们把输入（x）和路由器的权重矩阵（W）相乘：

然后，对输出应用 SoftMax 函数，为每个专家生成一个概率分布

路由器用这个概率分布来挑出给定输入最适合的专家。

最后，我们把每个路由器的输出和每个选中的专家的输出相乘，再把结果加起来。

咱们把所有步骤串起来，看看输入是怎么通过路由器和专家的：

3、路由的复杂性

不过，这个简单的函数常常会导致路由器总是选同一个专家，因为有些专家可能学得比其他专家快：

这不仅会导致专家选择的分布不均，还会使一些专家几乎不被训练，从而在训练和推理过程中引发问题。

为了解决这个问题，我们希望在训练和推理过程中，所有专家都能得到平等的重视，这被称为负载均衡。这样可以防止模型过度依赖少数几个专家，从而避免过拟合。

四、负载均衡

为了平衡专家的重要性，我们需要关注路由器，因为它是决定在特定时间选择哪些专家的主要组件。

1、KeepTopK

一种负载均衡路由器的方法是通过一个简单的扩展，称为 KeepTopK。通过引入可训练的高斯噪声，有助于打破对特定专家的偏好，使选择更加随机化，这样我们可以防止总是选择相同的专家：

然后，除了你想要激活的顶部 k 个专家（例如 2 个）之外，其他专家的权重会被设置为 −∞：

通过将这些权重设置为 −∞，SoftMax 函数在这些权重上的输出将导致概率为 0：

SoftMax 函数的性质： SoftMax 函数将一个向量转换为概率分布，每个元素的值在 0 和 1 之间，且所有元素的和为 1。 SoftMax 函数的公式为：

其中，是输入向量，是向量的长度，是向量的第个元素。设置权重为：

当某个元素的值为时，的值为 0。

因此，对于任何包含的向量，SoftMax 函数将这些元素的输出概率设置为 0。具体步骤

假设我们有中间向量，其中某些元素被设置为：

应用 SoftMax 函数：

具体计算：由于e^{-\infty} = 0，上述表达式简化为：\text{SoftMax}(\math

计算具体值：\text{SoftMax}(\mathbf{z}_{\text{top-k}}) = \left, \frac{1.6487}{1.6487 + $2.22$55}, 0, \frac{2.2255}{1.6487 + 2.2255}, 0 \right] = \left[ 0, 0.4256, 0,0.5744, 0 \right]

结果

通过将权重设置为，SoftMax 函数确保了这些专家的激活概率为 0，从而只激活顶部个专家。

优点

负载均衡： 确保每个专家都有机会被激活，避免某些专家过度使用而其他专家闲置。

防止过拟合： 通过引入噪声，防止模型过度依赖少数几个专家，从而提高模型的泛化能力。

灵活性： 可以根据具体需求调整的值，灵活控制专家的选择数量。

KeepTopK 策略是许多大型语言模型（LLM）仍在使用的一种方法，尽管有许多有前景的替代方案。需要注意的是，KeepTopK 也可以在不添加额外噪声的情况下使用。

Token Choice

KeepTopK 策略将每个词路由到几个选定的专家。这种方法称为 Token Choice，允许给定的词被发送到一个专家（top-1 路由）或多个专家（top-k 路由）。

或者发送到多个专家（top-k 路由）：

辅助损失

一个主要的好处是，它允许对各个专家的贡献进行加权和整合。

为了在训练过程中实现专家的更均匀分布，辅助损失（也称为负载均衡损失）被添加到了网络的常规损失中。

它增加了一个约束，迫使专家具有相等的重要性。这个辅助损失的第一个组成部分是在整个批次中对每个专家的路由器值进行求和：

这为我们提供了每个专家的重要性分数，这些分数表示无论输入是什么，给定专家被选中的可能性。

我们可以利用这些分数来计算变异系数（CV），它告诉我们专家之间的重要性分数的差异程度。

例如，如果重要性分数差异很大，变异系数（CV）将会很高：

相反，如果所有专家的重要性分数相似，变异系数（CV）将会很低（这是我们所期望的）：

利用这个 CV 分数，我们可以在训练过程中更新辅助损失，使其尽可能降低 CV 分数（从而给予每个专家相等的重要性）：

最后，辅助损失作为单独的损失在训练过程中进行优化。

2、专家容量

不平衡不仅存在于被选择的专家中，还存在于发送到专家的令牌分布中。

例如，如果输入令牌不成比例地发送到一个专家而不是另一个专家，这可能会导致训练不足。

这里，问题不仅仅在于使用了哪些专家，还在于对它们的使用程度。

一个解决方案是限制给定专家可以处理的令牌数量，即专家容量（Expert Capacity）。当一个专家达到其容量时，后续的令牌将被发送给下一位专家。

如果两个专家都达到容量了，这个词就不会被任何专家处理，而是直接送到下一层。这就叫做词溢出。

简单来说，就好比有两个专家在处理任务，但他们都忙不过来了，手上的活儿已经堆得满满的。这时候再来的新任务，他们就没办法处理了，只能把这个任务往后传，让下一层的专家或者别的机制来处理。这就叫词溢出。

3、用 Switch Transformer 简化 MoE

最早解决 MoE 训练不稳定性问题（如负载均衡）的基于 Transformer 的 MoE 模型之一是 Switch Transformer。它极大地简化了架构和训练过程，同时提高了训练的稳定性。

Switch Transformer

Switch Transformer 是 Google 在 2022 年发表的一篇论文中提出的一种模型，它简化了 MoE 的路由算法，减少了计算量和通信量，同时支持 bfloat16 精度进行训练。基于 T5-Base 和 T5-Large 设计的模型在相同的算力下训练速度提升了 7 倍；同时发布了 1.6 万亿参数的 MoE 模型，相比 T5-XXL 模型训练速度提升了 4 倍。

模型介绍

Switch Transformer 是一个 T5 模型（编码器-解码器），它用切换层替换了传统的前馈神经网络层。切换层是一个稀疏的 MoE 层，它为每个词元选择一个专家（Top-1 路由）。

简化稀疏路由

在传统的 MoE 层中，网络接收一个输入词元，输出会路由到 top-K 个专家进行计算。Switch Transformer 采用了一种简化的路由策略，即每次只发给一个专家，这样可以显著减少路由的计算量，同时保证模型的性能。

高效稀疏路由

Switch Transformer 通过以下方式实现高效稀疏路由：

路由计算量减少： 只有一个专家被激活，减少了计算量。

专家容量减半： 专家中的 batch_size 至少减半，进一步减少了计算量。

简化路由实现： 减少了传统 MoE 方法中通信的代价。

分布式 Switch 实现

Switch Transformer 的分布式实现中，所有核心被分配给数据分区维度，这也对应于模型中的专家数量。对于每个核心的每个词元，路由器会本地计算分配给专家的分配。输出是一个大小为的二进制矩阵，该矩阵在第一维度上进行分区，并确定专家的分配。

负载均衡损失

为了促使每个专家都可以拿到近似均匀分布的样本，Switch Transformer 引入了负载均衡损失。当时，损失是最小的。

切换层

简单来说，就是路由器在挑专家的时候，就是按照老规矩来，把输入数据和专家的权重相乘，然后用 SoftMax 函数处理一下，看看哪个专家最合适。没有啥复杂的操作，就是常规操作。

Switch Transformer 用了一种很简单的办法来挑专家。以前的方法是让每个词元找好几个专家来处理，这种方法叫 top-k 路由。但 Switch Transformer 说，其实每个词元找一个专家就足够了，这种方法叫 top-1 路由。这样做的好处是，计算量变小了，通信成本也降低了，而且训练起来更稳定。

容量因子

容量因子就像是给每个专家分配任务的“额度”。这个值很重要，因为它决定了每个专家能处理多少个词元。Switch Transformer 在这方面做了改进，通过直接调整容量因子来控制每个专家的处理能力。这样做的好处是，可以更好地平衡每个专家的负载，避免有些专家忙不过来，而有些专家却很闲。

tokens per batch： 每个批次中的令牌总数。

number of experts： 专家的数量。

capacity factor： 容量因子，通常大于 1，以提供额外的缓冲空间。

专家容量的组成部分其实挺简单的。专家容量主要由两个因素决定：每个批次中的令牌数量和专家的数量。

如果我们把容量因子调大，每个专家就能处理更多的词元。这听起来好像挺好的，因为专家们能干更多的活儿。

但是，如果把这个因子调得太大，就会浪费很多计算资源，因为有些专家可能会闲着没事干。反过来，如果容量因子太小，专家们就会忙不过来，处理不了所有的词元，这就叫令牌溢出。结果就是，模型的性能会变差，因为有些词元没被处理好。所以，容量因子得找个合适的值，这样才能让专家们既不过载，也不闲置，模型性能也能保持在最佳状态。