AGoQ通过量化激活和梯度，实现大模型训练内存降低52%且提速34%。
AGoQ: Activation and Gradient Quantization for Memory-Efficient Distributed Training of LLMs
- 按层分配比特位数，实现接近4比特的激活存储
- 8比特梯度存储与通信，内存减少52%，训练速度提升1.34倍
- 适合大规模分布式训练场景，兼容主流LLaMA架构
量化是降低大语言模型（LLM）训练时GPU内存需求的关键方法。然而，现有方法在4比特激活和8比特梯度场景下效果不佳，易导致收敛缓慢或精度下降。为此，我们提出AGoQ，包含两项新技术：1）一种分层感知的激活量化算法，根据层类型和流水线阶段动态分配比特位数，实现接近4比特的激活存储；2）一种梯度量化算法，通过8比特梯度存储和保持精度的8比特全部归约通信，降低内存占用并缩短通信时间。我们在两个GPU集群（最多64块显卡）上对不同规模的LLM进行大量实验，结果表明，AGoQ在80亿至320亿参数的LLaMA模型上，相比最先进的训练系统Megatron-LM（含或不含ZeRO）、COAT和DeepSpeed，内存最高降低52%，训练速度最快提升1.34倍，同时在预训练中无收敛损失，在下游任务上精度相当。
原文摘要 · Abstract (English)
Quantization is a key method for reducing the GPU memory requirement of training large language models (LLMs). Yet, current approaches are ineffective for 4-bit activations and 8-bit gradients, which would easily cause slow convergence or accuracy loss. To address this, we introduce AGoQ, incorporating two new techniques: 1) a layer-aware activation quantization algorithm that allocates appropriate bit-widths for activations of various layers based on their types and pipeline stages to achieve near 4-bit activation storage, and 2) a gradient quantization algorithm that reduces memory usage and shortens communication time by employing 8-bit gradient storage and precision-preserving 8-bit All-Reduce communication. We conduct extensive experiments using different sizes of LLMs on two GPU clusters (up to 64 GPUs), and the experimental results show that our AGoQ reduces the memory by up to 52\% and achieves up to 1.34$\times$ improvement of training speed compared to state-of-the-art training systems Megatron-LM (w/ or w/o ZeRO), COAT and DeepSpeed with 8B to 32B LLaMA models, while achieving convergence loss on pretraining and comparable accuracy on downstream tasks with LLaMA architectures.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。