分布式训练库MiCS实现近线性扩展

近线性扩展巨型模型训练

最先进的语言模型拥有数十亿参数。为了在可控时间内训练这些模型，需要将工作负载分布到大型计算集群中。理想情况下，训练时间应随集群规模线性减少。但由于协调集群节点工作所需的通信会抵消并行化带来的收益，实现线性扩展十分困难。

与传统分布式训练框架（如DeepSpeed和FairScale）将所有GPU分区模型状态不同，MiCS创建模型状态的多个副本，并在GPU子集内分区每个副本。根据模型大小，副本可适配单个计算节点（具有GPU间高速连接的机器）或多个节点。

在MiCS中，频繁的通信操作（如参数收集）被限制在GPU子集内。当通过新节点添加新副本来扩展集群时，频繁通信操作的延迟保持固定，而不随集群规模增长。

MiCS还减少了节点间传输的数据量，并采用梯度同步计划将所有工作节点的昂贵梯度同步分摊。

实验结果显示，在不同规模BERT模型上，使用p3dn.24xlarge实例集群时，吞吐量和扩展效率均有显著提升。MiCS实现近线性可扩展性（如下图矩形框所示），相比DeepSpeed-v0.5.6内置的ZeRO三阶段优化器，吞吐量提升达2.82倍。

即使在我们先前研究的较低GPU数量下，MiCS也显示出优于早期ZeRO第三阶段优化的性能。

当单个模型状态副本的内存需求超过单节点GPU内存总量时，需要在跨多个节点的GPU上存储副本。这种情况下，必须依赖效率较低的节点间通信。

MiCS通过减少节点间通信的参与GPU数量来优化通信。例如标准环状通信模式涉及四个GPU的节点间通信，而MiCS通过并行执行这些操作，使每个节点间通信仅涉及两个GPU，交换仅一半的信息，随后在本地聚合数据。

MiCS在分区组内多个微批次中累积梯度，直到处理完最后一个微批次。每个微步骤中，可在GPU子集（即分区组）内累积每个模型副本的完整梯度集。处理完最后微批次后，每个GPU与代表模型状态相同部分的其他GPU同步梯度。

这种设计将复制组的同步开销分摊到多个微步骤中，显著提升大规模集群的可扩展性和训练吞吐量性能。