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

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Lin, Wenxiang, Huang, Juntao, Zhang, Luhan, Li, Laili, Bao, Xiang, Zhang, Mengyang, Wang, Bing, Shi, Shaohuai
Format: Preprint
Veröffentlicht: 2026
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866909032677113856
author Lin, Wenxiang
Huang, Juntao
Zhang, Luhan
Li, Laili
Bao, Xiang
Zhang, Mengyang
Wang, Bing
Shi, Shaohuai
author_facet Lin, Wenxiang
Huang, Juntao
Zhang, Luhan
Li, Laili
Bao, Xiang
Zhang, Mengyang
Wang, Bing
Shi, Shaohuai
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2605_00539
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle AGoQ: Activation and Gradient Quantization for Memory-Efficient Distributed Training of LLMs
Lin, Wenxiang
Huang, Juntao
Zhang, Luhan
Li, Laili
Bao, Xiang
Zhang, Mengyang
Wang, Bing
Shi, Shaohuai
Computation and Language
Distributed, Parallel, and Cluster Computing
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.
title AGoQ: Activation and Gradient Quantization for Memory-Efficient Distributed Training of LLMs
topic Computation and Language
Distributed, Parallel, and Cluster Computing
url https://arxiv.org/abs/2605.00539