# Colossal-AI

A deep learning training system for large AI models that combines data, pipeline, tensor and ZeRO-style parallelism so a model can be split and trained across many GPUs at lower cost.

Colossal-AI provides a set of parallelism strategies, including ZeRO-style optimizer and gradient partitioning, 3D parallelism and heterogeneous training, for distributing large model training across GPU clusters.

In the project's own benchmark on Llama-like models, a 7-billion-parameter model trained at 17.13 samples per second on eight H200 GPUs using ZeRO stage 2 data parallelism, and a 70-billion-parameter model trained at 3.27 samples per second across sixteen H200 GPUs combining ZeRO, tensor and pipeline parallelism; the same configurations on B200 GPUs reached higher throughput and more TFLOPS per GPU. The project has also been used as the basis for other open-source work, including the Open-Sora video generation model.

It is written in Python and distributed under the Apache-2.0 license.

Website: https://www.colossalai.org