# CO2: Efficient Distributed Training with Full Communication-Computation Overlap

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

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 53 references that are in the map.
- First posted on arXiv 29 Jan 2024.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Feature Pyramid Networks for Object Detection](https://mlsanity.com/paper/1612.03144/) (2016)
- [ShapeNet: An Information-Rich 3D Model Repository](https://mlsanity.com/paper/1512.03012/) (2015)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books](https://mlsanity.com/paper/1506.06724/) (2015)
- [Local SGD Converges Fast and Communicates Little](https://mlsanity.com/paper/1805.09767/) (2018)
- [Masked Autoencoders for Point Cloud Self-supervised Learning](https://mlsanity.com/paper/2203.06604/) (2022)
- [8-bit Optimizers via Block-wise Quantization](https://mlsanity.com/paper/2110.02861/) (2021)
- [cosFormer: Rethinking Softmax in Attention](https://mlsanity.com/paper/2202.08791/) (2022)
- [Don't Use Large Mini-Batches, Use Local SGD](https://mlsanity.com/paper/1808.07217/) (2018)
- [Stochastic Gradient Push for Distributed Deep Learning](https://mlsanity.com/paper/1811.10792/) (2018)
- [Revisiting Distributed Synchronous SGD](https://mlsanity.com/paper/1702.05800/) (2017)
- [SlowMo: Improving Communication-Efficient Distributed SGD with Slow Momentum](https://mlsanity.com/paper/1910.00643/) (2019)
- [Hierarchically Gated Recurrent Neural Network for Sequence Modeling](https://mlsanity.com/paper/2311.04823/) (2023)
- [The Devil in Linear Transformer](https://mlsanity.com/paper/2210.10340/) (2022)
- [Sharper Convergence Guarantees for Asynchronous SGD for Distributed and Federated Learning](https://mlsanity.com/paper/2206.08307/) (2022)
- [1-bit Adam: Communication Efficient Large-Scale Training with Adam's Convergence Speed](https://mlsanity.com/paper/2102.02888/) (2021)
- [Slow and Stale Gradients Can Win the Race: Error-Runtime Trade-offs in Distributed SGD](https://mlsanity.com/paper/1803.01113/) (2018)
- [Asynchronous SGD Beats Minibatch SGD Under Arbitrary Delays](https://mlsanity.com/paper/2206.07638/) (2022)
- [Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models](https://mlsanity.com/paper/2401.04658/) (2024)
- [Understanding Top-k Sparsification in Distributed Deep Learning](https://mlsanity.com/paper/1911.08772/) (2019)
- [Staleness-aware Async-SGD for Distributed Deep Learning](https://mlsanity.com/paper/1511.05950/) (2015)
- [Toeplitz Neural Network for Sequence Modeling](https://mlsanity.com/paper/2305.04749/) (2023)
- [TransNormerLLM: A Faster and Better Large Language Model with Improved TransNormer](https://mlsanity.com/paper/2307.14995/) (2023)
- [Parallel SGD: When does averaging help?](https://mlsanity.com/paper/1606.07365/) (2016)
- [Asynchronous Stochastic Optimization Robust to Arbitrary Delays](https://mlsanity.com/paper/2106.11879/) (2021)
- [rTop-k: A Statistical Estimation Approach to Distributed SGD](https://mlsanity.com/paper/2005.10761/) (2020)
- [1-bit LAMB: Communication Efficient Large-Scale Large-Batch Training with LAMB's Convergence Speed](https://mlsanity.com/paper/2104.06069/) (2021)

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

- [A Survey of Efficient Reasoning for Large Reasoning Models: Language, Multimodality, and Beyond](https://mlsanity.com/paper/2503.21614/) (2025)
- [MoM: Linear Sequence Modeling with Mixture-of-Memories](https://mlsanity.com/paper/2502.13685/) (2025)
- [MBA-RAG: a Bandit Approach for Adaptive Retrieval-Augmented Generation through Question Complexity](https://mlsanity.com/paper/2412.01572/) (2024)
- [Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunch](https://mlsanity.com/paper/2501.18512/) (2025)
- [FlashMoE: Fast Distributed MoE in a Single Kernel](https://mlsanity.com/paper/2506.04667/) (2025)
- [Native Hybrid Attention for Efficient Sequence Modeling](https://mlsanity.com/paper/2510.07019/) (2025)
- [Linear Attention Sequence Parallelism](https://mlsanity.com/paper/2404.02882/) (2024)
- [ACCO: Accumulate While You Communicate for Communication-Overlapped Sharded LLM Training](https://mlsanity.com/paper/2406.02613/) (2024)
- [Distributed Sign Momentum with Local Steps for Training Transformers](https://mlsanity.com/paper/2411.17866/) (2024)
- [Taming Latency and Bandwidth: A Theoretical Framework and Adaptive Algorithm for Communication-Constrained Training](https://mlsanity.com/paper/2507.17346/) (2025)
- [Partial Parameter Updates for Efficient Distributed Training](https://mlsanity.com/paper/2509.22418/) (2025)
- [LOSCAR-SGD: Local SGD with Communication-Computation Overlap and Delay-Corrected Sparse Model Averaging](https://mlsanity.com/paper/2605.20866/) (2026)
- [LAYUP: Asynchronous decentralized gradient descent with LAYer-wise UPdates](https://mlsanity.com/paper/2410.05985/) (2024)
- [EDiT: A Local-SGD-Based Efficient Distributed Training Method for Large Language Models](https://mlsanity.com/paper/2412.07210/) (2024)
- [Memory and Bandwidth are All You Need for Fully Sharded Data Parallel](https://mlsanity.com/paper/2504.03655/) (2025)
- [Accelerating Decentralized Optimization via Overlapping Local Steps](https://mlsanity.com/paper/2601.01493/) (2026)
- [Cosine-Gated Adam-Decay: Drop-In Staleness-Aware Outer Optimization for Decoupled DiLoCo](https://mlsanity.com/paper/2605.09126/) (2026)

## Abstract (arXiv)

The fundamental success of large language models hinges upon the efficacious implementation of large-scale distributed training techniques. Nevertheless, building a vast, high-performance cluster featuring high-speed communication interconnectivity is prohibitively costly, and accessible only to prominent entities. In this work, we aim to lower this barrier and democratize large-scale training with limited bandwidth clusters. We propose a new approach called CO2 that introduces local-updating and asynchronous communication to the distributed data-parallel training, thereby facilitating the full overlap of COmunication with COmputation. CO2 is able to attain a high scalability even on extensive multi-node clusters constrained by very limited communication bandwidth. We further propose the staleness gap penalty and outer momentum clipping techniques together with CO2 to bolster its convergence and training stability. Besides, CO2 exhibits seamless integration with well-established ZeRO-series optimizers which mitigate memory consumption of model states with large model training. We also provide a mathematical proof of convergence, accompanied by the establishment of a stringent upper bound. Furthermore, we validate our findings through an extensive set of practical experiments encompassing a wide range of tasks in the fields of computer vision and natural language processing. These experiments serve to demonstrate the capabilities of CO2 in terms of convergence, generalization, and scalability when deployed across configurations comprising up to 128 A100 GPUs. The outcomes emphasize the outstanding capacity of CO2 to hugely improve scalability, no matter on clusters with 800Gbps RDMA or 80Gbps TCP/IP inter-node connections.
