Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

CO2: Efficient Distributed Training with Full Communication-Computation Overlap

Weigao Sun, Zhen Qin, Weixuan Sun, Shidi Li, Dong Li, Xuyang Shen, Yu Qiao, Yiran Zhong

TL;DR

CO2 tackles the challenge of scalable distributed training on bandwidth-limited clusters by enabling full overlap of communication and computation through local updates and asynchronous model synchronization after every $\tau$ local steps. It stabilizes asynchronous outer updates with a staleness gap penalty and outer momentum clipping, and provides a convergence bound of $O\left(\frac{1}{\sqrt{G T \tau}}\right) + O\left(\frac{G \tau}{T}\right)$ under standard assumptions. Empirically, CO2 delivers strong convergence and generalization across computer vision and natural language processing tasks on up to 128 GPUs, achieving near-100% scalability on both RoCE and TCP/IP interconnects, and integrates smoothly with ZeRO-series optimizers to reduce memory usage. The work offers a practical, theory-backed approach to democratize large-scale training under limited hardware and networking resources.

Abstract

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.

CO2: Efficient Distributed Training with Full Communication-Computation Overlap

TL;DR

CO2 tackles the challenge of scalable distributed training on bandwidth-limited clusters by enabling full overlap of communication and computation through local updates and asynchronous model synchronization after every local steps. It stabilizes asynchronous outer updates with a staleness gap penalty and outer momentum clipping, and provides a convergence bound of under standard assumptions. Empirically, CO2 delivers strong convergence and generalization across computer vision and natural language processing tasks on up to 128 GPUs, achieving near-100% scalability on both RoCE and TCP/IP interconnects, and integrates smoothly with ZeRO-series optimizers to reduce memory usage. The work offers a practical, theory-backed approach to democratize large-scale training under limited hardware and networking resources.

Abstract

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.
Paper Structure (32 sections, 1 theorem, 18 equations, 8 figures, 10 tables, 1 algorithm)

This paper contains 32 sections, 1 theorem, 18 equations, 8 figures, 10 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Method
  4. The Overview
  5. Staleness Gap Penalty
  6. Outer Momentum Clipping
  7. Convergence Analysis
  8. Experiments
  9. Experimental Setup
  10. Convergence and Generalization
  11. Scalability
  12. Ablation study
  13. Conclusion
  14. Experiment Details
  15. Hardware and Software
  16. ...and 17 more sections

Key Result

Theorem 1

If we take $\lambda=1/\Lambda_t,\gamma_t=\gamma$ and ignore the clip operation, such that $\bar{\lambda} = \alpha \lambda, \frac{ \bar{\lambda} \gamma}{1-\beta}=\sqrt{\frac{G}{T \tau}}$ and $T \tau \geq G L^2\left(1+\sqrt{3} \max \left\{\frac{3 \tau(1-\beta-\alpha)}{\alpha}, \frac{4 \tau \beta}{1-\

Figures (8)

  • Figure 1: Visualization of CO2 and SGD. We exemplify the mechanism of CO2 with a local step count $\tau=2$. This configuration dictates that the outer update starts after every two local steps, concurrently launching an AAR operation on model parameters. This strategy is made to make the full overlap of AAR communication with local computation possible. CO2 can effectively reduce the wall time required for training compared to the conventional SGD in DDP paradigm.
  • Figure 2: Validation curves for image classification tasks. Three models ResNet-50, ViT, and VVT are trained on ImageNet-1K for 90, 300 and 300 epochs, respectively. Our CO2 exhibits robust convergence and good generalization performance when compared to other existing methods.
  • Figure 3: Validation curves for autoregressive language tasks. We train GPT-2 on OpenWebText for 100K steps in three sizes: 125M (Small), 355M (Medium), and 770M (Large). CO2 exhibits robust convergence and the best generalization performance when compared to other existing methods.
  • Figure 4: (a): Scalability of CO2. Throughput (words/sec) results on distinctive inter-node network configurations are presented. CO2 exhibits pecfect $100\%$ scalability on both configurations. (b): Effects of $\tau$. Training speed and generalization performance results w.r.t. $\tau$ are presented. A larger value of $\tau$ leads higher communication efficiency but worse generalization behaviors.
  • Figure 5: Training curves for image classification tasks. Three models ResNet-50, ViT and VVT are trained on ImageNet-1K for 90, 300 and 300 epochs, respectively. Our CO2 exhibits robust convergence compared to other existing methods.
  • ...and 3 more figures

Theorems & Definitions (1)

  • Theorem 1