
DeepSeek-V3解析 3: 无辅助损失的负载均衡
在本文中,我们解释了DeepSeekMoE中使用的无辅助损失负载均衡方法,这是DeepSeek模型中采用的主要架构创新之一。具体来说,我们首先介绍了混合专家(MoE)的基础知识,强调了负载均衡的重要性,并回顾了先前的解决方案,包括辅助损失方法和专家选择。然后,我们解释了DeepSeek的无损失负载均衡方法及其性能。我们的主要结论是,DeepSeek的无损失方法避免了引入梯度干扰,同时保持了因果性,
目录
DeepSeek如何突破MoE中的隐藏瓶颈
作者使用ChatGPT创建的图片。
这是我们DeepSeek-V3系列的第三篇文章,我们将探讨DeepSeek [1, 2, 3]模型中与混合专家(Mixture-of-Experts, MoE)相关的另一个关键架构突破:无辅助损失的负载均衡[5]。
在本文中,我们将深入探讨DeepSeek如何解决MoE的隐藏瓶颈 - 负载均衡,同时消除梯度干扰并保持因果性,为基于专家的模型的效率设立了新标准。
本系列有3篇文章,请在站内搜索:
-
DeepSeek-V3解析1:多头潜在注意力
-
DeepSeek-V3解析2: DeepSeekMoE
-
DeepSeek-V3解析3: 无辅助损失的负载均衡
本文目录:
-
背景:介绍混合专家(MoE)的基础知识,解释为什么负载均衡很重要,并回顾先前的工作,包括辅助损失方法和专家选择。
-
DeepSeek的无辅助损失负载均衡:解释其工作原理。
-
评估:讨论无辅助损失负载均衡的性能。
-
总结。
背景
Transformer中的MoE(混合专家)
MoE代表混合专家,在Transformer模型的上下文中,这通常指的是将每隔几个Transformer层中的FFN替换为多个FFN,每个FFN充当一个专家。因此,当输入token到来时,门控操作将选择top-K个专家并将输入token仅路由到所选的FFN,因此只有被选中的专家会被激活。
这在下图中显示,其中左侧标准Transformer层中的FFN子层被MoE层替换。
图1. Transformer中MoE层(红框中高亮)的示意图。图片编辑自[6]。
负载均衡及其重要性
在我们之前的餐厅类比中,我们用一家提供多种菜系的餐厅来解释MoE的概念,其中每位厨师充当专家,一位主厨充当门控操作,将每道菜分配给具有适当技能的特定厨师。
为确保这样一个系统有效运作,我们需要:
-
每位专业厨师必须掌握自己菜系所需的技能(例如,制作饺子的厨师必须知道如何制作饺子),同时他们可以共同处理所有菜品。
-
主厨对所有专业厨师的专长有很好的了解,并能有效地分配订单。
在MoE中,前者对应于专家专业化与知识共享之间的权衡,我们在DeepSeekMoE中已详细讨论过。后者反映了负载均衡的重要性,这是本文的主要主题。
那么是什么使负载均衡如此重要呢?
原因是当负载不平衡发生时,MoE无法很好地运作,最常见的问题被称为"路由崩溃",即只有一小部分专家接收大部分输入token,而其他专家则未被充分利用。
结果,大部分计算由过载的专家进行,导致硬件利用率出现瓶颈,因为专家通常分布在多个GPU核心上。
路由崩溃还可能由于梯度冲突导致训练不稳定。由于过载的专家接收更多的输入token,它们也会累积更大的梯度并比未充分加载的专家学习得更快。因此,来自过载专家和未充分加载专家的梯度可能在幅度和方向上都存在分歧,使训练过程更难收敛。
最后,MoE中的负载不平衡还可能导致性能不佳和泛化能力差,因为未充分加载的专家没有获得足够的训练token来学习有意义的知识。
由于负载均衡在MoE中如此关键,已提出了不同的技术来处理这个问题。在这些先前的努力中,最常用的策略是添加辅助损失以实现负载均衡和专家选择。
使用辅助损失进行负载均衡
改善负载均衡的一种常见策略是在模型训练中引入辅助损失函数,在原始目标之上。
图2. 用于强制负载均衡的辅助损失示例。图片编辑自[5]。
上图显示了辅助损失函数的一个例子,其中
-
N是专家数量,T是token数量,K是每个输入token激活的专家数量。
-
s{i, t}是门控的输出,已通过Softmax归一化到[0, 1],表示第t个token选择第i个专家的概率。向量u_t是第t个token的输入隐藏状态,而e_i是第i个专家的质心,可以看作是过去路由到第i个专家的平均token嵌入。因此,**s{i, t}衡量当前输入与第i个专家接收的平均token的接近程度。**
-
因此,P_i可以看作是整个输入序列选择第i个专家的平均概率。
-
f_i表示路由到第i个专家的token比例。
注意,f_i是不可微的,因此最小化上述损失函数实际上变成了最小化s{i, t}。此外,由于f_i依赖于s{i, t},应用于s_{i, t}的调整也会影响f_i,结果分配给每个专家的负载将得到调整。
然而,使用这种辅助损失来平衡负载是有代价的,因为它的梯度可能会干扰语言建模目标的梯度,导致模型性能下降,特别是在极度不平衡的情况下,过载专家的f_i和P_i都变得非常大。
因此,使用这种方法平衡负载需要仔细权衡辅助损失。为了更清楚地看到这一点,[5]的作者进行了一项实验,通过训练具有不同alpha值的模型,结果如下图所示,其中y轴显示模型性能(以困惑度表示),x轴显示MaxVio,这是一个指示负载不平衡程度的指标,其中更高的MaxVio意味着更严重的负载不平衡(i代表第i个专家):
图3. 辅助损失控制训练中负载均衡与模型性能之间的困境。图片来自[5]。
如上图所示,当alpha太小时(alpha = 0),MaxVio保持较高,这意味着辅助损失不够有效,无法实现负载均衡目标。另一方面,当alpha变得太大时(alpha = 0.01),模型最终产生了更高的困惑度。
总之,辅助损失控制的负载均衡是一把双刃剑,如果alpha没有仔细调整,可能会对模型性能产生负面影响。在实践中,由于资源限制,在LLM训练期间调整alpha具有挑战性,这进一步复杂化了优化过程。
上图还将提出的无损失方法放在相同的困惑度-MaxVio图表中,其中它实现了低困惑度和低MaxVio,显示了无损失方法的有效性。
专家选择
我们在这里想要提到的另一项先前工作是专家选择[7],它提出了一种简单而有效的负载均衡方法,通过将路由策略从"token选择"转变为"专家选择"。
更具体地说,MoE路由中的门控分数通常使用Softmax在亲和矩阵之上计算,如图2所示。传统的路由方法从token维度应用Softmax来为每个token选择专家,因此这些方法被称为"token选择"。问题是,在这种机制下,我们无法控制每个专家将接收多少token,这最终会导致负载不平衡问题。
另一方面,专家选择通过从专家维度应用Softmax来选择路由到每个专家的token。因此,每个专家接收的token在设计上是完全平衡的,因此不需要辅助损失来进行负载均衡。在原始论文[7]中,这种方法展示了更好的模型性能和更快的训练速度。
然而,专家选择的一个限制是未来token泄露问题,因为每个专家在看到所有token路由分数后决定处理哪个token,这违反了因果性,在自回归任务(如文本生成和机器翻译)中可能是一个严重问题。
DeepSeek的无辅助损失负载均衡
为了解决负载均衡问题,同时不引入梯度干扰,DeepSeek提出了一种名为无损失均衡的新技术,通过直接调整门控分数s_{i, t}。
如我们之前提到的,当我们最小化图2中所示的辅助损失时,最终变成了通过调整s_{i, t}来最小化P_i。
因此,如果我们可以直接调整s_{i, t},理论上我们应该能够达到类似于应用辅助损失的效果。
为此,在每个专家的门控分数中添加了一个专家特定的偏置,如下图所示。注意,b_i不用于最终的门控分数(正如我们稍后将看到的,这个偏置也是不可微的),但它用于TopK来选择专家:
图4. 向门控分数引入偏置b_i。图片来自[5]。
上述偏置b_i的计算方式非常直观,如下图所示:我们首先获取分配给每个专家的token的平均数量及其平均值,然后获取每个专家分配的token与平均值之间的差异,偏置由这个差异(或误差)的符号乘以一个固定的更新率确定,这是一个可调的超参数。在后面的章节中,我们将看到更多关于这个超参数影响的实验。
图5. DeepSeek无损失负载均衡算法。图片来自[5]。
现在我们可以在下表中总结不同负载均衡方法的优势和局限性:
图6. 不同负载均衡方法的比较。图片来自[5]。
我们在图3中看到,提出的方法在模型性能和负载均衡方面实现了更好的权衡,但仍有许多方面需要检验。在下一节中,我们将仔细查看实验结果。
评估
基本上,有三个重要问题需要回答:
-
提出的方法能否在性能和负载均衡之间实现更好的权衡?
-
图5中的更新率u的影响是什么?
-
我们能否进一步优化偏置更新规则(考虑到它如此简单)?
性能与负载均衡
为了回答第一个问题,作者对1B和3B模型进行了实验,比较了损失控制和无损失负载均衡的困惑度和MaxVio,结果如下图所示:
图7. 损失控制和无损失负载均衡的比较。图片来自[5]。
上述结果与我们在图3中看到的类似,提出的方法同时实现了更低的困惑度和更低的MaxVio。
除了评估最终检查点,作者还展示了训练过程中的MaxVio曲线,以便我们能更好地理解提出的方法在整个训练过程中的表现,如下图所示:
图8. 训练过程中的MaxVio曲线。图片来自[5]。
如上图所示,在1B和3B设置中,无损失方法在训练过程中都表现出更好的负载均衡能力,显示了该方法的稳定性。
超参数(更新率)的影响
如图5所示,提出的方法引入了一个新的超参数u,称为更新率,因此一个自然的问题是这个超参数将如何影响无损失方法的有效性。更具体地说,我们需要了解无损失方法对u的值是敏感还是稳健,以及如何选择一个值来最好地利用提出的方法。
如前所述,向门控分数添加偏置项在概念上类似于直接对门控分数应用梯度更新,而不依赖于通过损失函数的反向传播。在这种情况下,更新率u的作用类似于梯度更新中的步长。从这个角度来看,我们可以预期类似的影响:小更新率可能导致收敛缓慢,而过大的更新率可能导致不稳定和波动。
在原始论文中,作者进行了实验,使用从1e-4到1e-2的不同更新率,结果如下图所示:
图9. 更新率对训练负载均衡的影响。图片来自[5]。
正如我们所预期的,当u太小时(在这种情况下为1e-4),MaxVio下降得更慢,但过大的u也会产生负面影响,由于更大的波动,导致整个训练过程中MaxVio更高。
替代偏置更新规则
为了回答第三个问题,作者尝试了几种替代策略,并将它们与提出的版本进行比较:
-
变体1:使用e_i的值计算偏置,而不是仅使用其符号,即从b_i = b_i +u∗sign(e_i)变为b_i = b_i +u∗e_i。
-
变体2:使用乘法偏置而不是加法偏置。
其中变体2可以更正式地描述如下:
根据他们的实验,变体1导致略微更好的负载均衡,但并未改善模型性能:
图10. 变体1的性能。图片来自[5]。
而变体2甚至显示出略微较差的模型性能:
图11. 变体2的性能。图片来自[5]。
以上所有结果表明,最简单的策略反而是最好的。
总结
在本文中,我们解释了DeepSeekMoE中使用的无辅助损失负载均衡方法,这是DeepSeek模型中采用的主要架构创新之一。
具体来说,我们首先介绍了混合专家(MoE)的基础知识,强调了负载均衡的重要性,并回顾了先前的解决方案,包括辅助损失方法和专家选择。然后,我们解释了DeepSeek的无损失负载均衡方法及其性能。
我们的主要结论是,DeepSeek的无损失方法避免了引入梯度干扰,同时保持了因果性,其有效性通过原始论文的实证结果得到了证明。
感谢阅读!
更多推荐
所有评论(0)