Accelerating Large Language Model Training with Hybrid GPU-based Compression

📄 arXiv: 2409.02423v1 📥 PDF

作者: Lang Xu, Quentin Anthony, Qinghua Zhou, Nawras Alnaasan, Radha R. Gulhane, Aamir Shafi, Hari Subramoni, Dhabaleswar K. Panda

分类: cs.DC, cs.AI

发布日期: 2024-09-04

💡 一句话要点

提出基于GPU压缩的混合并行训练方法，加速大语言模型训练。

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

关键词: 大语言模型 分布式训练 数据并行 张量并行 流水线并行 GPU压缩 混合并行 MPI通信

📋 核心要点

  1. 现有大模型训练依赖数据密集型通信，梯度、激活等信息的传输开销巨大，成为性能瓶颈。
  2. 针对不同并行维度通信数据的稀疏性差异，提出混合压缩方案，差异化地应用压缩策略。
  3. 实验表明，该方法在保证模型收敛的前提下，显著提升了训练速度和GPU利用率。

📝 摘要（中文）

数据并行（DP）、张量并行（TP）和流水线并行（PP）是目前广泛采用的三种加速和高效训练大型语言模型（LLM）的策略。然而，这些方法依赖于数据密集型的通信例程来收集、聚合和重新分发梯度、激活和其他重要的模型信息，这带来了显著的开销。与基于GPU的压缩库共同设计的MPI库已被证明可以显著减少消息大小，并利用互连带宽，从而在保持可接受的精度的同时提高训练效率。\n\n在这项工作中，我们研究了在使用3D并行和ZeRO优化进行分布式LLM训练的背景下，压缩辅助MPI集合通信的有效性。我们在Lassen超级计算机上扩展到192个V100 GPU。首先，我们对所有集合通信启用了一种朴素的压缩方案，观察到GPT-NeoX-20B训练的每个GPU的TFLOPS增加了22.5%，每秒样本数增加了23.6%。然而，这种策略忽略了每个并行维度中通信的消息之间的稀疏性差异，从而引入了更多的误差，导致训练损失的下降。因此，我们针对每个并行维度采用了混合压缩设置，并相应地调整了压缩强度。鉴于其低秩结构，我们在执行DP All-reduce时对梯度应用激进的压缩。我们采用较温和的压缩来保持精度，同时在TP和PP中通信激活、优化器状态和模型参数。使用调整后的混合压缩方案，我们证明了在达到基线损失收敛的同时，每个GPU的TFLOPS增加了17.3%，每秒样本数增加了12.7%。

🔬 方法详解

问题定义：论文旨在解决大规模语言模型（LLM）训练中，由于数据并行、张量并行和流水线并行导致的大量数据通信所带来的性能瓶颈问题。现有方法在梯度、激活等信息的传输过程中，通信开销巨大，严重影响训练效率。

核心思路：论文的核心思路是利用GPU加速的压缩技术，减少MPI通信的数据量。更进一步，针对不同并行策略（DP、TP、PP）下通信数据的特性，采用混合压缩策略，即对不同类型的数据采用不同强度的压缩，以在保证模型收敛的前提下，最大化压缩收益。

技术框架：整体框架基于3D并行和ZeRO优化，在现有的LLM训练流程中，嵌入了压缩辅助的MPI集合通信。主要模块包括：数据并行（DP）梯度压缩模块、张量并行（TP）激活压缩模块、流水线并行（PP）参数压缩模块。每个模块根据数据的特性选择合适的压缩算法和压缩比。

关键创新：关键创新在于混合压缩策略。不同于以往对所有通信数据采用统一压缩策略，该方法根据不同并行维度通信数据的稀疏性差异，自适应地调整压缩强度。例如，对DP中的梯度进行激进压缩，而对TP和PP中的激活和参数进行温和压缩，从而在保证模型精度和收敛性的前提下，最大化压缩带来的性能提升。

关键设计：论文的关键设计包括：1) 针对DP梯度，利用其低秩特性，采用高压缩比的压缩算法；2) 针对TP和PP中的激活和参数，采用低压缩比的压缩算法，以保证精度；3) 针对不同的压缩算法，选择合适的GPU加速库，以提高压缩和解压缩的效率；4) 通过实验确定不同并行维度下最佳的压缩比。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在GPT-NeoX-20B模型的训练中，采用朴素压缩方案后，每个GPU的TFLOPS提升了22.5%，每秒样本数提升了23.6%。而采用混合压缩方案后，在保证模型收敛的前提下，每个GPU的TFLOPS提升了17.3%，每秒样本数提升了12.7%。证明了混合压缩方案的有效性。

🎯 应用场景

该研究成果可广泛应用于各种大规模语言模型的训练场景，尤其是在资源受限的条件下，可以有效提高训练效率，降低训练成本。此外，该方法还可以推广到其他深度学习模型的训练中，具有重要的实际应用价值和广阔的应用前景。

📄 摘要（原文）

Data Parallelism (DP), Tensor Parallelism (TP), and Pipeline Parallelism (PP) are the three strategies widely adopted to enable fast and efficient Large Language Model (LLM) training. However, these approaches rely on data-intensive communication routines to collect, aggregate, and re-distribute gradients, activations, and other important model information, which pose significant overhead. Co-designed with GPU-based compression libraries, MPI libraries have been proven to reduce message size significantly, and leverage interconnect bandwidth, thus increasing training efficiency while maintaining acceptable accuracy. In this work, we investigate the efficacy of compression-assisted MPI collectives under the context of distributed LLM training using 3D parallelism and ZeRO optimizations. We scaled up to 192 V100 GPUs on the Lassen supercomputer. First, we enabled a naïve compression scheme across all collectives and observed a 22.5\% increase in TFLOPS per GPU and a 23.6\% increase in samples per second for GPT-NeoX-20B training. Nonetheless, such a strategy ignores the sparsity discrepancy among messages communicated in each parallelism degree, thus introducing more errors and causing degradation in training loss. Therefore, we incorporated hybrid compression settings toward each parallel dimension and adjusted the compression intensity accordingly. Given their low-rank structure (arXiv:2301.02654), we apply aggressive compression on gradients when performing DP All-reduce. We adopt milder compression to preserve precision while communicating activations, optimizer states, and model parameters in TP and PP. Using the adjusted hybrid compression scheme, we demonstrate a 17.3\% increase in TFLOPS per GPU and a 12.7\% increase in samples per second while reaching baseline loss convergence.