Colossal-AI User Guide

Colossal-AI is a unified deep learning system for large-scale parallel training and inference. It lets you scale a model that runs on your laptop to a multi-GPU cluster by changing only a few lines of code. Instead of rewriting your training loop for every parallelism strategy, you pick a plugin (a parallelism recipe) and wrap your model, optimizer, dataloader, and scheduler with a single Booster object.

This guide shows you how to install Colossal-AI, launch distributed jobs, train models with each parallelism strategy, run accelerated inference, and use the higher-level applications (chat/RLHF, domain LLM fine-tuning, MoE, evaluation, and retrieval-augmented Q&A).

Note

Colossal-AI runs on Linux only. The installer raises an error on Windows. If you are on Windows, use the Windows Subsystem for Linux (WSL). See Troubleshooting.

Table of Contents

1. Quick Start
2. Requirements
3. Installation
4. Launching Distributed Jobs
5. The Booster Training Workflow
6. Choosing a Parallelism Plugin
7. Optimizers
8. Saving and Loading Checkpoints
9. Mixed Precision, LoRA, and Memory Tricks
10. Accelerated Inference
11. Applications
12. Common Tasks (Recipes)
13. Troubleshooting
14. FAQ
15. Command & API Reference

1. Quick Start

The smallest possible Colossal-AI training script. It initializes the distributed environment, wraps your model in a Booster with a plugin, and runs a normal-looking PyTorch loop.

import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin import TorchDDPPlugin
from colossalai.nn.optimizer import HybridAdam

# 1. initialize the distributed environment (reads RANK/WORLD_SIZE/MASTER_ADDR from env)
colossalai.launch_from_torch()

# 2. build your model, optimizer, criterion, dataloader, scheduler as usual
model = MyModel()
optimizer = HybridAdam(model.parameters(), lr=1e-3)
criterion = lambda outputs, targets: torch.nn.functional.cross_entropy(outputs, targets)
dataloader = ...
lr_scheduler = ...

# 3. pick a plugin and create the Booster
plugin = TorchDDPPlugin()
booster = Booster(plugin=plugin)

# 4. boost everything — this injects the parallelism strategy
model, optimizer, criterion, dataloader, lr_scheduler = booster.boost(
    model, optimizer, criterion, dataloader, lr_scheduler
)

# 5. train with a standard loop, but use booster.backward(...) instead of loss.backward()
for epoch in range(10):
    for inputs, targets in dataloader:
        inputs, targets = inputs.cuda(), targets.cuda()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        booster.backward(loss, optimizer)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

Launch it on 4 GPUs of your current machine:

colossalai run --nproc_per_node 4 train.py

Distributed environment is initialized, world size: 4

Tip

The only structural difference from a single-GPU PyTorch script is steps 3–4 (choose a plugin, call booster.boost) and using booster.backward(loss, optimizer) instead of loss.backward(). To scale up, you change the plugin, not the loop.

2. Requirements

Important

The pip package installs without compiling any CUDA/C++ kernels by default. Fused kernels (required by fused optimizers such as FusedAdam) are either built at runtime on first use or built at install time when you set BUILD_EXT=1. See Installation.

3. Installation

Install from PyPI

Stable (no kernel build)

pip install colossalai

By default no PyTorch extensions are compiled. CUDA kernels are built lazily the first time you use a feature that needs them.

Stable (build kernels now)

BUILD_EXT=1 pip install colossalai

Compiles the fused CUDA/C++ kernels during installation. Use this if you want fused optimizers available immediately or want to fail fast on a broken toolchain.

Nightly

pip install colossalai-nightly

A weekly build that tracks main, giving you unreleased features and fixes.

Install from source

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

# install without building kernels (built lazily at runtime)
pip install .

# OR install and compile CUDA/C++ kernel fusion now
BUILD_EXT=1 pip install .

Successfully installed colossalai-0.5.0

Note

If you are on CUDA 10.2 you must manually download the NVIDIA cub library before building. See the Admin Guide for the exact commands and build internals.

Verify your installation

colossalai check -i

#### Installation Report ####
...
CUDA Version: ...
PyTorch Version: ...
CUDA extensions: ...

4. Launching Distributed Jobs

Colossal-AI ships a launcher CLI, colossalai run, that starts one process per GPU and wires up the torch.distributed environment variables your script reads with colossalai.launch_from_torch().

Single node

# 4 GPUs on this machine, default port 29500
colossalai run --nproc_per_node 4 train.py

# 4 GPUs at a custom port
colossalai run --nproc_per_node 4 --master_port 29550 train.py

Multiple nodes

By host list

colossalai run --host host1,host2 --master_addr host1 --nproc_per_node 4 train.py

By hostfile

A hostfile lists one hostname per line:

worker-0
worker-1
worker-2

colossalai run --hostfile ./hostfile --master_addr host1 --nproc_per_node 4 train.py

Subset of hostfile

# only these hosts
colossalai run --hostfile ./hostfile --master_addr host1 --include host1,host2 --nproc_per_node 4 train.py

# all hosts except these
colossalai run --hostfile ./hostfile --master_addr host1 --exclude host2 --nproc_per_node 4 train.py

Warning

--include and --exclude are mutually exclusive — using both prints an error and exits. Both are only effective together with --hostfile.

Initializing the environment inside your script

Pick the launch_* helper that matches how the job is started:

# torchrun / colossalai run
colossalai.launch_from_torch(seed=42)

# SLURM
colossalai.launch_from_slurm(host="node01", port=29500, seed=42)

Tip

Set seed= in the launch call to make every process use the same random seed. The default is 1024.

5. The Booster Training Workflow

The Booster is the central API. It takes your ordinary PyTorch objects and a plugin, and returns "boosted" versions that transparently apply the chosen parallelism, mixed precision, and checkpointing strategy.

from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin

plugin = GeminiPlugin(placement_policy="auto")
booster = Booster(plugin=plugin)

model, optimizer, criterion, dataloader, lr_scheduler = booster.boost(
    model, optimizer, criterion, dataloader, lr_scheduler
)

The methods you actually call

Pipeline-parallel loops are different

When your plugin enables pipeline parallelism (pp_size > 1), do not write the normal model(input) / booster.backward(...) pattern. Use execute_pipeline, which drives the forward and backward micro-batch schedule for you:

outputs = booster.execute_pipeline(
    data_iter,           # iter(dataloader) or iter([batch])
    model,
    criterion=lambda outputs, inputs: outputs.loss,
    optimizer=optimizer,
    return_loss=True,
    return_outputs=False,
)
loss = outputs["loss"]   # only populated on the last pipeline stage
optimizer.step()
optimizer.zero_grad()

Caution

Mixing the conventional loss.backward() path with a pipeline-parallel plugin causes incorrect results or runtime errors. Pick one: execute_pipeline for pipeline plugins, booster.backward for everything else.

6. Choosing a Parallelism Plugin

The plugin decides how your model is distributed. All plugins share the Booster interface, so swapping strategies is a one-line change. Pick based on model size and the kind of parallelism you need.

TorchDDPPlugin — start here

from colossalai.booster.plugin import TorchDDPPlugin
plugin = TorchDDPPlugin()

Pure data parallelism. Every GPU holds a full model replica. Good when the model fits comfortably in GPU memory and you just want more throughput.

LowLevelZeroPlugin — shrink optimizer memory

from colossalai.booster.plugin import LowLevelZeroPlugin
# stage 1 shards optimizer states; stage 2 also shards gradients
plugin = LowLevelZeroPlugin(stage=2, precision="bf16", cpu_offload=True)

Tip

cpu_offload=True moves optimizer state to CPU RAM, freeing GPU memory at the cost of some speed — a good first lever when you hit out-of-memory.

GeminiPlugin — train models bigger than one GPU

from colossalai.booster.plugin import GeminiPlugin
plugin = GeminiPlugin(
    placement_policy="auto",      # let Gemini decide what lives on GPU vs CPU
    precision="bf16",
    offload_optim_frac=1.0,       # offload all optimizer state to CPU
)

Gemini implements ZeRO-3 with a chunk-based memory manager that automatically moves parameters between GPU and CPU (and optionally NVMe) so you can train models far larger than a single GPU's memory.

HybridParallelPlugin — the full toolbox

from colossalai.booster.plugin import HybridParallelPlugin
plugin = HybridParallelPlugin(
    tp_size=2,                    # tensor parallel degree (required)
    pp_size=2,                    # pipeline parallel degree (required)
    zero_stage=1,
    precision="bf16",
    enable_all_optimization=True, # turn on fused norm / flash attention / etc.
    num_microbatches=4,           # needed when pp_size > 1
)

tp_size × pp_size × (data-parallel degree, inferred from world size) must divide your total GPU count. With pp_size > 1, use execute_pipeline (see above) and set num_microbatches or microbatch_size.

Note

enable_all_optimization=True switches on Shardformer's transformer optimizations (fused normalization, flash attention, JIT fusion). It works best with standard Hugging Face transformer architectures.

MoeHybridParallelPlugin — Mixture of Experts

from colossalai.booster.plugin import MoeHybridParallelPlugin
plugin = MoeHybridParallelPlugin(
    tp_size=1, pp_size=1, ep_size=4,   # ep_size = expert-parallel degree (required)
    precision="bf16",
)

The full option table for every plugin (dozens of tuning knobs each) is in the Admin Guide → Plugin Configuration Reference.

7. Optimizers

Colossal-AI provides drop-in optimizers, several of which are required to unlock memory savings. Import them from colossalai.nn.optimizer.

from colossalai.nn.optimizer import HybridAdam
optimizer = HybridAdam(model.parameters(), lr=1e-3)

Important

When you use a fused optimizer such as FusedAdam, the fused CUDA kernel must be available. Either install with BUILD_EXT=1 or let Colossal-AI build it at runtime (the first step will pause to compile).

8. Saving and Loading Checkpoints

Always save and load through the Booster, not torch.save/torch.load — only the Booster understands sharded and distributed state produced by ZeRO/tensor-parallel plugins.

# save (sharded, HF-transformers-style folder)
booster.save_model(model, "ckpt/model", shard=True, size_per_shard=1024, use_safetensors=True)
booster.save_optimizer(optimizer, "ckpt/optimizer", shard=True)
booster.save_lr_scheduler(lr_scheduler, "ckpt/scheduler.pt")

# load
booster.load_model(model, "ckpt/model")
booster.load_optimizer(optimizer, "ckpt/optimizer")
booster.load_lr_scheduler(lr_scheduler, "ckpt/scheduler.pt")

Tip

Sharded + safetensors checkpoints are directly compatible with the Hugging Face ecosystem, making it easy to share or resume models.

9. Mixed Precision, LoRA, and Memory Tricks

Mixed precision

Most plugins control precision through their precision= argument ("fp16", "bf16", and on supported plugins "fp8"). If you are not using a plugin that controls precision, set it on the Booster directly:

booster = Booster(mixed_precision="bf16")

Valid Booster precision strings: "fp16" (PyTorch AMP), "fp16_apex" (NVIDIA Apex), "bf16", "fp8".

LoRA fine-tuning

from peft import LoraConfig
lora_config = LoraConfig(task_type="CAUSAL_LM", r=8, lora_alpha=32, lora_dropout=0.1)
model = booster.enable_lora(model, lora_config=lora_config)
# ... train ...
booster.save_lora_as_pretrained(model, "ckpt/lora_adapter")

Note

LoRA requires the Hugging Face peft library and a plugin that supports LoRA. You can also pass quantize=True to combine LoRA with 4-bit (QLoRA-style) quantization.

Gradient accumulation

with booster.no_sync(model=model):     # skip gradient all-reduce
    for micro_batch in accumulation_steps[:-1]:
        loss = criterion(model(micro_batch), targets)
        booster.backward(loss, optimizer)
# final micro-batch outside no_sync triggers the sync

For more memory-saving features — NVMe offload, gradient checkpointing, lazy initialization, sequence parallelism — see the Admin Guide.

10. Accelerated Inference

Colossal-Inference provides a high-throughput inference engine with continuous batching, paged KV-cache, tensor/pipeline parallelism, speculative decoding, and CUDA-graph capture.

from colossalai.inference import InferenceConfig, InferenceEngine
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

config = InferenceConfig(
    max_batch_size=8,
    max_input_len=256,
    max_output_len=256,
    dtype="fp16",
    tp_size=1,
    do_sample=True,
    temperature=0.8,
    top_k=50,
    top_p=0.95,
)

engine = InferenceEngine(model, tokenizer, config, verbose=True)
outputs = engine.generate(prompts=["Tell me about large-scale model training."])
print(outputs)

Frequently used InferenceConfig fields:

The complete InferenceConfig field list is in the Admin Guide.

11. Applications

Beyond the core library, the applications/ directory ships complete, runnable solutions. Each has its own requirements.txt.

ColossalChat — ChatGPT-style RLHF

A full RLHF pipeline (applications/ColossalChat, package coati) supporting supervised fine-tuning and several preference-optimization algorithms:

• SFT (supervised fine-tuning)
• PPO (reward model + policy optimization)
• GRPO (group relative policy optimization, including a distributed Ray-based reasoning pipeline in coati/distributed)
• DPO, KTO (direct/Kahneman-Tversky preference optimization)

The distributed reinforcement-learning entry point is rl_example.py:

python rl_example.py \
  --model Qwen/Qwen2.5-7B \
  --dataset data.jsonl \
  -t 4 -i 4 \                 # 4 trainer workers, 4 inferencer workers
  -b vllm \                   # inference backend
  -a GRPO \                   # algorithm
  -ibs 8 -tbs 8 \             # inference / train batch sizes
  -e 1                        # episodes

Note

ColossalChat's RL training uses Ray to coordinate separate trainer and inferencer pools. Install its requirements (pip install -r applications/ColossalChat/requirements.txt), which include ray, transformers, vllm, and fastapi.

Colossal-LLaMA — domain LLM pretraining & fine-tuning

applications/Colossal-LLaMA provides scripts to continue-pretrain and fine-tune LLaMA-family models on your own corpus, including tokenizer extension and dataset preparation utilities.

ColossalMoE — Mixture-of-Experts training

applications/ColossalMoE demonstrates training MoE models (e.g. Mixtral) with expert parallelism via MoeHybridParallelPlugin.

ColossalEval — model evaluation

applications/ColossalEval runs dataset-based and GPT-graded evaluation across standard benchmarks (MMLU, CMMLU, AGIEval, GAOKAO, CEval, and more).

ColossalQA — retrieval-augmented Q&A

applications/ColossalQA builds a LangChain-based RAG system over your documents, with a web UI demo (examples/webui_demo).

Tip

Ready-made model training examples (LLaMA 1/2/3, GPT, BERT, OPT, PaLM, Grok-1, Mixtral, DeepSeek, ViT, Stable Diffusion, DreamBooth) live under examples/. They are the fastest way to see a complete, working configuration for a specific model.

12. Common Tasks (Recipes)

Train a model too big for one GPU

import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin
from colossalai.nn.optimizer import HybridAdam

colossalai.launch_from_torch()
plugin = GeminiPlugin(placement_policy="auto", precision="bf16", offload_optim_frac=1.0)
booster = Booster(plugin=plugin)

model = build_large_model()
optimizer = HybridAdam(model.parameters(), lr=1e-4)
model, optimizer, *_ = booster.boost(model, optimizer)

colossalai run --nproc_per_node 8 train.py

Combine tensor + pipeline parallelism for a transformer

from colossalai.booster.plugin import HybridParallelPlugin
plugin = HybridParallelPlugin(tp_size=2, pp_size=2, zero_stage=1,
                              precision="bf16", num_microbatches=8,
                              enable_all_optimization=True)

Drive training with booster.execute_pipeline(...) because pp_size > 1.

Resume training from a sharded checkpoint

booster.load_model(model, "ckpt/model")
booster.load_optimizer(optimizer, "ckpt/optimizer")
booster.load_lr_scheduler(lr_scheduler, "ckpt/scheduler.pt")

Run multi-node training over a hostfile

colossalai run --hostfile ./hostfile --master_addr worker-0 --nproc_per_node 8 train.py

Run with Docker

cd ColossalAI
docker build -t colossalai ./docker
docker run -ti --gpus all --rm --ipc=host colossalai bash

Warning

Building the image requires the NVIDIA Container Runtime to be the default Docker runtime (the build compiles CUDA kernels). See the Admin Guide for details.

13. Troubleshooting

Error: Windows is not supported yet. Please try again within the Windows Subsystem for Linux (WSL).

• Cause: You are installing on native Windows. Colossal-AI's installer (setup.py) explicitly blocks Windows.
• Solution: Install and run inside WSL2 with a CUDA-capable Linux distribution, or use a Linux machine / container.
• Prevention: Use the provided Docker image, which is Linux-based.

Error: [extension] PyTorch is not found while BUILD_EXT=1. You need to install PyTorch first in order to build CUDA extensions

• Cause: You set BUILD_EXT=1 but PyTorch is not installed in the active environment.
• Solution: pip install torch (matching your CUDA version) before installing Colossal-AI with BUILD_EXT=1.
• Prevention: Install PyTorch first, or omit BUILD_EXT=1 and let kernels build lazily at runtime.

Error: [extension] Could not find any kernel compatible with the current environment.

• Cause: No prebuilt kernel matches your CUDA/GPU; the build environment is incomplete.
• Solution: Ensure ninja and a matching CUDA toolkit are installed, then reinstall with BUILD_EXT=1. Set TORCH_CUDA_ARCH_LIST to your GPU's architecture if auto-detection fails.
• Prevention: Use the official Docker image.

Error: Could not find {VAR} in the torch environment ... (or SLURM/OpenMPI variant)

• Cause: You called launch_from_torch()/launch_from_slurm()/launch_from_openmpi() but the expected environment variables are not set — usually because the script was started directly with python instead of through colossalai run/torchrun/srun/mpirun.
• Solution: Launch the script with the matching launcher (e.g. colossalai run --nproc_per_node N train.py).
• Prevention: Match the launch_from_* helper to your launcher.

Error: Cannot find index file in {checkpoint} / No weigth is loaded into the model.

• Cause: Loading a sharded checkpoint directory that is missing its index file, or pointing at the wrong path.
• Solution: Point load_model at the directory produced by save_model(..., shard=True), which contains the index file. For single-file checkpoints, pass the file path.
• Prevention: Always load with the same shard semantics you used to save.

Error: device type must be cuda for nccl backend

• Cause: Pipeline parallelism point-to-point communication ran on a non-CUDA tensor with the NCCL backend.
• Solution: Move data to CUDA before the step and ensure your accelerator is a GPU; NCCL requires CUDA tensors.

Error: AssertionError: torch.distributed is not initialized, please use colossalai.launch to create the distributed environment

• Cause: You created a plugin/Booster before initializing the distributed environment.
• Solution: Call colossalai.launch_from_torch() (or another launch_*) before building plugins.

Out-of-memory (CUDA OOM)

• Solutions, in order of escalation:
  1. Switch from TorchDDPPlugin to LowLevelZeroPlugin(stage=2).
  2. Add cpu_offload=True.
  3. Move to GeminiPlugin with offload_optim_frac/offload_param_frac.
  4. Add tensor/pipeline parallelism via HybridParallelPlugin.
  5. Use bf16 precision and a memory-efficient optimizer (Adafactor, CAME, GaLore).

Important

Several open issues track plugin-combination edge cases (e.g. certain optimizers with LowLevelZeroPlugin + small initial_scale + CPU offload, or use_fp8=True shape mismatches). If you hit an assertion deep inside a plugin, check the project's GitHub issues for your exact combination before assuming a bug in your code.

14. FAQ

Q: Do I have to rewrite my training loop to use Colossal-AI? No. Keep your PyTorch loop. Add colossalai.launch_from_torch(), wrap your objects with booster.boost(...), and call booster.backward(loss, optimizer) instead of loss.backward(). To change parallelism, change the plugin.

Q: Which plugin should I start with? Start with TorchDDPPlugin if the model fits on a GPU. Move to LowLevelZeroPlugin when optimizer memory is tight, GeminiPlugin when the model itself doesn't fit, and HybridParallelPlugin for the largest transformer models.

Q: What's the difference between colossalai.launch and Booster? launch sets up the distributed process group (one-time, at startup). Booster wraps your model/optimizer with the parallelism strategy (per training script). You use both.

Q: Do I need to build CUDA kernels? Not for basic use — they build lazily at runtime. Build them at install time with BUILD_EXT=1 if you want fused optimizers immediately or want to detect toolchain problems early.

Q: Can I run on Windows? Only through WSL2. Native Windows is blocked by the installer.

Q: How do I save a checkpoint that Hugging Face tools can read? Use booster.save_model(model, path, shard=True, use_safetensors=True) — it writes an HF-transformers-style sharded safetensors folder.

Q: Which precisions are supported? fp16, bf16, and fp8 (plugin-dependent), plus Apex fp16 via "fp16_apex" on the Booster.

Q: How do I do LoRA / QLoRA fine-tuning? booster.enable_lora(model, lora_config=...) with a peft.LoraConfig; add quantize=True for 4-bit QLoRA-style training.

Q: What models have ready-made examples? LLaMA 1/2/3, GPT, BERT, OPT, PaLM, Grok-1, Mixtral, DeepSeek, ViT, Stable Diffusion, and DreamBooth — under examples/.

Q: How is inference different from training? Use InferenceEngine + InferenceConfig from colossalai.inference, not the Booster. It adds continuous batching, paged KV-cache, and speculative decoding.

Q: Is there a CLI to check my environment? Yes: colossalai check -i prints an installation/compatibility report.

15. Command & API Reference

CLI: colossalai

colossalai run options

Public Python API

For architecture, the complete configuration reference, deployment, security, and internals, see the Admin Guide.