AGoQ: Activation and Gradient Quantization for Memory-Efficient Distributed Training of LLMs

📄 arXiv: 2605.00539v1 📥 PDF

作者: Wenxiang Lin, Juntao Huang, Luhan Zhang, Laili Li, Xiang Bao, Mengyang Zhang, Bing Wang, Shaohuai Shi

分类: cs.CL, cs.DC

发布日期: 2026-05-01

💡 一句话要点

AGoQ：通过激活和梯度量化实现LLM分布式训练的内存高效性

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

关键词: 量化训练 大规模语言模型 分布式训练 内存优化 梯度量化 激活量化 LLaMA All-Reduce

📋 核心要点

  1. 现有量化方法在4比特激活和8比特梯度下训练LLM时，存在收敛慢或精度损失的问题。
  2. AGoQ通过层感知激活量化和精度保持的梯度量化，在降低内存占用的同时，保证训练效率和模型精度。
  3. 实验表明，AGoQ在LLaMA模型上显著降低内存占用并提升训练速度，同时保持了预训练收敛性和下游任务精度。

📝 摘要（中文）

本文提出AGoQ，一种用于大规模语言模型（LLM）训练的激活和梯度量化方法，旨在降低GPU内存需求。现有方法在4比特激活和8比特梯度量化方面效果不佳，容易导致收敛缓慢或精度损失。AGoQ包含两项新技术：一是层感知的激活量化算法，基于层类型和流水线阶段为不同层的激活分配适当的比特宽度，从而实现接近4比特的激活存储；二是梯度量化算法，通过采用8比特梯度存储和精度保持的8比特All-Reduce通信，减少内存使用并缩短通信时间。在不同规模的LLM上，使用两个GPU集群（最多64个GPU）进行了大量实验。结果表明，与最先进的训练系统Megatron-LM（有或没有ZeRO）、COAT和DeepSpeed相比，AGoQ在8B到32B LLaMA模型上，内存减少高达52%，训练速度提高高达1.34倍，同时在预训练中实现了收敛损失，并在具有LLaMA架构的下游任务中实现了相当的精度。

🔬 方法详解

问题定义：论文旨在解决大规模语言模型（LLM）训练过程中GPU内存需求过高的问题。现有的量化方法在极低比特（如4比特激活和8比特梯度）下，往往会导致训练收敛速度慢、模型精度下降等问题，难以满足实际应用需求。

核心思路：AGoQ的核心思路是针对激活和梯度分别设计量化策略，在保证模型性能的前提下，尽可能降低内存占用。对于激活，采用层感知的量化方法，根据不同层的特性分配不同的比特宽度；对于梯度，采用精度保持的量化通信方法，减少通信开销。

技术框架：AGoQ包含两个主要模块：层感知激活量化和梯度量化。层感知激活量化模块根据层类型和流水线阶段，动态调整激活的量化比特宽度。梯度量化模块则采用8比特存储和All-Reduce通信，减少内存占用和通信时间。整体训练流程与标准的分布式训练流程类似，只是在激活和梯度计算完成后，会进行量化操作。

关键创新：AGoQ的关键创新在于层感知的激活量化策略和精度保持的梯度量化通信。层感知激活量化能够更精细地控制不同层的量化误差，避免信息损失。精度保持的梯度量化通信则能够在降低通信开销的同时，保证梯度信息的有效传递。与现有方法相比，AGoQ能够更好地平衡内存占用、训练速度和模型精度。

关键设计：层感知激活量化中，需要根据不同层的类型（如Transformer层、Embedding层等）和流水线阶段（如前向传播、反向传播等）设置不同的量化比特宽度。梯度量化通信中，需要设计合适的量化和反量化策略，以保证梯度信息的精度。具体的参数设置和量化策略的选择需要根据实际的模型和数据集进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，AGoQ在8B到32B LLaMA模型上，相比于Megatron-LM、COAT和DeepSpeed等先进系统，内存减少高达52%，训练速度提高高达1.34倍。同时，AGoQ在预训练中实现了收敛损失，并在下游任务中实现了与原始模型相当的精度。这些结果验证了AGoQ在降低内存占用和提升训练效率方面的有效性。

🎯 应用场景

AGoQ可应用于大规模语言模型的分布式训练，尤其是在GPU资源受限的环境下。该方法能够降低训练成本，加速模型迭代，并促进LLM在更多领域的应用，例如自然语言处理、机器翻译、文本生成等。未来，AGoQ有望扩展到其他类型的深度学习模型和硬件平台。

📄 摘要（原文）

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.