Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel
  Training
  The Transformer architecture has improved the performance of deep learning
models in domains such as Computer Vision and Natural Language Processing.
Together with better performance come larger model sizes. This imposes
challenges to the memory wall of the current accelerator hardware such as GPU.
It is never ideal to train large models such as Vision Transformer, BERT, and
GPT on a single GPU or a single machine. There is an urgent demand to train
models in a distributed environment. However, distributed training, especially
model parallelism, often requires domain expertise in computer systems and
architecture. It remains a challenge for AI researchers to implement complex
distributed training solutions for their models.
  In this paper, we introduce Colossal-AI, which is a unified parallel training
system designed to seamlessly integrate different paradigms of parallelization
techniques including data parallelism, pipeline parallelism, multiple tensor
parallelism, and sequence parallelism. Colossal-AI aims to support the AI
community to write distributed models in the same way as how they write models
normally. This allows them to focus on developing the model architecture and
separates the concerns of distributed training from the development process.
The documentations can be found at https://www.colossalai.org and the source
code can be found at https://github.com/hpcaitech/ColossalAI.
