状态最先进的语言模型拥有数十亿参数。在可控时间内训练这些模型需要将工作负载分配到大型计算集群上。理想情况下，训练时间应随着集群规模的扩大而线性减少。然而，线性扩展很难实现，因为协调集群节点工作所需的通信消耗了并行化带来的增益。
最近，我们在优化某机构DeepSpeed分布式训练库的通信效率方面付出了一些努力，显著提升了最多64个GPU的性能。然而，当我们在公有云环境中从几十个GPU扩展到数百个时，通信开销再次开始压倒效率增益。
在即将于2023年国际超大型数据库会议上发表的一篇论文中，我们提出了一种方法，使模型训练能在云中数百个GPU上高效扩展。我们将此方法称为MiCS，因为它通过最小化通信规模来降低通信开销。
具体来说，现有的分布式训练框架（如DeepSpeed和FairScale）将模型状态划分到所有GPU上，而MiCS则创建多个模型状态的副本，并将每个副本划分到一个GPU子集内。根据模型大小，一个副本可能适合单个计算节点（一台具有GPU间高速连接的机器），也可能适合多个节点。
因此，在MiCS中，频繁的通信操作（如参数收集）被限制在一个GPU子集内。这样，当我们通过在新节点上添加新副本来扩展集群时，频繁通信操作的通信延迟保持固定，而不是随着集群大小而增长。
我们还减少了在模型状态副本无法容纳于单个节点时节点间传输的数据量。最后，MiCS包含一个梯度同步调度，可在所有工作节点之间分摊昂贵的梯度同步成本。
实验结果显示，在由p3dn.24xlarge实例组成的集群上评估不同大小的BERT模型时，吞吐量和扩展效率均有显著提升。MiCS能够实现近乎线性的可扩展性（由下图中的矩形框表示），并且与DeepSpeed-v0.5.6内置的通信管理方法——三阶段零冗余优化器的第二和第三阶段相比，吞吐量最高提升了2.82倍。
我们还比较了MiCS与之前对ZeRO第三阶段的优化（见下图），结果表明即使在我们之前研究的较低GPU数量下，MiCS也能带来改进。我们在arXiv的预印本论文中更详细地报告了所有这些发现。
AWS P4d提供高达400Gbps的网络带宽用于高性能计算。不幸的是，分布式系统可能无法高效地充分利用400Gbps，因为通信开销——尤其是延迟——会随着向集群添加更多GPU而增加。
我们已经部署MiCS在p4d.24xlarge和p4de.24xlarge实例上训练参数高达1750亿的专有模型。当在16个p4de.24xlarge实例上训练一个序列长度为2048的1750亿参数模型时，我们能够在每个GPU上实现169万亿次浮点运算的性能。当我们在64个p4d.24xlarge实例上训练一个1000亿参数模型时，MiCS在每个GPU上维持超过170万亿次浮点运算。
当集群规模从128个GPU扩展到512个GPU时，MiCS实现了线性扩展效率的99.4%。相比之下，DeepSpeed ZeRO的第三阶段仅实现了72%的弱扩展效率，并在每个GPU上达到62万亿次浮点运算后饱和。
规模感知的模型分区
默认情况下，DeepSpeed将模型状态分区到所有设备上，这种策略降低了集群中每个GPU的内存消耗，但在训练中产生了大量的通信开销。更重要的是，开销随集群规模而增加，这导致大规模下的可扩展性显著下降。
MiCS没有将模型状态分区到所有GPU，而是将集群中的GPU划分为多个组，并在每个组内分区模型状态。我们称这些组为分区组。每个组持有模型状态的完整副本。下图给出了一个由两个连续GPU组成的分区组的例子。持有相同部分模型状态的GPU形成另一种组，即复制组。
在分区组内划分模型状态，将最频繁的通信——参数收集和梯度同步——限制在固定数量的GPU内。此策略有效控制了通信开销，使其不随集群规模增长。
分层通信策略
当单个模型状态副本的内存需求大于单个节点内GPU内存总量时，我们需要将副本存储在跨越多个节点的GPU上。在这种情况下，我们必须依赖效率较低的节点间通信。
传输数据的体积和集体通信中的延迟由消息大小和参与者数量决定。特别是，通信量与(p-1)/p成正比，其中p表示参与者数量，如果参与者使用标准的环形通信模式，延迟与参与者数量呈线性依赖关系。
消息大小在不损害数据完整性的情况下无法减少，但我们可以减少节点间通信中的参与者数量。这将通信量因子降低到(p-k)/p，延迟降低p/(p/k + k)倍，其中k是单个节点上的GPU数量。
考虑下面一个简单的例子，涉及两个节点，每个节点有两个GPU。标准的环形通信模式会通过将消息从每个GPU传递到下一个来跨节点聚合数据，因此单个节点间通信涉及四个GPU。
相比之下，MiCS并行执行这些节点间操作，因此每个节点间通信只涉及两个GPU，它们只交换我们想要通信的一半信息。然后每个节点在本地聚合节点间数据以组装完整消息。在这种情况下，通信量因子从¾降低到½。
两跳梯度同步
在所有工作节点之间同步梯度是一项昂贵的操作，是为了保持工作节点处理相同的模型状态所必需的。在大型神经网络的训练期间，批量大小通常受GPU内存限制。梯度累积是一种技术，它将一批样本分割成几个微批次，这些微批次将在多个微步中顺序运行。
使用MiCS，我们可以在分区组内通过多个微批次累积梯度，直到处理完最后一个微批次。也就是说，对于每个微步，我们可以在一个GPU子集内累积每个模型副本的完整梯度集。然后，在最后一个微批次处理完后，每个GPU与代表相同模型状态部分的其他GPU同步梯度。
这使我们能够将复制组之间的同步开销分摊到多个微步上。下图给出了一个包含四个微步的训练中两跳梯度同步的例子。
由于这三种技术，MiCS在大型集群上显示出极佳的可扩展性，并提供出色的训练吞吐性能，使我们能够在AWS p4de.24xlarge机器上实现新的最先进性能。我们正在努力将MiCS开源供公众使用，相信它将大大减少在Amazon EC2平台上进行大型模型训练的时间和成本。请参阅我们的预印本以获取对我们系统及其性能分析的更详细解释。
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）
