English

Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training

Machine Learning 2023-10-06 v3 Artificial Intelligence Computation and Language Computer Vision and Pattern Recognition Distributed, Parallel, and Cluster Computing

Abstract

The success of Transformer models has pushed the deep learning model scale to billions of parameters. Due to the limited memory resource of a single GPU, However, the best practice for choosing the optimal parallel strategy is still lacking, since it requires domain expertise in both deep learning and parallel computing. The Colossal-AI system addressed the above challenge by introducing a unified interface to scale your sequential code of model training to distributed environments. It supports parallel training methods such as data, pipeline, tensor, and sequence parallelism, as well as heterogeneous training methods integrated with zero redundancy optimizer. Compared to the baseline system, Colossal-AI can achieve up to 2.76 times training speedup on large-scale models.

Keywords

Cite

@article{arxiv.2110.14883,
  title  = {Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training},
  author = {Shenggui Li and Hongxin Liu and Zhengda Bian and Jiarui Fang and Haichen Huang and Yuliang Liu and Boxiang Wang and Yang You},
  journal= {arXiv preprint arXiv:2110.14883},
  year   = {2023}
}
R2 v1 2026-06-24T07:15:15.420Z