Enhancing Stability for Large Language Models Training in Constrained Bandwidth Networks

📄 arXiv: 2407.01614v3 📥 PDF

作者: Yun Dai, Tejas Dharamsi, Byron Hsu, Tao Song, Hamed Firooz

分类: cs.LG, cs.AI

发布日期: 2024-06-28 (更新: 2024-10-06)

💡 一句话要点

提出改进的ZeRO++算法，解决低带宽网络下大语言模型训练的收敛性问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型训练 分布式训练 ZeRO++ 低带宽网络 收敛性 数据并行 参数分区

📋 核心要点

  1. 现有ZeRO++在低带宽网络下训练大模型时，分层分区方案易产生竞争条件，导致收敛不稳定。
  2. 论文提出改进的分区算法，旨在解决ZeRO++的收敛性问题，同时保持训练效率。
  3. 实验表明，改进算法在Falcon和Llama-2模型上实现了可靠收敛，且吞吐量和训练速度提升98%。

📝 摘要（中文）

训练具有数十亿参数的超大型语言模型（LLM）是一项计算密集型任务，对当前的数据并行训练系统提出了严峻挑战。ZeRO++等技术使得在低带宽集群上高效地分布式训练此类巨型模型成为可能，但由于其分层分区（hpZ）方案中潜在的竞争条件，可能导致收敛问题，从而影响机器间的通信。本文首先展示了这些竞争条件如何在训练具有数十亿参数的模型时导致不稳定。然后，提出了一种对分区算法的修改，解决了这些收敛性挑战，同时保持了具有竞争力的训练效率。在训练数十亿参数的Falcon模型和Llama-2模型上的经验评估表明，更新后的算法能够可靠地收敛这些大型模型，而标准的ZeRO++ hpZ无法收敛。更新后的算法能够以98%的吞吐量和模型训练速度提升实现更大的模型的稳健训练，而不会牺牲收敛质量。

🔬 方法详解

问题定义：论文旨在解决在低带宽网络环境下，使用ZeRO++等数据并行技术训练超大型语言模型时遇到的收敛性问题。现有的ZeRO++分层分区（hpZ）方案在减少跨机器通信的同时，引入了潜在的竞争条件，导致训练过程不稳定，难以收敛。

核心思路：论文的核心思路是通过修改ZeRO++的分区算法，消除或减轻这些竞争条件，从而提高训练的稳定性。具体来说，通过调整参数的划分和通信方式，避免多个进程同时竞争同一资源，从而减少训练过程中的冲突。

技术框架：论文提出的方法是对现有ZeRO++框架的改进，主要集中在参数分区的算法上。整体训练流程仍然是标准的数据并行训练，包括数据划分、模型复制、梯度计算、参数更新等步骤。改进后的分区算法替换了原有的hpZ方案，并与现有的ZeRO++的其他优化技术兼容。

关键创新：论文的关键创新在于对ZeRO++分层分区算法的改进，通过重新设计参数划分策略，减少了进程间的竞争，从而提高了训练的稳定性。与原始的hpZ方案相比，改进后的算法能够更好地适应低带宽网络环境，并支持更大规模模型的训练。

关键设计：论文中关于参数设置、损失函数和网络结构等技术细节没有特别强调，重点在于分区算法的改进。具体的改进细节（如具体的参数划分策略、通信方式的调整等）在论文中可能有所描述，但摘要中未提及。需要查阅原文才能了解更详细的技术细节。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，改进后的算法在训练数十亿参数的Falcon模型和Llama-2模型时，能够实现可靠的收敛，而标准的ZeRO++ hpZ算法无法收敛。同时，改进后的算法在保持收敛质量的前提下，实现了98%的吞吐量和模型训练速度提升。这些结果验证了该方法在低带宽网络环境下训练大模型的有效性。

🎯 应用场景

该研究成果可应用于各种需要分布式训练超大型语言模型的场景，尤其是在计算资源有限、网络带宽受限的环境下。例如，可以使用廉价的GPU集群训练更大规模的模型，从而提升模型的性能和泛化能力。此外，该方法还可以降低训练成本，加速大语言模型在各个领域的应用。

📄 摘要（原文）

Training extremely large language models (LLMs) with billions of parameters is a computationally intensive task that pushes the limits of current data parallel training systems. While techniques like ZeRO++ have enabled efficient distributed training of such giant models on inexpensive low-bandwidth clusters, they can suffer from convergence issues due to potential race conditions in the hierarchical partitioning (hpZ) scheme employed to reduce cross-machine communication. In this work, we first show how these race conditions cause instability when training models with billions of parameters. We then propose a modification to the partitioning algorithm that addresses these convergence challenges while maintaining competitive training efficiency. Empirical evaluation on training the multi-billion parameters Falcon Models and Llama-2 models demonstrates the updated algorithm's ability to achieve reliable convergence on these massive models, where stock ZeRO++ hpZ fails to converge. The updated algorithm enables robust training of larger models with 98\% throughput and model training speed improvement without sacrificing the quality of convergence.