Research map / Architectures and efficiency / RNNs and normalization
CO2: Efficient Distributed Training with Full Communication-Computation Overlap
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes CO2, a local-updating method with one-step asynchronous all-reduce of parameters that fully overlaps communication with computation for distributed training on bandwidth-limited clusters, plus a staleness gap penalty and outer momentum clipping.
Problem
Distributed data-parallel training of large models is bottlenecked by inter-node communication on clusters with limited bandwidth, and existing methods do not fully overlap communication with computation without hurting final accuracy.
Approach
Combines local updating (Local-SGD style) with one-step asynchronous all-reduce of model parameters so that synchronization of outer iterates overlaps with the next round of local computation; adds a staleness gap penalty on the outer momentum (ratio of outer displacement to maximum inner displacement) and coordinate-wise outer momentum clipping for stability.
Key results
- On RDMA with TransNormer-LLM (7B) on 128 A100 GPUs, CO2 (tau=12) achieves scalability ratio 1.08 versus 0.93 for AdamW from 16 to 128 GPUs.
- On TCP/IP (~80Gbps), CO2 (tau=48) reaches 304736 words/sec at 128 GPUs versus 22810 for AdamW, scalability ratio 1.06 vs 0.82.
- CO2 matches or beats AdamW perplexity on GPT-2 Small/Medium/Large (7.37, 6.36, 6.13 vs 7.44, 6.61, 6.26) and ResNet-50 top-1 77.14 vs 76.92 for SGD.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use CO2 instead of the strongest reported baseline | ResNet-50 · ImageNet-1K val | Top-1 accuracy: 77.12 ($\pm$ 0.11) -> 77.14 ($\pm$ 0.09) | No clear effect |
| Use CO2 instead of SGD | ViT (Base) · ImageNet-1K val | Top-1 accuracy: 81.33 ($\pm$ 0.04) -> 80.95 ($\pm$ 0.08) | Hurt |
| Use CO2 instead of SGD | VVT (Large) · ImageNet-1K val | Top-1 accuracy: 83.64 ($\pm$ 0.06) -> 83.38 ($\pm$ 0.06) | Hurt |
| Use CO2 instead of SGD | VVT (Large) · ADE20K | accuracy: 47.82 ($\pm$ 0.05) -> 47.80 ($\pm$ 0.11) | No clear effect |
| Use CO2 instead of SlowMo | Point-MAE · ShapeNet validation | accuracy: 68.69 ($\pm$ 0.32) -> 68.89 ($\pm$ 0.39) | Helped |
| Use CO2 instead of the best baseline | GPT-2 (Medium) · OpenWebText | perplexity: 6.41 ($\pm$ 0.65) -> 6.36 ($\pm$ 0.66) | Helped |
| Use CO2 instead of Adamw | TN-LLM (7B) · WikiText-103 | perplexity: 16.82 ($\pm$ 0.86) -> 16.78 ($\pm$ 0.95) | Helped |
| Use CO2 instead of Adamw | RoBERTa (Large) · Wikipedia+BookCorpus | perplexity: 3.96 ($\pm$ 0.37) -> 3.95 ($\pm$ 0.96) | No clear effect |
| Use CO2 (τ=12) instead of Adamw | TransNormer-LLM (7B) · RDMA, 128 GPUs | throughput (words/sec): 289106 ($\pm$ 423) -> 307557 ($\pm$ 617) | Helped |
| Use CO2 (τ=48) instead of Adamw | TransNormer-LLM (7B) · TCP/IP, 128 GPUs | throughput (words/sec): 22810 ($\pm$ 526) -> 304736 ($\pm$ 729) | Helped |
| Add staleness gap penalty (ablation) | CO2 on GPT-2 (Small) · 100K steps | validation perplexity: 7.56 -> 7.39 | Helped |
| Add outer momentum clipping (ablation) | CO2 on GPT-2 (Small) · 100K steps | validation perplexity: 7.42 -> 7.39 | Helped |
Surprising findings
- CO2 with the staleness gap penalty and clipping achieves slightly better validation perplexity than AdamW on GPT-2 despite one-step asynchronous staleness.
- Scalability ratios above 1 are observed, which the authors attribute mainly to measurement fluctuations.
Limitations
- Staleness gap penalty may not help with every learning rate schedule (authors note possible sub-optimality with some schedules).
- Requires tuning of local step count tau per task and cluster; larger tau can increase perplexity.
- Convergence guarantee relies on constant learning rate and ignores outer momentum clipping in the analysis.
Experimental setup
Datasets: ImageNet-1K ADE20K ShapeNet OpenWebText WikiText-103 Wikipedia BookCorpus
Benchmarks: ImageNet-1K ADE20K ModelNet40 OpenWebText WikiText-103
Metrics: top-1 accuracy mIoU-like segmentation accuracy perplexity throughput (words/sec) scalability ratio
Baselines: SGD AdamW Local-SGD Overlap-Local-SGD SlowMo
Builds on: Local-SGD SlowMo Overlap-Local-SGD ZeRO
Closest papers by meaning
- Rethinking Memory and Communication Costs for Efficient Data Parallel Training of Large Language Models (NeurIPS 2024 · similarity 0.87)
- AsyncMesh: Fully Asynchronous Optimization for Data and Pipeline Parallelism (NeurIPS 2026 · similarity 0.87)
- Mixtures of Subspaces for Bandwidth Efficient Context Parallel Training (NeurIPS 2025 · similarity 0.86)
- HALoS: Hierarchical Asynchronous Local SGD over Slow Networks for Geo-Distributed Large Language Model Training (ICML 2025 · similarity 0.86)
- FlowMoE: A Scalable Pipeline Scheduling Framework for Distributed Mixture-of-Experts Training (NeurIPS 2025 · similarity 0.85)
- Communication Efficient Distributed Training with Distributed Lion (NeurIPS 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.