Pipette: Automatic Fine-grained Large Language Model Training Configurator for Real-World Clusters

📄 arXiv: 2405.18093v1 📥 PDF

作者: Jinkyu Yim, Jaeyong Song, Yerim Choi, Jaebeen Lee, Jaewon Jung, Hongsun Jang, Jinho Lee

分类: cs.DC, cs.LG

发布日期: 2024-05-28

备注: published at DATE 2024

🔗 代码/项目: GITHUB

💡 一句话要点

Pipette：面向真实集群的自动细粒度大语言模型训练配置器

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

关键词: 大语言模型训练 3D并行 集群优化 性能建模 内存估计

📋 核心要点

  1. 现有方法忽略了真实集群中异构互连速度的差异，导致次优的并行配置方案。
  2. Pipette通过更精确的性能建模、内存估计和细粒度的GPU分配，自动寻找更优的LLM训练配置。
  3. 实验表明，Pipette在大型集群上能够显著加速LLM的训练过程，优于现有技术。

📝 摘要（中文）

训练大型语言模型（LLM）极具挑战，因为它需要巨大的计算和内存容量。为了解决这些问题，通常使用具有3D并行性的GPU集群，该并行性沿数据批次、流水线阶段和层内张量维度分割模型。然而，3D并行性的使用带来了额外的挑战，即找到每个维度上的最佳分割方式并将分割后的模型映射到GPU上。先前的研究试图自动找到最佳配置，但许多研究缺乏几个重要方面。例如，异构互连速度的性质通常被忽略。虽然互连的峰值带宽通常被设置为相等，但实际达到的带宽在真实集群中每个链路都不同。加上没有正确考虑通信的关键路径建模，它们很容易陷入次优配置。此外，它们常常未能考虑每个GPU的内存需求，经常推荐无法执行的解决方案。为了应对这些挑战，我们提出了Pipette，这是一种用于真实集群的自动细粒度LLM训练配置器。通过设计更好的性能模型以及内存估计器和细粒度的单个GPU分配，Pipette实现了满足内存约束的更快配置。我们在大型集群上评估了Pipette，表明它比现有技术提供了显着的加速。Pipette的实现可在https://github.com/yimjinkyu1/date2024_pipette 获得。

🔬 方法详解

问题定义：现有的大语言模型训练通常采用3D并行策略，但如何针对真实集群的异构互连环境，自动找到最优的并行配置（数据并行、流水线并行、张量并行）是一个难题。现有方法通常忽略互连速度的异构性，并且关键路径建模不准确，导致配置方案次优，甚至超出GPU内存限制。

核心思路：Pipette的核心思路是构建更精确的性能模型，该模型能够考虑到真实集群中不同GPU之间互连带宽的差异。同时，Pipette还包含一个内存估计器，用于确保生成的配置方案满足每个GPU的内存约束。通过细粒度的GPU分配，Pipette能够充分利用集群资源，找到更优的训练配置。

技术框架：Pipette的整体框架包含以下几个主要模块：1) 性能建模模块：用于预测不同并行配置下的训练速度，考虑了异构互连带宽的影响。2) 内存估计模块：用于估计每个GPU的内存使用情况，确保配置方案满足内存约束。3) 细粒度GPU分配模块：用于将模型的不同部分分配到不同的GPU上，充分利用集群资源。4) 优化搜索模块：用于搜索最优的并行配置方案，目标是最小化训练时间，同时满足内存约束。

关键创新：Pipette的关键创新在于其更精确的性能模型，该模型能够考虑到真实集群中异构互连带宽的差异。与现有方法相比，Pipette的性能模型更加准确，能够更好地预测不同并行配置下的训练速度。此外，Pipette还引入了细粒度的GPU分配策略，能够更灵活地利用集群资源。

关键设计：Pipette的性能模型基于对通信开销的精确建模，考虑了不同GPU之间互连带宽的差异。内存估计模块则通过分析模型的结构和数据依赖关系，预测每个GPU的内存使用情况。优化搜索模块采用了一种基于约束优化的算法，用于搜索最优的并行配置方案。具体的参数设置和损失函数等技术细节在论文中进行了详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Pipette在大型GPU集群上的实验结果表明，与现有技术相比，Pipette能够显著加速LLM的训练过程。具体的性能提升幅度取决于集群的规模和模型的复杂度，但总体而言，Pipette能够提供显著的加速效果，并确保生成的配置方案满足内存约束。

🎯 应用场景

Pipette可应用于各种需要大规模LLM训练的场景，例如自然语言处理、机器翻译、文本生成等。通过自动优化训练配置，Pipette可以显著降低训练成本，缩短开发周期，加速LLM在实际应用中的部署。该研究对于推动LLM的普及和应用具有重要意义。

📄 摘要（原文）

Training large language models (LLMs) is known to be challenging because of the huge computational and memory capacity requirements. To address these issues, it is common to use a cluster of GPUs with 3D parallelism, which splits a model along the data batch, pipeline stage, and intra-layer tensor dimensions. However, the use of 3D parallelism produces the additional challenge of finding the optimal number of ways on each dimension and mapping the split models onto the GPUs. Several previous studies have attempted to automatically find the optimal configuration, but many of these lacked several important aspects. For instance, the heterogeneous nature of the interconnect speeds is often ignored. While the peak bandwidths for the interconnects are usually made equal, the actual attained bandwidth varies per link in real-world clusters. Combined with the critical path modeling that does not properly consider the communication, they easily fall into sub-optimal configurations. In addition, they often fail to consider the memory requirement per GPU, often recommending solutions that could not be executed. To address these challenges, we propose Pipette, which is an automatic fine-grained LLM training configurator for real-world clusters. By devising better performance models along with the memory estimator and fine-grained individual GPU assignment, Pipette achieves faster configurations that satisfy the memory constraints. We evaluated Pipette on large clusters to show that it provides a significant speedup over the prior art. The implementation of Pipette is available at https://github.com/yimjinkyu1/date2024_pipette.