亚马逊近线性大规模模型训练加速库MiCS来了！

最先进的语言模型有数十亿个参数。在可管理的时间内训练这些模型需要在大型计算群集上分配工作负载。理想情况下，训练时间将随着群集大小的扩展而线性减少。然而，线性扩展很难实现，因为协调群集节点工作所需的通信会侵蚀并行化的收益。

最近，AWS在优化Microsoft DeepSpeed分布式训练库的通信效率方面付出了一些努力，显著提高了多达64个GPU的性能。然而，在公共云环境中，当我们从数十个GPU扩展到数百个GPU时，通信开销再次开始压倒效率收益。

为了在云中数百个GPU上有效地扩展模型训练，我们提出了MiCS，一种最大限度地降低通信规模以降低通信开销的方法。

具体来说，在现有的分布式训练框架（如DeepSpeed和FairScale）在所有GPU上划分模型状态的情况下，MiCS会制作模型状态的多个副本，并在GPU的子集中划分每个副本。根据模型大小，副本可以安装在单个计算节点上--GPU之间具有高速连接的单个计算机--或多个节点上。

因此，在MiCS中，频繁的通信操作，如参数收集，被限制在GPU的子集上。这样，当我们通过跨新节点添加新副本来扩展群集时，频繁通信操作的通信延迟保持固定，而不是随着群集的大小而增长。

如果模型状态的副本不适合单个节点，我们还减少了节点之间传输的数据量。最后，MiCS包括一个梯度同步计划，该计划在所有工作线程之间摊销昂贵的梯度同步。

我们的实验结果表明，在由p3dn.24x大实例组成的集群上评估的不同大小BERT模型的吞吐量和扩展效率显著提高。MiCS能够实现近线性可扩展性（下图中的矩形帧表示），并提供高达2.82倍的吞吐量，与内置在DeepSpeed-v0.5.6中的通信管理方法ZeRO的三级零冗余优化器（Zero）的第二和第三状态相比。