ColossalAI: Providing Efficient Large-Scale AI Model Training Solutions

General Introduction

ColossalAI is an open source platform developed by HPC-AI Technologies to provide an efficient and cost-effective solution for large-scale AI model training and inference. By supporting multiple parallelization strategies, heterogeneous memory management, and mixed-precision training, ColossalAI is able to significantly reduce the time and resource consumption of model training and inference. Whether in data parallelism, tensor parallelism, or pipeline parallelism, ColossalAI provides powerful tools and libraries to help researchers and developers realize efficient training and inference of large-scale models on multi-GPU clusters.

Function List

• Data parallelism, tensor parallelism, pipelined parallelism and other parallel strategies support
• Mixed Precision Training and Zero Redundancy Optimizer (ZeRO)
• Heterogeneous memory management to support efficient training of large models
• Support for multiple domain-specific models, such as Open-Sora, Colossal-LLaMA, etc.
• Providing user-friendly tools for distributed training and inference
• Integration of high-performance kernel, KV cache, paging attention and sequential batch processing
• Easy configuration of parallel training through configuration files
• Provide rich examples and documentation to help you get started quickly
• Provides multiple installation options for Docker images and building from source code

Using Help

Installation Guide

Installation from PyPI

You can easily install Colossal-AI with the following command:

pip install colossalai

By default, PyTorch extensions are not built during installation. If you need to build PyTorch extensions, you can set theBUILD_EXT=1::

BUILD_EXT=1 pip install colossalai

Additionally, we release NIGHTLY versions every week, allowing you to access the latest unreleased features and bug fixes. Installation is as follows:

pip install colossalai-nightly

Installation from source

git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI
pip install .

CUDA/C++ kernels are not compiled by default. colossalAI will build them at runtime. Enable CUDA kernel fusion if needed:

BUILD_EXT=1 pip install .

For CUDA 10.2 users, you can manually download the cub library and copy it to the appropriate directory before installing it.

Using Docker

Pulling Images from DockerHub

You can get the information directly from theDockerHub pagePull the Docker image.

Build your own image

cd ColossalAI
docker build -t colossalai ./docker

Starts the container in interactive mode:

docker run -ti --gpus all --rm --ipc=host colossalai bash

Functional operation flow

data parallelism

Data parallelism is the process of dividing a dataset into multiple subsets and training the model in parallel on multiple GPUs.ColossalAI makes it easy for users to train in data parallel with a simplified data parallelism profile:

from colossalai.nn.parallel import DataParallel
model = DataParallel(model)

tensor parallelism

Tensor parallelism is the process of dividing the parameter tensor of a model into multiple sub-tensors and computing them in parallel on multiple GPUs.ColossalAI provides implementations of 1D, 2D, 2.5D, and 3D tensor parallelism:

from colossalai.nn.parallel import TensorParallel
model = TensorParallel(model, parallel_mode='1D')

parallel assembly lines

Pipeline parallelism is the division of a model into multiple stages, each executed by one or more GPUs.ColossalAI provides easy pipeline parallelism configuration:

from colossalai.pipeline.parallel import PipelineParallel
model = PipelineParallel(model, num_stages=4)

Mixed precision training

Mixed-precision training significantly reduces video memory usage and speeds up training by using a combination of 16-bit floating point numbers (FP16) and 32-bit floating point numbers (FP32) during training:

from colossalai.amp import convert_to_amp
model, optimizer, criterion = convert_to_amp(model, optimizer, criterion)

Zero Redundancy Optimizer (ZeRO)

The ZeRO optimizer significantly reduces the graphics memory footprint by distributing the optimizer states, gradients, and parameters across multiple GPUs:

from colossalai.zero import ZeroOptimizer
optimizer = ZeroOptimizer(optimizer, model)

Real-world applications

Open-Sora

Open-Sora is ColossalAI's complete solution for video generation modeling, including model parameters, training details, and the ability to generate 16-second 720p HD videos with one click:

# 训练
python train.py
# 推理
python infer.py

For more information, please seeOpen-SoraThe

Colossal-LLaMA

Colossal-LLaMA provides an open-source solution for domain-specific Large Language Models (LLMs) that can achieve comparable results to mainstream large models with a small amount of training money:

# 训练
python train_llama.py
# 推理
python infer_llama.py

For more information, please seeColossal-LLaMAThe