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

TL;DR

Colossal-AI presents a unified, modular framework for large-scale distributed training that integrates data, pipeline, tensor, and sequence parallelism with heterogeneous training and memory-saving strategies. Its design enables flexible combinations of acceleration techniques, automatic parallelization, and memory reuse to scale models to billions of parameters. Empirical evaluations show up to 2.76x speedups over established baselines like Megatron-LM and DeepSpeed across diverse models and hardware. The work contributes practical tooling for democratizing large-scale training and lays groundwork for hardware-aware automatic parallelization and wider ecosystem integration.

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.

Figures (14)

• Figure 1: Architecture of Colossal-AI
• Figure 2: Architecture of The Transformer Layer
• Figure 3: Existing parallelism for distributed training
• Figure 4: Megatron-LM MLP Module
• Figure 5: Scaling Performance of Tensor Parallelism in Theoretic Analysis ($h=1024, s=512, b=32$)