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

1. Problem Statement

Training Large Language Models (LLMs) requires massive GPU memory, primarily consumed by model parameters, optimizer states, gradients, and activations.

• Memory Bottleneck: Activations typically occupy the largest portion of memory during training, scaling linearly with sequence length and batch size.
• Limitations of Existing Methods:
  • Activation Recomputation/Offloading: Reduces memory but introduces significant computational overhead (recomputing) or I/O latency (offloading).
  • Standard Quantization: While 8-bit quantization (INT8/FP8) is common, applying it to 4-bit activations or 8-bit gradients often leads to severe accuracy loss, slow convergence, or training instability due to numerical overflow during accumulation and All-Reduce operations.
  • Gradient Communication: Current systems often store gradients in high precision (FP32) and only quantize them for communication, failing to reduce the memory footprint of stored gradients.

2. Methodology: AGoQ Framework

The authors propose AGoQ, a holistic framework integrating two novel techniques to achieve near 4-bit activation storage and 8-bit gradient storage/communication without sacrificing model convergence.

A. Layer-Aware Activation Quantization (LAAQ)

Instead of uniformly quantizing all layers to 4 bits (which causes accuracy loss), LAAQ dynamically assigns bit-widths based on layer types and pipeline stages.

• Theoretical Error Analysis: The authors derive error bounds for different operations (RMSNorm, SiLU, Attention, GEMM).
  • Findings: Quantizing Attention module activations (Q, K, V) introduces significantly larger gradient errors compared to other modules. Therefore, Attention activations are not quantized (kept in higher precision).
  • Strategy: For other layers (RMSNorm, SiLU, FFN), the system uses 4-bit block-wise quantization.
  • Recomputation vs. Caching: The analysis shows that for operations like RMSNorm and SiLU, storing only the quantized input and recomputing intermediate values during the backward pass yields a tighter error bound than caching the intermediate values in quantized form. Thus, AGoQ stores only 4-bit inputs and recomputes intermediates.
• Dynamic Bit-width Compensation (DBCA-PP):
  • In Pipeline Parallelism (PP), different pipeline stages hold different numbers of activation batches (uneven memory usage).
  • AGoQ exploits this by assigning higher bit-widths (e.g., 5, 6, or 8 bits) to stages with fewer active batches and lower bit-widths (4 bits) to stages with more batches. This balances the total memory usage while maintaining an average of ~4 bits per element, compensating for precision loss where memory is available.

B. Precision-Preserved Gradient Quantization (QuanGrad)

This technique reduces both memory usage and communication time by using 8-bit (FP8) gradients throughout the process.

• Local Accumulation: To prevent overflow when accumulating gradients from multiple mini-batches, the system dequantizes the FP8 main gradient to FP16/BF16, adds the new local gradient, and then re-quantizes the sum back to FP8.
• Precision-Preserved All-Reduce: Standard FP8 All-Reduce causes overflow during the reduction step. AGoQ decomposes the All-Reduce operation into:
  1. All-to-All: Send compressed FP8 gradients to all devices.
  2. Local Reduce: Each device dequantizes received data to FP32, performs the local sum, and re-quantizes to FP8.
  3. All-Gather: Gather the reduced results.
  • This ensures that the accumulation happens in high precision (FP32) while communication remains in low precision (FP8), preventing overflow and maintaining convergence.

C. Kernel Fusion

To eliminate the overhead of quantization/dequantization operations, AGoQ fuses these element-wise operations with adjacent GEMM (General Matrix Multiply) kernels. This leverages Tensor Cores for the heavy lifting while keeping quantization logic on CUDA cores within the same kernel, minimizing latency.

3. Key Contributions

1. Layer-Aware Quantization Strategy: A theoretical and practical approach to identify which layers can be safely compressed to 4 bits (excluding Attention) and which require recomputation strategies to minimize error.
2. Dynamic Bit-width Compensation: A novel algorithm for Pipeline Parallelism that dynamically adjusts quantization precision based on memory availability at each stage, maximizing memory efficiency without increasing peak usage.
3. Precision-Preserved Gradient Communication: A method to perform 8-bit All-Reduce by splitting the operation to allow high-precision accumulation, solving the overflow issues that previously prevented low-bit gradient training.
4. System Integration: Full integration into Megatron-LM, compatible with 8-bit Adam optimizers, supporting both NVIDIA GPUs and Huawei Ascend NPUs.

4. Experimental Results

Experiments were conducted on clusters with up to 64 GPUs (NVIDIA A6000/Pro 6000) using LLaMA models (7B, 8B, 13B, 34B) and OLMo-1B.

• Memory Reduction:

  • AGoQ reduces peak memory usage by up to 52% compared to full-precision training (Megatron-LM).
  • Compared to the state-of-the-art COAT system, AGoQ reduces activation memory by an additional 30% and gradient memory by 75%.
  • It enables training longer sequence lengths (e.g., 80K tokens) that cause Out-Of-Memory (OOM) errors in baseline systems.

• Training Speed:

  • Achieves up to 1.34× speedup in end-to-end training time compared to Megatron-LM (with or without ZeRO) and DeepSpeed.
  • The speedup is more pronounced with longer sequence lengths due to reduced communication overhead and memory constraints.

• Convergence and Accuracy:

  • Convergence: Training loss curves on LLaMA2-7B and LLaMA3-8B closely match the full-precision baseline. In contrast, naive FP8 All-Reduce methods showed significantly higher loss.
  • Downstream Tasks: Zero-shot accuracy on datasets (ARC, HellaSwag, PIQA, etc.) remains comparable to the baseline, with some tasks showing slight improvements.
  • Stability: Gradient norms remain stable throughout training, confirming the effectiveness of the precision-preserved accumulation.

5. Significance

• Enabling Long-Context Training: By drastically reducing memory footprints, AGoQ makes it feasible to train LLMs with very long context windows (e.g., 64K–80K tokens) on existing hardware without requiring massive clusters.
• Bridging the Gap: It successfully pushes quantization to 4-bit for activations and 8-bit for gradients, a regime previously considered too unstable for LLM training, without compromising model quality.
• Scalability: The method is highly scalable across different parallelism strategies (DP, TP, PP) and hardware architectures, offering a practical path toward training trillion-parameter models on limited resources.

In summary, AGoQ represents a significant advancement in systems for LLM training, moving beyond simple compression to a mathematically grounded, adaptive quantization framework that balances memory efficiency, communication speed, and model accuracy.