# BAGUA: Scaling up Distributed Learning with System Relaxations

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

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 51 of its 89 references that are in the map.
- First posted on arXiv 3 Jul 2021.

## Builds on (most cited first, 30 of 51)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Local SGD Converges Fast and Communicates Little](https://mlsanity.com/paper/1805.09767/) (2018)
- [signSGD: Compressed Optimisation for Non-Convex Problems](https://mlsanity.com/paper/1802.04434/) (2018)
- [Can Decentralized Algorithms Outperform Centralized Algorithms? A Case Study for Decentralized Parallel Stochastic Gradient Descent](https://mlsanity.com/paper/1705.09056/) (2017)
- [Sparsified SGD with Memory](https://mlsanity.com/paper/1809.07599/) (2018)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning](https://mlsanity.com/paper/1807.06629/) (2018)
- [TernGrad: Ternary Gradients to Reduce Communication in Distributed Deep Learning](https://mlsanity.com/paper/1705.07878/) (2017)
- [The Convergence of Sparsified Gradient Methods](https://mlsanity.com/paper/1809.10505/) (2018)
- [Decentralized Stochastic Optimization and Gossip Algorithms with Compressed Communication](https://mlsanity.com/paper/1902.00340/) (2019)
- [Don't Use Large Mini-Batches, Use Local SGD](https://mlsanity.com/paper/1808.07217/) (2018)
- [Gradient Sparsification for Communication-Efficient Distributed Optimization](https://mlsanity.com/paper/1710.09854/) (2017)
- [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)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [Asynchronous Decentralized Parallel Stochastic Gradient Descent](https://mlsanity.com/paper/1710.06952/) (2017)
- [D2: Decentralized Training over Decentralized Data](https://mlsanity.com/paper/1803.07068/) (2018)
- [On Biased Compression for Distributed Learning](https://mlsanity.com/paper/2002.12410/) (2020)
- [ATOMO: Communication-efficient Learning via Atomic Sparsification](https://mlsanity.com/paper/1806.04090/) (2018)
- [Memory-Efficient Pipeline-Parallel DNN Training](https://mlsanity.com/paper/2006.09503/) (2020)
- [AISHELL-2: Transforming Mandarin ASR Research Into Industrial Scale](https://mlsanity.com/paper/1808.10583/) (2018)
- [Local SGD with Periodic Averaging: Tighter Analysis and Adaptive Synchronization](https://mlsanity.com/paper/1910.13598/) (2019)
- [Asynchronous Stochastic Gradient Descent with Delay Compensation](https://mlsanity.com/paper/1609.08326/) (2016)
- [DoubleSqueeze: Parallel Stochastic Gradient Descent with Double-Pass Error-Compensated Compression](https://mlsanity.com/paper/1905.05957/) (2019)

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

- [Scaling Expert Language Models with Unsupervised Domain Discovery](https://mlsanity.com/paper/2303.14177/) (2023)
- [Fine-tuning Language Models over Slow Networks using Activation Compression with Guarantees](https://mlsanity.com/paper/2206.01299/) (2022)
- [On Efficient Training of Large-Scale Deep Learning Models: A Literature Review](https://mlsanity.com/paper/2304.03589/) (2023)
- [Persia: An Open, Hybrid System Scaling Deep Learning-based Recommenders up to 100 Trillion Parameters](https://mlsanity.com/paper/2111.05897/) (2021)
- [Data Debugging with Shapley Importance over End-to-End Machine Learning Pipelines](https://mlsanity.com/paper/2204.11131/) (2022)
- [BlueFog: Make Decentralized Algorithms Practical for Optimization and Deep Learning](https://mlsanity.com/paper/2111.04287/) (2021)
- [Merak: An Efficient Distributed DNN Training Framework with Automated 3D Parallelism for Giant Foundation Models](https://mlsanity.com/paper/2206.04959/) (2022)
- [Training and Serving System of Foundation Models: A Comprehensive Survey](https://mlsanity.com/paper/2401.02643/) (2024)
- [KVServe: Service-Aware KV Cache Compression for Communication-Efficient Disaggregated LLM Serving](https://mlsanity.com/paper/2605.13734/) (2026)
- [CGX: Adaptive System Support for Communication-Efficient Deep Learning](https://mlsanity.com/paper/2111.08617/) (2021)
- [Optimal Complexity in Non-Convex Decentralized Learning over Time-Varying Networks](https://mlsanity.com/paper/2211.00533/) (2022)
- [Investigation of Energy-efficient AI Model Architectures and Compression Techniques for "Green" Fetal Brain Segmentation](https://mlsanity.com/paper/2405.15778/) (2024)
- [Fully First-Order Methods for Decentralized Bilevel Optimization](https://mlsanity.com/paper/2410.19319/) (2024)

## Abstract (arXiv)

Recent years have witnessed a growing list of systems for distributed data-parallel training. Existing systems largely fit into two paradigms, i.e., parameter server and MPI-style collective operations. On the algorithmic side, researchers have proposed a wide range of techniques to lower the communication via system relaxations: quantization, decentralization, and communication delay. However, most, if not all, existing systems only rely on standard synchronous and asynchronous stochastic gradient (SG) based optimization, therefore, cannot take advantage of all possible optimizations that the machine learning community has been developing recently. Given this emerging gap between the current landscapes of systems and theory, we build BAGUA, a MPI-style communication library, providing a collection of primitives, that is both flexible and modular to support state-of-the-art system relaxation techniques of distributed training. Powered by this design, BAGUA has a great ability to implement and extend various state-of-the-art distributed learning algorithms. In a production cluster with up to 16 machines (128 GPUs), BAGUA can outperform PyTorch-DDP, Horovod and BytePS in the end-to-end training time by a significant margin (up to 2 times) across a diverse range of tasks. Moreover, we conduct a rigorous tradeoff exploration showing that different algorithms and system relaxations achieve the best performance over different network conditions.