# TAGC: Optimizing Gradient Communication in Distributed Transformer Training

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 33 references that are in the map.
- First posted on arXiv 8 Apr 2025.

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

- [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)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [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)
- [VisionLLM: Large Language Model is also an Open-Ended Decoder for Vision-Centric Tasks](https://mlsanity.com/paper/2305.11175/) (2023)
- [Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training](https://mlsanity.com/paper/1712.01887/) (2017)
- [PACT: Parameterized Clipping Activation for Quantized Neural Networks](https://mlsanity.com/paper/1805.06085/) (2018)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Qsparse-local-SGD: Distributed SGD with Quantization, Sparsification, and Local Computations](https://mlsanity.com/paper/1906.02367/) (2019)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Text-to-Audio Generation using Instruction-Tuned LLM and Latent Diffusion Model](https://mlsanity.com/paper/2304.13731/) (2023)
- [Liger Kernel: Efficient Triton Kernels for LLM Training](https://mlsanity.com/paper/2410.10989/) (2024)
- [High-Dimensional Stochastic Gradient Quantization for Communication-Efficient Edge Learning](https://mlsanity.com/paper/1910.03865/) (2019)
- [PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers](https://mlsanity.com/paper/2102.03161/) (2021)
- [Progressive Gradient Flow for Robust N:M Sparsity Training in Transformers](https://mlsanity.com/paper/2402.04744/) (2024)
- [Accelerating Distributed Deep Learning using Lossless Homomorphic Compression](https://mlsanity.com/paper/2402.07529/) (2024)

## Built on by (most cited first, 2 of 2)

- [TACO: Efficient Communication Compression of Intermediate Tensors for Scalable Tensor-Parallel LLM Training](https://mlsanity.com/paper/2604.24088/) (2026)
- [CIERA: Cross-Iteration Exponent Reuse for Lossless Allgather in Sharded MoE Training](https://mlsanity.com/paper/2609.04609/) (2026)

## Abstract (arXiv)

The increasing complexity of large language models (LLMs) necessitates efficient training strategies to mitigate the high computational costs associated with distributed training. A significant bottleneck in this process is gradient synchronization across multiple GPUs, particularly in the zero-redundancy parallelism mode. In this paper, we introduce Transformer-Aware Gradient Compression (TAGC), an optimized gradient compression algorithm designed specifically for transformer-based models. TAGC extends the lossless homomorphic compression method by adapting it for sharded models and incorporating transformer-specific optimizations, such as layer-selective compression and dynamic sparsification. Our experimental results demonstrate that TAGC accelerates training by up to 15% compared to the standard Fully Sharded Data Parallel (FSDP) approach, with minimal impact on model quality. We integrate TAGC into the PyTorch FSDP framework, the implementation is publicly available at https://github.com/ipolyakov/TAGC.