arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2409.02423cs.DCcs.AI
收藏

利用基于GPU的混合压缩加速大型语言模型训练

Accelerating Large Language Model Training with Hybrid GPU-based Compression

The Ohio State University(俄亥俄州立大学)

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

AI总结:

本文针对大规模LLM分布式训练中的通信瓶颈,提出按并行维度差异化设置压缩强度的混合GPU压缩方案,在192块V100 GPU上实现GPT-NeoX-20B训练吞吐量提升17.3%且保持损失收敛。

摘要:

数据并行(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%。然而,这种策略忽略了每个并行维度中通信消息之间的稀疏性差异,从而引入了更多误差并导致训练损失退化。因此,我们针对每个并行维度采用混合压缩设置,并相应调整压缩强度。鉴于梯度的低秩结构(arXiv:2301.02654),我们在执行DP All-reduce时对梯度应用激进压缩。在TP和PP中通信激活值、优化器状态和模型参数时,我们采用较温和的压缩以保持精度。使用调整后的混合压缩方案,我们展示了每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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑