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

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

## Key facts

- Cited by 96 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 45 references that are in the map.
- First posted on arXiv 28 Oct 2021.

## Builds on (most cited first, 21 of 21)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [Sequence Parallelism: Long Sequence Training from System Perspective](https://mlsanity.com/paper/2105.13120/) (2021)
- [PyTorch Distributed: Experiences on Accelerating Data Parallel Training](https://mlsanity.com/paper/2006.15704/) (2020)
- [Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines](https://mlsanity.com/paper/2107.06925/) (2021)
- [Maximizing Parallelism in Distributed Training for Huge Neural Networks](https://mlsanity.com/paper/2105.14450/) (2021)

## Built on by (most cited first, 30 of 96)

- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [LlamaFactory: Unified Efficient Fine-Tuning of 100+ Language Models](https://mlsanity.com/paper/2403.13372/) (2024)
- [A Comprehensive Overview of Large Language Models](https://mlsanity.com/paper/2307.06435/) (2023)
- [Scaling Data-Constrained Language Models](https://mlsanity.com/paper/2305.16264/) (2023)
- [ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models](https://mlsanity.com/paper/2310.10505/) (2023)
- [Open-Sora 2.0: Training a Commercial-Level Video Generation Model in 200k](https://mlsanity.com/paper/2503.09642/) (2025)
- [LLM Post-Training: A Deep Dive into Reasoning Large Language Models](https://mlsanity.com/paper/2502.21321/) (2025)
- [Parrot: Efficient Serving of LLM-based Applications with Semantic Variable](https://mlsanity.com/paper/2405.19888/) (2024)
- [Response Length Perception and Sequence Scheduling: An LLM-Empowered LLM Inference Pipeline](https://mlsanity.com/paper/2305.13144/) (2023)
- [FP8-LM: Training FP8 Large Language Models](https://mlsanity.com/paper/2310.18313/) (2023)
- [OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework](https://mlsanity.com/paper/2405.11143/) (2024)
- [Advancing Transformer Architecture in Long-Context Large Language Models: A Comprehensive Survey](https://mlsanity.com/paper/2311.12351/) (2023)
- [Model Compression and Efficient Inference for Large Language Models: A Survey](https://mlsanity.com/paper/2402.09748/) (2024)
- [USP: A Unified Sequence Parallelism Approach for Long Context Generative AI](https://mlsanity.com/paper/2405.07719/) (2024)
- [Understanding LLMs: A Comprehensive Overview from Training to Inference](https://mlsanity.com/paper/2401.02038/) (2024)
- [RLinf: Flexible and Efficient Large-scale Reinforcement Learning via Macro-to-Micro Flow Transformation](https://mlsanity.com/paper/2509.15965/) (2025)
- [RenAIssance: A Survey into AI Text-to-Image Generation in the Era of Large Model](https://mlsanity.com/paper/2309.00810/) (2023)
- [A Survey on Transformer Compression](https://mlsanity.com/paper/2402.05964/) (2024)
- [Over-Tokenized Transformer: Vocabulary is Generally Worth Scaling](https://mlsanity.com/paper/2501.16975/) (2025)
- [Reinforced MLLM: A Survey on RL-Based Reasoning in Multimodal Large Language Models](https://mlsanity.com/paper/2504.21277/) (2025)
- [A Survey on Efficient Training of Transformers](https://mlsanity.com/paper/2302.01107/) (2023)
- [Human-Imperceptible Retrieval Poisoning Attacks in LLM-Powered Applications](https://mlsanity.com/paper/2404.17196/) (2024)
- [Peacock: A Family of Arabic Multimodal Large Language Models and Benchmarks](https://mlsanity.com/paper/2403.01031/) (2024)
- [On Efficient Training of Large-Scale Deep Learning Models: A Literature Review](https://mlsanity.com/paper/2304.03589/) (2023)
- [Reinforcement Learning with Human Feedback for Realistic Traffic Simulation](https://mlsanity.com/paper/2309.00709/) (2023)
- [INTERS: Unlocking the Power of Large Language Models in Search with Instruction Tuning](https://mlsanity.com/paper/2401.06532/) (2024)
- [VeOmni: Scaling Any Modality Model Training with Model-Centric Distributed Recipe Zoo](https://mlsanity.com/paper/2508.02317/) (2025)
- [Large Sequence Models for Sequential Decision-Making: A Survey](https://mlsanity.com/paper/2306.13945/) (2023)
- [Protein Language Models and Structure Prediction: Connection and Progression](https://mlsanity.com/paper/2211.16742/) (2022)
- [EMMA-500: Enhancing Massively Multilingual Adaptation of Large Language Models](https://mlsanity.com/paper/2409.17892/) (2024)

## Abstract (arXiv)

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.