找回密码
 立即注册
首页 业界区 业界 探秘Transformer系列之(29)--- DeepSeek MoE

探秘Transformer系列之(29)--- DeepSeek MoE

蓟晓彤 5 天前
探秘Transformer系列之(29)--- DeepSeek MoE


目录

  • 探秘Transformer系列之(29)--- DeepSeek MoE

    • 0x00 概述
    • 0x01 难点

      • 1.1 负载均衡
      • 1.2 辅助损失函数

        • 1.2.1 定义
        • 1.2.2 损失函数
        • 1.2.3 业界鼻祖
        • 1.2.4 GShard
        • 1.2.5 Switch Transformers
        • 1.2.6 通用思路
        • 1.2.7 小结

      • 1.3 专家策略

        • 1.3.1 专家能力
        • 1.3.2 专家选择


    • 0x02 DeepSeek V1

      • 2.1 背景&动机
      • 2.2 解决方案
      • 2.3 负载均衡

        • 2.3.1 专家级负载损失(Expert-Level Balance Loss)
        • 2.3.2 设备级负载损失(Device-Level Balance Loss)

      • 2.4 实现

    • 0x03 DeepSeek V2

      • 3.1 负载均衡
      • 3.2 实现

    • 0x04 DeepSeek V3

      • 4.1 架构

        • 4.1.1 Sigmoid

          • 概念
          • 对比

        • 4.1.2 专家分组

      • 4.2 负载均衡

        • 4.2.1 无辅助损失负载均衡
        • 4.2.2 sequence粒度的负均衡损失
        • 4.2.3 路由

      • 4.3 vs Llama 4
      • 4.4 实现

    • 0x05 其它探索

      • 5.1 MOA
      • 5.2 MoD

        • 5.2.1 概要
        • 5.2.2 思路
        • 5.2.3 Budget
        • 5.2.4 路由机制
        • 5.2.5 采样
        • 5.2.6 路由讨论

      • 5.3 MoE 与大型多模态模型的结合
      • 5.4 LoRA 与专家混合

        • 5.4.1 分类
        • 5.4.2 LoRAMoE

          • 动机
          • 方案
          • 训练

        • 5.4.3 HydraLoRA

          • 动机
          • 观察
          • 方案
          • 训练&推理


      • 5.5 高效微调

    • 0xFF 参考


0x00 概述

MoE有两个鲜明特点:

  • 动态路由:使用门控网络(Gating Network)来决定每个输入应由哪些专家处理。
  • 稀疏激活:对于每个输入,只有部分专家被激活,大大减少了计算量。
也正是由于这两个特点导致了几大难题:

  • 负载均衡。某些专家的过度使用导致负载分布不均。
  • 路由网络退化。由于门控网络路由决策的过拟合,探索能力下降。
  • 参数爆炸。专家数量增加导致过高的内存和存储需求。
  • 通信瓶颈。在分布式系统中,专家之间的高通信开销尤其突出。
  • 内存碎片化。不高效的内存使用导致训练期间出现内存不足错误。
  • 在模型精度(trim token)和硬件效能(zero padding)之间如何*衡(Tradeoff )。
这些难题中,又以负载均衡最为典型,我们接下来进行深入分析,看看业界如何应对负载均衡问题。
注:

  • 全部文章列表在这里,估计最终在35篇左右,后续每发一篇文章,会修改此文章列表。cnblogs 探秘Transformer系列之文章列表
  • 本系列是对论文、博客和代码的学习和解读,借鉴了很多网上朋友的文章,在此表示感谢,并且会在参考中列出。因为本系列参考文章实在太多,可能有漏给出处的现象。如果原作者或者其他朋友发现,还请指出,我在参考文献中进行增补。
0x01 难点

我们先给出一个例子来进行说明。
我们需要开发应用,但是此应用需要有前台开发和后台开发,那么招聘程序员有两种选择:

  • 雇佣一位同时精通前台开发和后台开发的程序员,这样他可以完成所有开发。这类似于标准 Transformer 模型,由单个 FFN 子层处理所有输入 token。
  • 雇佣多位各有所长的程序员,比如前台开发程序员和后台开发程序员,再加上一位开发经理来分配工作。这类似于 MoE 方法,每个程序员充当专家,开发经理则作为门控机制(Gating)来选择专家。
为确保该系统高效运作,需满足以下条件:

  • 每位程序员必须精通自身工作所需技能,同时所有程序员需能共同完成应用开发。
  • 开发经理需充分了解所有程序员的专长,并能高效分配任务。
1.1 负载均衡

虽然稀疏门控G(x;Θ)能在不增加计算成本的情况下显著扩展模型参数空间,但其性能高度依赖门控机制的有效性 。MoE的公式是“每遇到一个Token,就去找相应的Expert来计算”,但实际训练时是反过来的:先给每个Expert分配好相应的算力,然后将Token分配(Route)到所属的Expert中并行计算,这也就为什么负责打分的门控机制也被称为Router的原因。因为门控机制无法控制发给专家的token的概率,所以在实际操作中,会存在专家间工作负载分布不均衡的情况。某些专家被频繁使用(接收到了很多token)而其他专家却很少被调用(接收的token寥寥无几)。我们管这种现象叫expert负载不均衡。这不仅不符合MoE的设计初衷(术业有专攻),还影响计算效率(例如引起分布式训练中各卡通讯时的负载不均)。
要理解专家负载均衡,可以从batch size的角度出发。通常来讲,较大的 batch size 推理性能更好,但是由于样本在MoE层激活专家时需要并行,MoE层中每个专家的实际batch size会减少。举个例子,假设当前 batch 有10个token,其中5个token被路由到了某个专家网络,而其他5个token被路由到了其它5个不同的专家网络,这就会让各专家网络获得的 batch size 不均匀,一个专家的batch size是另外5个专家的5倍,导致另外5个专家利用率不足和硬件资源浪费。并且,如果我们把所有的token都发送给少数几个头部专家,训练效率将会变得低下:

  • 假如前面几个token增加了这些被选择专家的门控权重,这反过来导致它们以更高的门控权重更频繁地被选择。这导致它们训练更多,它们的门控权重再次增加。这种情况会不断自我强化,因为受青睐的专家训练得更快、更充分,它的效果一直在被优化,因此选择它们的频率也会更高。那么此时,这几个少数专家就会过载,每次需要计算大量token,其他的专家模型就得不到训练而被闲置,就会导致工作负荷不足的专家会因为训练 tokens 不足,而难以学习有效知识。进而导致模型失衡,性能下降。
  • 由于梯度冲突(gradient conflict),路由奔溃也会导致训练不稳定。超负荷工作的专家接收更多 input token,他们积累的梯度也会更大,因此,超负荷工作专家与负荷不足专家的梯度在幅值和方向上均可能发生偏离,导致训练难以收敛。
因此,MoE模型中需要控制专家均衡。为了对抗这种不*衡,人们实施了各种负载均衡技术,目的是确保所有专家在训练过程中大致承担相等的责任。我们在前一篇就介绍过,在门控中加入噪声机制可以避免反复路由到同一个专家,让门控在做出选择时跳过几个最优专家模型,去选择其余的专家模型,从而更好地协同工作。具体如下。

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

仅仅加入噪声是不够的,人们针对这一问题提出了多种解决方案。其中,最常用的策略是为负载均衡添加辅助损失函数(Auxiliary Loss)和基于专家来施加的一些策略,比如专家能力和专家选择(Expert Choice)。
本节后续介绍的知识点都是解决负载均衡的相关技术。
1.2 辅助损失函数

为了缓解负载失衡,人们在模型训练的目标函数基础上引入了辅助损失函数(auxiliary load balancing loss ),鼓励给予所有专家同等的重视,以促进每批次中令牌(token)在专家间的均匀分配,这些损失项被添加到训练目标中。
1.2.1 定义

论文"A Survey on Mixture of Experts"给出辅助损失函数的定义如下图,N 是专家数量,T 是 token 数量,K 是每个 input token 激活的专家数量。为确保T个令牌(token)在N个专家间的均匀分布,应最小化负载均衡损失函数。下图也给出了损失函数的最优情况,此时,每个专家都获得等量的分发token,Di = 1/N,以及相等的门控概率比例 Pi = 1/N。这样可以维持所有专家间的负载均衡,确保工作量均匀分布。
2.jpeg

1.2.2 损失函数

因为辅助损失函数是在目标损失函数基础上引进的,所以我们先来看看损失函数。
最早的损失函数来自论文“Adaptive Mixture of Local Experts”。有两种损失函数,分别对应下图的标号1和2。其中,对于数据c,\(p^c_i\)是第 i 个专家占的比重,由门控网络控制,\(o^c_i\)是第 i 个专家的输出,\(d^c\)是最终期望的整体输出。
第一种方案有一个非常严重的问题,即不同的 expert 会互相影响,导致专家网络之间的强烈耦合。因为是所有专家网络的权重加总来共同计算损失的,一个专家权重的变化会影响到其他专家网络的loss。这种耦合可能会导致多个专家网络被用于处理每条样本,而不是专注于它们各自擅长的子任务。
为了解决这个问题,论文重新定义了损失函数,以鼓励专家网络之间的相互竞争,即方案2。直觉上改进就是每个 expert 单独计算 loss,多个 expert 的 loss 加权得到整体的 loss。这意味着,每个expert在处理特定样本的目标是独立于其他expert的。如果门控网络和expert都使用这个新的loss进行梯度下降训练,系统会倾向于将某类特定样本分配给特定的expert。因为当一个expert在给定样本上的的loss小于所有expert的*均loss时,它对该样本的门控score会增加;当它的表现不如*均loss时,它的门控score会减少。这种机制鼓励expert之间的竞争,而不是合作,从而提高了学习效率和泛化能力。
3.jpeg

我们接下来看看如何增加辅助损失函数。
1.2.3 业界鼻祖

论文“Outrageously large neural networks: The sparsely-gated mixture-of-experts layer”就首创设计了额外的损失函数来促使所有专家具有同等的重要性,也首创了具有辅助负载*衡损失的可微分启发式方法,通过选择概率对专家输出进行加权,使门控过程可微分,从而能够通过梯度优化门控函数。这种方法随后成为MoE领域的主流研究范式。
论文先为每个专家定义了其相对于一批训练样本的重要性 Importance(X) ,即该专家在这批样本中门控值的总和,在批次中经常选择的专家将具有很高的重要性分数。然后定义了一个额外的损失函数 \(L_{importance}(X)\) ,这个额外损失函数被添加到模型的整体损失函数中。这个损失函数等于重要性值集合的CV(coefficient of variation)*方(变异系数,可以度量一组数据的离散程度),乘以一个手动调整的缩放因子 \(w_{importance}\) 。这个额外的损失鼓励所有专家在所有批次中具有相等的重要性。
虽然这种损失函数可以确保同等重要性,但专家可能仍然会收到数量截然不同的示例。例如,一位专家可能会收到一些权重较大的示例,而另一位专家可能会收到许多权重较小的示例。这可能会导致分布式硬件上的内存和性能问题。为了解决这个问题(鼓励每个expert拿到相同数量的样本进行计算),论文引入了第二个损失函数 \(L_{load}\) 来进一步确保负载均衡。
具体计算公式推导如下图所示。
4.jpeg

1.2.4 GShard

对于损失函数的设计,有一个简单的想法:例如有 S 个token,E 个专家,假设\(c_e\)是第 e 个专家接受的token数量。如果MoE是负载均衡的,那么每个专家收到的token一样多,token比例\(\frac{c_e}{S}\)应该是相同的。因此GShard把每个专家收到的token的*方和定义为负载均衡损失(见下图标号1),所有专家收到token比例都相等时,\(l_{aux}\)得到最小值。但是\(c_e\)是从top-k函数中导出,不是可微的。而且,这样的辅助损失函数并不包含Gating函数的参数,无法训练进行梯度更新。于是GShard中引入了一个做法, 把*方项的一个分量\(\frac{c_e}{S}\)替代成Gating softmax的均值\(m_e\)(见下图标号2),因为\(c_e\)和\(m_e\)有一定依赖关系,对\(m_e\)的调整也会影响\(c_e\),从而实现对各专家负载分配的调节。
具体推导过程见下图。
5.jpeg

为了进一步防止某些专家超负荷,GShard在训练时又引入以下概念:

  • 专家容量 (Expert Capacity),约为 \(C \approx \frac{2N}{E}\)(如果共有 N 个 token,E 个专家)。一旦某个专家被分配的 token 超过了这个容量,就可能丢弃一些 token(或者通过残差连接导入到下一层再处理)。
  • 本地分组 (Local Groups):不是所有 token 都在全局竞争,而是先进行分组后再分配给专家。这样能缩小混乱程度。
1.2.5 Switch Transformers

为了进一步防止 token 被丢弃,Switch Transformer 引入了简化版的辅助损失(类似 GShard 的辅助损失)来*衡重要性分数并在专家之间执行负载均衡,以希望每个专家从批次中获得的token数量大致相等。具体如下图所示,其中 α 是一个超参数,用于在训练过程中微调此损失的重要性。值过高会影响主要损失函数,而值过低则无法有效进行负载*衡。
该辅助损失函数不再计算变异系数,而是将分配的 token 数量与每个专家的路由概率进行加权比较。该损失相对于 P 是可微分的,并且可以很容易地纳入训练中,它鼓励分配给每个专家的token分数和分配给每个专家的路由概率分数都是 1/N,这意味着每个专家同样重要,并且接收到*衡数量的token。
6.jpeg

1.2.6 通用思路

苏剑林大神给出了构建Aux Loss的通用思路:首先基于F构建符合要求的损失,然后在实现时将F替换成P+sg[F−P]。此处F是Expert当前的负载分布,而P则相当于F的一个光滑*似。sg[]是stop gradient算子,特点是保持前向输出不变,但强制梯度为零。假设 Q是均匀分布\(Q = (1/n,1/n,...,1/n)\),负载均衡等价于F=Q。下式就是一个比较直观的Aux Loss:

\[L_{aux} = \frac{1}{2}||F-Q||^2  = \frac{1}{2}\sum^n_{i=1}(F_i-1/n)^2 \tag 1\]
如果F是\(argtop_k\)的输出,则上式并不是一个能直接用的可导目标,可以通过STE(Straight-Through Estimator)技巧来在在反向传播的时候将F替换成P。
另外,常见辅助负载均衡损失公式如下,其中n是expert个数,p是router输出的概率,f是每个expert在所有token上被topk取到的概率。

\[L_{aux} =  = \sum^n_{i=1}f_ip_i \tag2\]
公式(1)和公式(2)的梯度是等价的,这说明公式(2)可以使专家负载均衡化。
1.2.7 小结

下图是各种辅助损失函数及其典型系数配置概述。我们以粗体来标注每种辅助损失的作者,用下划线来表示对原始方案进行修改的优化者。
7.jpeg

1.3 专家策略

1.3.1 专家能力

因为GShard和Switch Transformers都提到了专家能力(Expert Capacity,也叫做专家容量),而GShard的实现相对简单,因此我们以Switch Transformers为例来仔细学习下。
Switch Transformers的模型是面向TPU,其要求所有张量形状都是在编译期间内静态固定的。但因为路由是动态的,因此路由到每个专家的张量形状也是动态的。为了解决这一问题,Switch Transformers使用了专家容量(Expert Capacity)。如下图所示,专家容量为每个batch中总 Token 数除以专家数,然后再乘以容量因子(Capacity Factor),即可得到专家容量(每个专家对应的 Token 数)。该变量决定了可以路由到任何一个专家的最大token数量。如果当前专家接收到的token数已经超过了容量,那么它就不再接收token了,此时我们称这些路由到已满容量的专家的token为溢出token。
8.jpeg

针对专家容量,Switch Transformers也引入了两种处理策略:

  • Zero padding。因为存在容量,所以Switch Transformers给每个专家设定了一个Expert buffer用来接收在容量范围内的token。如果某些专家的Expert buffer没有填充满,则使用0向量填充,我们称为Zero padding。具体而言,最终每个expert上的输入数据维度为(E, C, M),其中C表示capacity。这样可以保证每个expert上要处理的输入数据维度是一样的,这有利于硬件层面的后续处理。
  • Drop tokens。如果向单个专家发送了太多token(即超出专家能力),致使某个专家发生溢出情况,Switch Transformers会跳过这些token的计算。这些token被通过残差连接直接传递到下一层。因为这些token没有经过专家处理,因此是一种信息损失。
Zero padding和Drop tokens是一体两面。缩小容量会减少Zero padding,但是会因为 token 溢出而导致模型性能下降。扩大容量可以缓解Drop tokens,但过大的容量会引起更严重的zero padding问题(影响到矩阵的稀疏程度),无效计算越多,浪费计算资源。另外,负载越不均衡,需要 Padding 的 Token 也就越多,无效计算越多,丢弃的token也可能更多。
从上图也可以看出,容量因子越大,需要 Padding 的 Token 也就越多,无效计算越多;负载越不均衡,需要 Padding 的 Token 也就越多,无效计算越多。为了更好的负载均衡,作者同样添加了 Load Balancing Loss。上图有 6 个 Token,3 个专家,我们分别看看两种不同的专家容量设置情况。

  • 容量因子为 1.0:如上图左侧所示,对应的专家容量为 2:

    • Expert 1有 3 个 Token,则需要丢弃一个token,或者把该token通过残差连接直接传到下一层。
    • Expert 2 有 2 个 Token,恰好等于专家容量。
    • Expert 3只有 1 个 Token,需要填充 1 个空的 Token。

  • 容量因子为 1.5:如上图右侧所示,对应的专家容量为 3:

    • Expert 1 有 3 个 Token,恰好等于专家容量。
    • Expert 2 只有 2 个 Token,需要填充 1 个空的 Token。
    • Expert 3 只有 1 个 Token,需要填充 2 个空的 Token。

对于填充,我们可以使用掩码来进行处理,假设 S = batch_size * seq_len,E代表专家数目,C表示capacity(Expert buffer),则之前的门控网络的输入是:(S, embedding_size),输出是 (S, E)。含有容量处理的门控网络的输入是(S, embedding_size),输出是:

  • combine_weights =  (S, E, C)。表示对 S  个token中每个 token 而言,它去到每个专家的概率。而这个概率按照该token在buffer中的位置(C)存放,不是目标位置的地方则用0填充。比如[[[p1, 0, 0]], [0, p2, 0], [0,0,0]],表示第一个token会去两个专家,去第一个专家的概率是p1,去第二个专家的概率是p2。
  • dispatch_mask =  (S, E, BF)。combine_weights 为0的地方设为False,为1的地方设为True,填充时候会用到mask。
我们用实际例子来看看。假设的输入
• Token 数 S = 3
• Expert 数 E = 4
• Buffer 长度 C = 3(每个专家最多接收 3 个 token)
• Embedding 维度 M = 512

reshape_input (
来源:程序园用户自行投稿发布,如果侵权,请联系站长删除
免责声明:如果侵犯了您的权益,请联系站长,我们会及时删除侵权内容,谢谢合作!
您需要登录后才可以回帖 登录 | 立即注册