SE-MoE：可拓展分布式MoE训练及推理框架（百度）

文章链接：SE-MOE: A SCALABLE AND EFFICIENT MIXTURE-OF-EXPERTS DISTRIBUTED TRAINING AND INFERENCE SYSTEM

登载刊物：NAACL2021

源码链接：https://github.com/PaddlePaddle/Paddle

目录

1.背景介绍

近年来，大规模神经网络在许多机器学习任务中表现出色，例如自然语言处理（NLP）和计算机视觉（CV）。同时模型的参数规模也从数百亿个参数扩展。然而，这些密集激活的模型需要大量的计算资源和大量的训练时间，并且超大规模模型的推理性能难以满足实际需求。
2021年12月，最大的单个密集激活模型Megatron-Turing NLG（530 B）在2000多个NVIDIA A100 GPU上训练了3个月，这使得其成本更高，并进一步阻止了其发展成为具有更大参数规模的模型。

与密集激活的模型相比，MoE架构根据输入数据选择性地激活用于训练的参数的子集。给定稀疏性，计算成本关于模型的大小次线性地增加。
目前最大版本的GLaM 具有1.2T参数和每个MoE层总共64个专家，但来自输入批次的每个令牌仅激活95 B（1.2T的8%）参数的子网。与训练GPT-3（175 B）相比，节省了三分之二的电力成本，推理仅需一半的计算资源。尽管有这些好处，MoE模型在计算，通信和存储方面存在挑战和限制。

计算挑战

虽然当参数规模由于专家的增加而扩大时，计算成本保持不变，但训练和推理中的计算面临以下限制。

• 一方面，由于专家选择的不平衡，MoE往往会使培训效果变差。因此，已经提出了许多解决方案，例如添加辅助损耗、使用随机专家和噪声Routing策略。除了限制专家的能力外，它还可以避免低效的训练和计算资源的浪费。
• 另一方面，各层路由选择和辅助损耗的计算更注重调度而非计算，在不重用GPU等高速计算设备的情况下，给CPU设备带来更多负担。同时，计算中引入了大量冗余操作，例如Host2Device（H2D）和Device2Host（D2H）。

通讯挑战

大多数正在进行的研究更加关注由门控网络学习引起的Routing策略的不平衡。然而，由于MoE中参数的激活与输入数据非常密切相关，因此尽管有有效的Routing方法，但当数据不平衡时，经常会发生负载不平衡。
对于跨设备通信，负载不平衡导致各设备步调不一致，造成相互等待同步通信，尤其是在多任务训练中。同时，以Switch Transformer 为例，每个MoE层在前向和后向阶段都需要四次All-to-All通信，这经常跨越节点或集群。当底层网络拓扑无法感知时，机器间通信更容易发生路由冲突和阻塞，导致性能迅速下降。

存储限制

MoE架构受到计算设备中的存储器存储的显著限制。对于密集激活的模型，模型规模通常受到训练时间而不是记忆的限制。例如具有1万亿个参数的密集模型需要大约3个月的时间来使用3072个A100 GPU训练4500亿个令牌。相比之下，具有数万亿参数的MoE模型只需要几周时间来训练相同数量的令牌，因为其计算成本呈次线性增加。
然而，模型规模取决于设备存储器是否承受模型状态的占用。虽然设备的所有可用存储都包含GPU，CPU内存，SSD（固态驱动器）等中的高带宽内存（HBM），但它们之间的I/O延迟彼此不同，使计算等待参数等。构建一个统一高效的存储管理系统来支持稀疏激活的训练，打破存储墙，是一个具有挑战性的问题。

2.内容摘要

本文分析了MoE模型在计算、通信和存储方面的挑战和局限性。提出了一种名为SE-MoE的分布式系统，该系统利用HBM、CPU内存和ssd等异构存储资源来克服MoE模型的内存限制。他们还引入了2D预取调度和融合通信来提高存储效率。对于可伸缩推理，SE-MoE采用环内存方法来跨内存段有效地执行计算任务。实验证明了SE-MoE拥有更高的吞吐量和更低的内存占用。
主要贡献：

• SE-MoE新型分布式系统，能够将MoE模型扩展到数万亿个参数，充分利用包括HBM，CPU内存甚至SSD在内的集群，超越内存墙，实现高效的训练调度。此外，使用2D预取调度和融合通信提高异构存储效率。
• 动态图调度采用了一种新的基于环形存储器的推理方法，该方法可以尽可能多地重叠计算和通信，并进一步获得更有效的推理性能，而无需使用额外的机器用于更大规模的MoE模型。
• 在NLP和CV任务中采用了一些有效的训练方法来扩展多任务学习而无需额外的记忆并提高性能。这些方法包括负载平衡、嵌入分区和资源感知通信。

SE-MoE训练设计

为了用尽可能少的硬件资源训练具有相当多参数的模型，采用卸载策略并在设备中的存储资源的极端利用率下训练更大规模的模型是合适的解决方案。SE-MoE通过组合各种存储设备来存储稀疏和密集参数。同时，采用2D预取调度来减轻不同设备之间的移动所造成的训练性能的下降。

A. 分级存储

通常存储的参数状态由三部分组成：可训练参数、参数梯度和相应的优化器状态。
存储设备可分为三部分来存储参数状态：GPU节点、CPU节点和SSD节点。

• 由于dense参数密集用于计算，不占大部分存储空间，因此其参数状态都存储在GPU节点上，以避免频繁的数据移动。
• 相比之下，由于稀疏参数在训练中被选择性地激活并且与desnse参数相比占用大量存储空间，因此参数状态被放置在SSD节点上并且在适当的时间被传输到GPU节点以用于计算。

通过根据参数的计算和存储特点，将相应的参数状态合理地存储在分级存储器上，可以尽可能充分地利用设备的存储空间。

B. 2D预取调度策略

由于采用分级存储来保存稀疏参数状态和密集参数状态，模型训练期间不同设备之间传输参数状态的时间会显著增加，因此设计了一种2D预取调度策略来提高访问效率。如上图所示：

• 对于使用四分策略后的dense参数切片，通过NVLink的高速带宽，在水平维度上的行列之间进行通信后，可以预取完整的密集参数，实现数据并行的预期效果。
• 稀疏参数存储在SSD中，于是在CPU内存中建立了相应的缓存机制（比如LFU）来预存稀疏参数。当接收到预取请求时，优选的是从CPU高速缓存中检索所请求的稀疏参数，并且然后如果它们未存储在CPU高速缓存中则从SSD中检索它们。

C.融合传播

为了减少不连续的通信，对请求通信的参数采用了融合策略。通过参数管理单元，可以在通信前将参数片合并成较大的参数片，然后在通信后将其切割成相应的较小的参数片。

通信结束后，根据记录的切片索引重建整个参数状态。通过减少通信时间，减少的延迟开销可以更好地扩展到多种设备。

SE-MoE推理设计

有很多研究表明，MoE模型的训练效率明显高于密集模型。然而对于推理，与密集模型相比，大量参数（大多数是无效参数）引入了更大的存储负担。为了降低存储负担和提高推理效率，有两种方法很受欢迎：

• 知识蒸馏，可以减少模型大小并保持精度[1、2、3]。
• 并行化技术，包括专家切片、专家并行化、张量切片等。

当机器的数量或存储设备被限定时，这些方法在大参数量MoE模型下也将无计可施。于是文章提出了基于环形存储器的MoE推理架构。

如上图所示，整个推理过程分为六个步骤：

1. 图融合。对于超大规模分布式训练模型，采用源图合并和相应的分布式策略进行参数冗余消除。
2. 蒸馏和压缩。对教师网络中的众多专家进行蒸馏和压缩，得到专家数较少的学生网络。
3. 图形转换。将动态图转换为静态图，以便后续优化和部署。
4. 图形分割。根据推理资源和实际需求，选择合理的分布式策略，手动或自动地将静态图分解为多个分布式子图，并增加额外的通信。
5. 优化。对于分布式子图，采用核融合等相关的IR通道优化，进一步提高了推理性能。
6. 部署。将优化后的子图部署到服务器上提供服务。

环存储卸载

为了能够在有限资源的情况下进行大规模MoE模型的推理，有必要采用卸载策略来解决存储问题。然而，数据移动的速度不可避免地成为推理性能的瓶颈。因此，许多方法试图尽可能地覆盖推理计算背后的数据移动，从而可以减少计算等待时间。

上图展示了MoE模型的Transformer架构，每一层都是相互独立的，可以用于交错计算和卸载以实现重叠。假设在MoE推断模型中存在N个解码器层，将N个层的专家参数的N个副本存储在CPU设备中，并且将诸如嵌入的其他参数放置在GPU设备的密集缓冲器上。同时，GPU设备还缓存专家参数的K个副本。
如下图a所示，一旦与第i层相关的所有计算完成，则可以释放GPU存储器中对应的

      P 
     
    
      i 
     
    
   
  
    P_{i} 
   
  
Pi​参数，并开始从CPU存储器异步加载第（K+ i）层的 
 
  
   
    
    
      S 
     
     
     
       K 
      
     
       + 
      
     
       i 
      
     
    
   
  
    S_{K+i} 
   
  
SK+i​专家参数，以占用 
 
  
   
    
    
      P 
     
    
      i 
     
    
   
  
    P_{i} 
   
  
Pi​的空间。以这种方式，GPU设备上的专家参数的固定K个副本由计算释放负载来维护，并且它们被存储在环形存储器中以减轻存储器碎片。通过使用不同的CUDA流，来自CPU的专家加载和计算可以部分重叠，如图b所示。并且当MoE推断模型具有更多解码器层并且环存储器大小足够时，可以极大地最大化重叠。

环形内存卸载的调度和时间轴(a)。

1. 从SSD存储器中的文件加载N个参数副本
2. 从CPU存储器加载K个参数副本
3. 执行第i层计算
4. 释放第i个参数并触发异步复制以用 S K + i S_{K+i} SK+i​替换 P i P_{i} Pi​。

SE-MoE模型方法

MoE模型的独特架构在训练和推理中引入了新的固有问题。先将问题和解决措施汇总如下：

• 弹性MoE训练 → \rightarrow → 由于不均匀的输入数据而导致的负载不平衡。
• 资源感知通信 → \rightarrow → MoE有相当多的跨机器通信。
• 数据并行嵌入分区 → \rightarrow → 大多数任务中使用过大的词汇表而导致的存储限制。

A 弹性MoE训练

负载不均衡导致1.超过内存限制，因为单个任务节点不均匀的输入导致需要处理更大的批数据；
2.同步通信必须等待最慢的节点，这被称为“木桶效应”，并导致计算利用率下降。

弹性MoE训练事先统计估计的任务工作量，灵活调整训练节点的数量，保证各节点的负载均衡。
具体有两种做法：1.在存储压力小的情况下，将多个节点与轻型任务相结合（b）；2.添加额外的节点来处理繁重的任务（c）。

B 资源感知通信

在MoE模型的训练和推理期间，在专家并行性中的设备之间需要大量的All-to-All通信，All-to-All通信的多个进程将同时竞争有限的网络资源。并且由于更大的消息路径和流量成本，跨集群的数据交互比一个集群中的设备慢得多。因此需要提前规划资源安排路径，提高高速通道的利用率和减小交互竞争资源的潜在可能性。

下图就是一个资源感知通信后的路线优化

C 数据并行嵌入分区

在超大规模模型训练的实现中，嵌入表往往是整个模型参数中最大的参数，因此嵌入表的存储受到模型规模的限制。

假设我们在N个训练进程中划分维度为

     [ 
    
   
     V 
    
   
     ， 
    
   
     H 
    
   
     ] 
    
   
  
    [V，H] 
   
  
[V，H]的嵌入表，按行的方法将 
 
  
   
   
     [ 
    
    
    
      V 
     
    
      N 
     
    
   
     ， 
    
   
     H 
    
   
     ] 
    
   
  
    [ \frac{V}{N} ，H] 
   
  
[NV​，H]碎片分发给每个worker，这意味着每个进程只有部分词汇的嵌入表示。在查询嵌入表之前，每个进程的输入数据必须通过All-to-All通信彼此交换，以获得局部词汇表中的嵌入结果。之后，为了获得worker的输入数据的正确结果，通过All-to-All通信再次交换嵌入结果，该All-to-All通信可以被视为先前通信的逆过程。显然，对于后向阶段，需要将梯度交换到恢复嵌入表梯度。数据并行嵌入分区只引入了三个All-to-All通信和删除AllReduce同步嵌入表梯度数据并行。

实验结果

训练比较

与最先进的MoE系统DeepSpeed相比，SE-MoE在单节点训练中获得了近28%的加速，在具有超过1000亿个参数的MoE模型的多节点训练中获得了至少33%的加速。同时，SE-MoE将每个等级的GPU内存减少了近12 GB。

推理比较

与DeepSpeed相比，SE-MoE在具有超过2000亿个参数的MoE模型上获得了近13%的加速。

3.文章总结

SE-MoE训练和推理系统，可以很好地满足自然语言处理和CV任务的要求。未来将进一步探索一个统一的稀疏训练和推理系统，在多个维度考虑参数服务器和调度。统一的系统将有效地改善稀疏训练，以克服通信、计算和存储瓶颈。