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

arXiv: https://arxiv.org/abs/2605.00539 · HTML: https://mlsanity.com/paper/2605.00539/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 28 of its 43 references that are in the map.
- First posted on arXiv 1 May 2026.

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Crowdsourcing Multiple Choice Science Questions](https://mlsanity.com/paper/1707.06209/) (2017)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [8-bit Optimizers via Block-wise Quantization](https://mlsanity.com/paper/2110.02861/) (2021)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [FP8-LM: Training FP8 Large Language Models](https://mlsanity.com/paper/2310.18313/) (2023)
- [1-bit Adam: Communication Efficient Large-Scale Training with Adam's Convergence Speed](https://mlsanity.com/paper/2102.02888/) (2021)
- [Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes](https://mlsanity.com/paper/1807.11205/) (2018)
- [Memory Efficient Optimizers with 4-bit States](https://mlsanity.com/paper/2309.01507/) (2023)
- [COAT: Compressing Optimizer states and Activation for Memory-Efficient FP8 Training](https://mlsanity.com/paper/2410.19313/) (2024)
- [GACT: Activation Compressed Training for Generic Network Architectures](https://mlsanity.com/paper/2206.11357/) (2022)
- [Jetfire: Efficient and Accurate Transformer Pretraining with INT8 Data Flow and Per-Block Quantization](https://mlsanity.com/paper/2403.12422/) (2024)
- [Towards Scalable Distributed Training of Deep Learning on Public Cloud Clusters](https://mlsanity.com/paper/2010.10458/) (2020)
- [CompAct: Compressed Activations for Memory-Efficient LLM Training](https://mlsanity.com/paper/2410.15352/) (2024)
- [gZCCL: Compression-Accelerated Collective Communication Framework for GPU Clusters](https://mlsanity.com/paper/2308.05199/) (2023)
- [SSDTrain: An Activation Offloading Framework to SSDs for Faster Large Language Model Training](https://mlsanity.com/paper/2408.10013/) (2024)
- [Adacc: An Adaptive Framework Unifying Compression and Activation Recomputation for LLM Training](https://mlsanity.com/paper/2508.00806/) (2025)

## Abstract (arXiv)

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× 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.