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

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.

Training large language models (LLMs) requires massive GPU memory. Most of that memory is eaten up by activations and gradients — like trying to move a whole house with a tiny van. We introduce AGoQ, a method that squeezes both activations and gradients into much smaller formats without hurting the model’s learning ability. 1. For activations, we assign just 4 bits on average (instead of the usual 16) by treating different layers of the model according to their sensitivity. 2. For gradients, we store and communicate them using a special 8‑bit format that avoids overflow when summing gradients across many GPUs. The result: up to 52% less GPU memory and up to 34% faster training, while keeping the same accuracy as standard full‑precision training. AGoQ makes LLM training possible on more affordable hardware, opening up large‑scale AI research to groups with limited resources.