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

2026-05-01

AI Summary

Existing quantization methods employing 4-bit activations and 8-bit gradients struggle to balance memory efficiency with model convergence, often resulting in slow training or degraded accuracy. This work proposes AGoQ, which introduces a novel layer-aware activation bit-width allocation strategy combined with high-fidelity 8-bit gradient storage and an optimized All-Reduce communication mechanism. AGoQ overcomes the convergence barriers of low-bit training in distributed large language models. Evaluated on LLaMA models ranging from 8B to 32B parameters, AGoQ reduces memory consumption by up to 52% and accelerates training by 1.34× compared to Megatron-LM, COAT, and DeepSpeed, while preserving pretraining convergence and downstream task accuracy.

Abstract

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.

Problem

• quantization
• large language models
• activation
• gradient
• distributed training

Innovation

• activation quantization
• gradient quantization
• memory-efficient training
• distributed LLM training
• low-bit quantization