作者: 昇腾实战派
DeepSeek知识地图: https://blog.csdn.net/weixin_45216014/article/details/156450562

一、背景介绍

1.1 负载均衡及其重要性

在我们之前的餐厅类比中,我们用一个提供多种菜系的餐厅来解释MoE的概念:每位厨师扮演一个专家,主厨则充当门控操作,将每道菜分配给具备相应技能的特定厨师。
为了确保这样的系统高效运行,我们需要满足以下两点:

  1. 每位专业厨师必须掌握自己菜系所需的技能(例如,制作饺子的厨师必须知道如何包饺子),同时他们能够共同处理所有菜品。
  2. 主厨需要充分了解所有专业厨师的专长,并能够高效地分配订单。

在MoE中,前者对应专家专业化与知识共享之间的权衡,后者则反映了负载均衡的重要性。

那么,为什么负载均衡如此重要呢?

原因在于,当负载不均衡时,MoE无法有效运行,最常见的问题是“路由崩溃”,即只有一小部分专家接收到大多数输入标记,而其他专家则未被充分利用。

结果,大部分计算由过载的专家完成,导致硬件利用率出现瓶颈,因为专家通常分布在多个GPU核心上。

路由崩溃还会因梯度冲突导致训练不稳定。由于过载的专家接收到更多的输入标记,它们会积累更大的梯度并比未充分利用的专家学习得更快。因此,过载专家和未充分利用专家的梯度在大小和方向上可能产生分歧,使得训练过程更难收敛。

最后,MoE中的负载不均衡还可能导致性能不佳和泛化能力差,因为未充分利用的专家未能获得足够的训练标记来学习有意义的知识。

由于负载均衡在MoE中至关重要,人们提出了多种技术来解决这一问题。在这些先前的研究中,最常用的策略是添加辅助损耗以实现负载均衡以及专家选择法。

1.2 Auxiliary Loss For Load Balancing

改善负载均衡的常见策略是在模型训练的原始目标函数之上引入辅助损耗函数。下图为常用的辅助损失loss之一;
在这里插入图片描述

在这里插入图片描述

参数:

  • N 是专家的数量,T 是标记的数量,K 是每个输入标记激活的专家数量。
  • s_{i, t} 是门控的输出,通过Softmax归一化到[0, 1]区间,表示第t个标记选择第i个专家的概率。
  • u_t 是第t个标记的输入隐藏状态,而 e_i 是第i个专家的质心,可以看作是过去路由到第i个专家的标记嵌入的平均值。因此,s_{i, t} 衡量了当前输入与第i个专家接收的平均标记的接近程度。
  • 因此,P_i 可以看作是整个输入序列中选择第i个专家的平均概率。
  • f_i 表示路由到第i个专家的标记比例。
  • G表示非线性的门控函数,如softmax,sigmoid

需要注意,f_i 是不可微分的,因此最小化上述损耗函数实际上等价于最小化 s_{i, t}。此外,由于 f_i 依赖于 s_{i, t},对 s_{i, t} 的调整也会影响 f_i,从而调整分配给每个专家的负载,但是使用这种辅助损耗来平衡负载是有代价的,其梯度可能会干扰语言建模目标的梯度,导致模型性能下降,特别是在极端不均衡的情况下,过载专家的 f_iP_i 变得非常大时。因此,使用这种方法平衡负载需要谨慎地权衡辅助损耗的权重。

同时DeepSeek在V1,V2版MoE模型中,增加了专家级,设备级和设备通信级等平衡负载辅助loss。这些辅助loss只是为了做计算、通讯的负载均衡,对模型的效果调优并没有帮助。甚至这些辅助loss增加过多,loss太大会对主模型造成影响,导致主模型的效果有损。为了减轻多辅助负载均衡的loss对主模型的影响,在V3版把多辅助loss都精简掉了,通过引入一个可动态调节的bias来做到负载均衡。

作者:姜富春
链接:https://zhuanlan.zhihu.com/p/18565423596
来源:知乎

一、无辅助损失负载均衡

1.1 Auxiliary-Loss-Free Load Balancing公式介绍

为了在负载平衡和模型性能之间实现更好的权衡,deepseek v3首次引入了一种无辅助损失的负载平衡策略,以确保负载平衡。具体来说,我们为每个专家引入了一个偏置项b_i ,并将其添加到相应的亲和度分数 Si,t 中,以确定 top-K 路由:
在这里插入图片描述
在这里插入图片描述

上述偏差 b_i 的计算方式非常直观,如下图所示:首先计算每个专家分配的标记数量的平均值,然后计算每个专家分配的标记数量与平均值的差值,偏差由该差值(或误差)的符号乘以一个固定的更新率决定,更新率是一个可调的超参数。

在这里插入图片描述

需要注意的是,偏置项仅用于路由。门控值(将乘以 FFN 输出)仍然来源于原始亲和度分数 si,t 。在训练过程中,我们持续监控每个训练步骤中整个批次的专家负载。在每个step结束时,如果其对应的专家负载过重,我们将偏置项减少 γ ;如果其对应的专家负载过轻,我们将偏置项增加 γ ,其中是一个超参数,称为偏置更新速度。通过这种动态调整,DeepSeek-V3 在训练过程中保持了平衡的专家负载,并且比通过纯辅助损失来鼓励负载平衡的模型表现更好。

简单总结:

•Auxiliary-Loss-Free Load Balancing:路由FFN输出的值需要加上一个偏置项b,然后才用于计算top-k的路由结果

架构:创新的负载平衡策略和训练目标

在 DeepSeek-V2 的高效架构基础上,开创了一种用于负载平衡的辅助无损耗策略,最大限度地减少了因鼓励负载平衡而导致的性能下降

1.2 代码解读(mindspeed-llm)

1.2.1 使能参数:
--moe-router-load-balancing-type noaux_tc \
--moe-router-score-function sigmoid \
--moe-router-enable-expert-bias

每个参数的详细介绍如下:

MindSpeed-LLM\mindspeed_llm\training\arguments.py

在这里插入图片描述

在这里插入图片描述

调用逻辑如下图:

1.2.2 balance_loss使能

MindSpeed-LLM\mindspeed_llm\core\transformer\moe\router.py

在这里插入图片描述

在这里插入图片描述

参数moe-router-load-balancing-type=noaux_tc,进入对应分支,通过topk_softmax_with_capacity函数进行损失函数初始化,同时通过参数moe-router-score-function=sigmoid,传参到score_function,透传到topk_softmax_with_capacity中,选自sigmod函数

1.2.3 expert_bias初始化

通过参数moe-router-enable-expert-bias进行使能,并初始化expert_bias

在这里插入图片描述

MindSpeed-LLM\mindspeed_llm\core\transformer\moe\moe_utils.py

1.2.4 topk_softmax_with_capacity函数介绍:

在这里插入图片描述

在这里插入图片描述

根据score_function=sigmoid选择执行分支,获取当前的scores_for_routing,然后传入compte_topk函数,获取每个token分配的专家

二、序列维度辅助损失

2.1 Complementary Sequence-Wise Auxiliary Loss

尽管 DeepSeek-V3 主要依赖无辅助损失策略来实现负载平衡,但为了防止任何单个序列内的极端不平衡,我们还采用了补充的序列级平衡损失:

在这里插入图片描述

其中,平衡因子 α 是一个超参数,其值将被赋予 DeepSeek-V3 极其小的值;l(.) 表示指示函数;T表示序列中的token数量。序列级平衡损失鼓励每个序列上的专家负载保持平衡。

简单总结:

•Complementary Sequence-Wise Auxiliary Loss:尽管主要依赖于无辅助损失的策略来实现负载均衡,为了防止任何单个序列内的极端不平衡,还采用了一个互补的序列级平衡损失鼓励专家在每个句子中的投票平衡一些

2.2 代码解读(mindspeed-llm)

2.2.1 参数

需要与Auxiliary-Loss-Free Load Balancing的参数配合使用,代码待合入,PR为:https://gitee.com/ascend/MindSpeed-LLM/pulls/2331/files

--seq-aux

在这里插入图片描述

2.2.2 seq_aux_loss使能

MindSpeed-LLM\mindspeed_llm\core\transformer\moe\router.py

在这里插入图片描述

2.2.3 实现代码

MindSpeed-LLM\mindspeed_llm\core\transformer\moe\router.py

在这里插入图片描述

计算得出aux_loss,通过MoEAuxLossAutoScaler应用到前向和反向;

save_to_aux_losses_tracker保存aux_loss到日志

https://github.com/NVIDIA/Megatron-LM/blob/core_r0.8.0/megatron/core/transformer/moe/moe_utils.py#L101

在这里插入图片描述

参考文档:

https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf

https://arxiv.org/abs/2408.15664

https://zhuanlan.zhihu.com/p/25228000281

https://zhuanlan.zhihu.com/p/20761021507

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐