Skip to main content

Setup

Requirements:

If you encounter any problem about installation, you may want to raise an issue in this repository.

Download From PyPI​

You can install Colossal-AI with

pip install colossalai

Note: only Linux is supported for now

If you want to build PyTorch extensions during installation, you can use the command below. Otherwise, the PyTorch extensions will be built during runtime.

BUILD_EXT=1 pip install colossalai

Download From Source​

The version of Colossal-AI will be in line with the main branch of the repository. Feel free to raise an issue if you encounter any problem.

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

# install dependency
pip install -r requirements/requirements.txt

# install colossalai
BUILD_EXT=1 pip install .

If you don't want to install and enable CUDA kernel fusion (compulsory installation when using fused optimizer), just don't specify the BUILD_EXT:

pip install .

For Users with CUDA 10.2, you can still build ColossalAI from source. However, you need to manually download the cub library and copy it to the corresponding directory.

# clone the repository
git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI

# download the cub library
wget https://github.com/NVIDIA/cub/archive/refs/tags/1.8.0.zip
unzip 1.8.0.zip
cp -r cub-1.8.0/cub/ colossalai/kernel/cuda_native/csrc/kernels/include/

# install
BUILD_EXT=1 pip install .