一种新的分布式训练库实现了从数十个到数百个GPU扩展中近线性的效率。
顶尖的语言模型拥有数十亿个参数。为了在可控时间内训练这些模型，需要将工作负载分布到大型计算集群上。理想情况下，训练时间会随着集群规模的扩大而线性减少。然而，由于协调集群节点工作所需的通信会抵消并行化带来的增益，线性扩展很难实现。
最近，研究团队投入精力优化了某机构DeepSpeed分布式训练库的通信效率，显著提升了其在多达64个GPU上的性能。但是，当从数十个GPU扩展到数百个GPU时，在公有云环境中，通信开销再次开始压过效率增益。
在一篇将于2023年在某国际会议上展示的论文中，研究者提出了一种方法，使模型训练能够在云端的数百个GPU上高效扩展。该方法名为MiCS，其核心思想是通过最小化通信规模来降低通信开销。
具体来说，现有的分布式训练框架（如DeepSpeed和FairScale）会将模型状态划分到所有GPU上，而MiCS则会创建多个模型状态的副本，并将每个副本划分到一个GPU子集内。根据模型大小，一个副本可以放在单个计算节点（即一台机器，其GPU间具有高速连接）上，也可以跨多个节点放置。
因此，在MiCS中，频繁的通信操作（如参数收集）被限制在GPU子集内。这样一来，当通过跨新节点添加新副本来扩展集群时，这些频繁通信操作的延迟保持固定不变，而不是随着集群规模的增长而增加。同时，MiCS还减少了在单个节点无法容纳模型状态副本时，节点之间传输的数据量。最后，MiCS包含一个梯度同步调度机制，可以在所有工作节点间分摊昂贵的梯度同步成本。
实验结果表明，在不同规模的BERT模型上（使用某实例集群评估），MiCS在吞吐量和扩展效率方面均有显著提升。MiCS能够实现近乎线性的可扩展性，与DeepSpeed-v0.5.6内置的通信管理方法ZeRO的第二、三阶段相比，吞吐量最高可提升2.82倍。研究者还将MiCS与先前对ZeRO第三阶段的优化进行了比较，即使在之前研究过的较低GPU数量下，也证明了其改进之处。
某计算实例为高性能计算提供高达400Gbps的网络带宽。然而，由于通信开销（尤其是延迟，当集群中添加更多GPU时延迟会增加），分布式系统可能无法高效地充分利用400Gbps。
研究者已部署MiCS在某实例上训练高达1750亿参数的专有模型。在16台某实例上训练一个1750亿参数、序列长度为2048的模型时，每个GPU实现了169 teraflops的性能（达到理论峰值的54.2%）。当在64台某实例上训练一个1000亿参数的模型时，MiCS使每个GPU保持了超过170 teraflops的性能。当集群规模从128个GPU扩展到512个GPU时，MiCS实现了99.4%的线性扩展效率。相比之下，DeepSpeed ZeRO的第三阶段仅达到72%的扩展效率，且每个GPU的性能饱和在62 teraflops。
可感知规模的模型划分
默认情况下，DeepSpeed将所有设备上的模型状态进行分区，这种策略降低了集群中每个GPU的内存消耗，但会在训练中产生巨大的通信开销。更重要的是，这种开销会随着集群规模的增长而增加，导致大规模下的可扩展性显著下降。
MiCS没有将模型状态划分到所有GPU，而是将集群中的GPU分成多个分区组，并在每个组内划分模型状态。每个分区组持有模型状态的一个完整副本。持有相同部分模型状态的GPU则构成另一种组，称为复制组。在每个分区组内划分模型状态，将最频繁的通信（参数收集和梯度同步）限制在固定数量的GPU内，从而有效控制通信开销，使其不随集群规模增长。
分级通信策略
当单个模型状态副本所需内存超过单节点GPU内存总量时，需要将副本存储在跨多个节点的GPU上，此时必须依赖效率较低的节点间通信。
集体通信中传输的数据量和延迟取决于消息大小和参与节点数。通信量与 (p - 1)/p 成正比（p为参与节点数），如果采用标准的环形通信模式，延迟与参与节点数呈线性关系。在不影响数据完整性的前提下无法减小消息大小，但可以减少节点间通信的参与节点数。
考虑一个简单的例子，涉及两个节点，每个节点有两个GPU。标准的环形通信模式会将数据跨节点聚合，消息依次在每个GPU间传递，因此一次节点间通信涉及四个GPU。相比之下，MiCS并行执行这些节点间操作，使得每次节点间通信只涉及两个GPU，它们只交换所需通信信息的一半。然后，每个节点在本地聚合节点间数据以组装完整消息。在这种情况下，通信量因子从 3/4 降低到 1/2。
两跳梯度同步
在所有工作节点间同步梯度是一项昂贵的操作，但这是保持工作节点在相同模型状态下工作的必要条件。在大规模神经网络的训练中，批次大小通常受限于GPU内存。梯度累积是一种将一批样本分割成多个微批次（microbatches），并按顺序在多个微步（microsteps）中运行的技术。
利用MiCS，可以在每个分区组内，在多个微批次中累积梯度，直到最后一个微批次被处理完毕。也就是说，对于每个微步，可以在一个GPU子集（即一个分区组）内累积每个模型副本的完整梯度集。然后，在最后一个微批次处理完成后，每个GPU与代表相同模型状态的其他GPU同步梯度。这允许将跨复制组的同步开销分摊到多个微步中。
得益于此三种技术，MiCS在大型集群上展现出卓越的可扩展性，提供了优秀的训练吞吐性能，并使其能够在某实例上达到新的性能水平。研究者正致力于将MiCS开源，相信它将大大降低在特定云平台上训练大型模型的时间和成本。关于该系统更详细的解释和性能分析，可参考其预印本论文。FINISHED
