[论文解读] DiLoCo: Distributed Low-Communication Training of Language Models

Arthur Douillard, Qixuan Feng+7 more|arXiv (Cornell University)|Nov 14, 2023

Topic Modeling被引用 7

一句话总结

DiLoCo 提出了一种用于大规模语言模型的分布式训练框架，通过使用较大的内部优化步长（500步）结合 AdamW 优化器与外部 Nesterov 动量更新，显著减少了通信开销。在 C4 数据集上，其性能与完全同步训练相当，但通信数据量减少了 500 倍，且对数据分布、工作节点可用性变化以及异构设备表现出良好的鲁棒性。

ABSTRACT

Large language models (LLM) have become a critical component in many applications of machine learning. However, standard approaches to training LLM require a large number of tightly interconnected accelerators, with devices exchanging gradients and other intermediate states at each optimization step. While it is difficult to build and maintain a single computing cluster hosting many accelerators, it might be easier to find several computing clusters each hosting a smaller number of devices. In this work, we propose a distributed optimization algorithm, Distributed Low-Communication (DiLoCo), that enables training of language models on islands of devices that are poorly connected. The approach is a variant of federated averaging, where the number of inner steps is large, the inner optimizer is AdamW, and the outer optimizer is Nesterov momentum. On the widely used C4 dataset, we show that DiLoCo on 8 workers performs as well as fully synchronous optimization while communicating 500 times less. DiLoCo exhibits great robustness to the data distribution of each worker. It is also robust to resources becoming unavailable over time, and vice versa, it can seamlessly leverage resources that become available during training.

研究动机与目标

• 解决加速器未共置且通信带宽受限时训练大规模语言模型的挑战。
• 克服完全同步分布式训练所面临的可扩展性与可靠性问题，后者需要高度同步、高带宽的集群环境。
• 实现在多个隔离设备集群（‘岛屿’）之间进行训练，这些集群之间互联性差。
• 提升对数据分布变化、工作节点故障以及计算资源动态可用性的鲁棒性。
• 在最小化通信开销的同时，保持与完全同步训练相当的模型性能与效率。

提出的方法

• 采用受联邦平均启发的框架，使用较大的内部训练步长（H=500），以减少通信频率。
• 在每个工作节点的本地模型更新中使用 AdamW 作为内部优化器，与标准大规模语言模型训练实践保持一致。
• 在聚合所有工作节点的梯度后，使用 Nesterov 动量作为外部优化器来更新全局模型参数。
• 将模型更新解耦为局部训练阶段与低频的全局同步，将节点间通信频率降低至每 500 步一次。
• 允许工作节点在不同硬件类型和地理位置上运行，仅需周期性地同步参数。
• 在每次外部优化步骤后，将更新后的全局模型重新分发给所有工作节点，以实现迭代优化。

实验结果

研究问题

• RQ1 一种通信频率极低的分布式训练方法，是否能在大规模语言模型训练中实现与完全同步训练相当的性能？
• RQ2 该方法对不同工作节点数据分片之间的数据分布差异是否具有鲁棒性？
• RQ3 该方法能否有效应对训练过程中工作节点可用性动态变化的情况，如节点故障或新资源加入？
• RQ4 该方法在超过 8 个工作节点时是否仍能有效扩展？通信减少的极限在哪里？
• RQ5 当集群中使用非均匀或异构设备时，该方法是否能保持训练的稳定性和收敛性？

主要发现

• 在 C4 数据集上，DiLoCo 尽管通信数据量仅为完全同步训练的 1/500，其困惑度（perplexity）性能仍与之相当或更优。
• 该方法对工作节点分片间多样的数据分布具有鲁棒性，能保持稳定训练与收敛。
• DiLoCo 能够优雅地处理动态工作节点可用性变化，无缝集成新工作节点并从故障中恢复。
• 该方法支持在地理上分布的集群间进行训练，低带宽连接下通信频率从每几百毫秒一次降低至每几分钟一次。
• DiLoCo 保持了与完全同步模型相同的模型效率与推理速度，推理阶段无性能损失。
• 尽管 DiLoCo 在超过 8 个工作节点后收益递减，但其在异构、松散连接环境中展现出出色的训练速度效率与鲁棒性。