# 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](#troubleshooting).

---

## Table of Contents

1. [Quick Start](#1-quick-start)
2. [Requirements](#2-requirements)
3. [Installation](#3-installation)
4. [Launching Distributed Jobs](#4-launching-distributed-jobs)
5. [The Booster Training Workflow](#5-the-booster-training-workflow)
6. [Choosing a Parallelism Plugin](#6-choosing-a-parallelism-plugin)
7. [Optimizers](#7-optimizers)
8. [Saving and Loading Checkpoints](#8-saving-and-loading-checkpoints)
9. [Mixed Precision, LoRA, and Memory Tricks](#9-mixed-precision-lora-and-memory-tricks)
10. [Accelerated Inference](#10-accelerated-inference)
11. [Applications](#11-applications)
12. [Common Tasks (Recipes)](#12-common-tasks-recipes)
13. [Troubleshooting](#13-troubleshooting)
14. [FAQ](#14-faq)
15. [Command & API Reference](#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.

```python
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:

```bash
colossalai run --nproc_per_node 4 train.py
```

```text
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

| Requirement | Minimum |
|---|---|
| Operating system | Linux (Windows via WSL only) |
| Python | >= 3.7 |
| PyTorch | >= 2.2 |
| CUDA | >= 11.0 |
| NVIDIA GPU compute capability | >= 7.0 (V100 / RTX 20-series and newer) |

> [!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).

---

## 3. Installation

### Install from PyPI

=== "Stable (no kernel build)"

    ```bash
    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)"

    ```bash
    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"

    ```bash
    pip install colossalai-nightly
    ```

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

### Install from source

```bash
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 .
```

```text
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](hpcaitech__ColossalAI__AdminGuide.md) for the exact commands and build internals.

### Verify your installation

```bash
colossalai check -i
```

```text
#### 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

```bash
# 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"

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

=== "By hostfile"

    A hostfile lists one hostname per line:

    ```text
    worker-0
    worker-1
    worker-2
    ```

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

=== "Subset of hostfile"

    ```bash
    # 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:

| Launcher | Use when started by | Reads from environment |
|---|---|---|
| `colossalai.launch_from_torch()` | `colossalai run` / `torchrun` | `RANK`, `LOCAL_RANK`, `WORLD_SIZE`, `MASTER_ADDR`, `MASTER_PORT` |
| `colossalai.launch_from_slurm(host, port)` | SLURM (`srun`) | `SLURM_PROCID`, `SLURM_NPROCS` |
| `colossalai.launch_from_openmpi(host, port)` | OpenMPI (`mpirun`) | `OMPI_COMM_WORLD_RANK`, `OMPI_COMM_WORLD_LOCAL_RANK`, `OMPI_COMM_WORLD_SIZE` |
| `colossalai.launch(rank, world_size, host, port)` | Manual / custom | You supply everything |

```python
# 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.

```python
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

| Method | What it does |
|---|---|
| `booster.boost(model, optimizer, criterion, dataloader, lr_scheduler)` | Wraps and returns all five objects with the plugin's strategy applied. Only `model` is required; the rest are optional. |
| `booster.backward(loss, optimizer)` | Runs the backward pass. **Use this instead of `loss.backward()`** — it routes through the optimizer so gradient scaling/sharding works. |
| `booster.execute_pipeline(data_iter, model, criterion, optimizer)` | Runs a fused forward+backward step. **Required** when your plugin uses pipeline parallelism (see below). |
| `booster.no_sync(model=..., optimizer=...)` | Context manager that skips gradient sync — use it for gradient accumulation. Supported by DDP and Low-Level ZeRO-1. |
| `booster.enable_lora(model, lora_config=...)` | Wraps the model with LoRA adapters (uses Hugging Face PEFT). |
| `booster.save_model / load_model / save_optimizer / load_optimizer / save_lr_scheduler / load_lr_scheduler` | Checkpoint I/O that understands sharded/distributed state. See [Checkpoints](#8-saving-and-loading-checkpoints). |

### 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:

```python
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.

| Plugin | Strategy | Best for |
|---|---|---|
| `TorchDDPPlugin` | Plain PyTorch DistributedDataParallel | Models that fit on one GPU; pure data parallelism |
| `TorchFSDPPlugin` | PyTorch Fully Sharded Data Parallel | Sharded data parallelism using native PyTorch FSDP |
| `LowLevelZeroPlugin` | ZeRO stage 1 or 2 (shard optimizer / + gradients) | Medium models; cut optimizer-state memory with little code change |
| `GeminiPlugin` | ZeRO-3 with heterogeneous (CPU/NVMe) memory management | Very large models on limited GPU memory; automatic offload |
| `HybridParallelPlugin` | Tensor + pipeline + sequence + ZeRO, composed | Largest models; full 3D/4D parallelism for transformers |
| `MoeHybridParallelPlugin` | Hybrid parallelism + expert parallelism | Mixture-of-Experts models |

### TorchDDPPlugin — start here

```python
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

```python
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

```python
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

```python
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

```python
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](hpcaitech__ColossalAI__AdminGuide.md).

---

## 7. Optimizers

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

| Optimizer | Notes |
|---|---|
| `HybridAdam` | Adam that runs on **both CPU and GPU** — the recommended optimizer for Gemini/ZeRO with offload. |
| `CPUAdam` | Adam computed on CPU; pairs with CPU offload. |
| `FusedAdam`, `FusedLAMB`, `FusedSGD` | Fused-kernel versions (need `BUILD_EXT=1` or runtime build). |
| `Lamb`, `Lars` | Large-batch optimizers. |
| `GaLoreAdamW`, `GaLoreAdamW8bit`, `GaLoreAdafactor` | GaLore low-rank gradient projection (memory-efficient). |
| `Adafactor`, `CAME` | Memory-efficient adaptive optimizers. |
| `DistributedLamb`, `DistributedCAME`, `DistributedAdaFactor`, `DistGaloreAwamW` | Distributed (tensor-parallel-aware) variants. |

```python
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.

```python
# 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")
```

| Argument | Meaning |
|---|---|
| `shard=True` | Write a multi-file sharded checkpoint (a directory). `shard=False` writes a single file. |
| `size_per_shard` | Max shard size in MB (default `1024`). Only used when `shard=True`. |
| `use_safetensors=True` | Save in the `.safetensors` format (requires the `safetensors` package). |
| `use_async=True` | Save asynchronously to overlap I/O with compute. |

> [!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:

```python
booster = Booster(mixed_precision="bf16")
```

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

### LoRA fine-tuning

```python
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

```python
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](hpcaitech__ColossalAI__AdminGuide.md).

---

## 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.

```python
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:

| Field | Default | Purpose |
|---|---|---|
| `max_batch_size` | `8` | Maximum concurrent requests. |
| `max_input_len` / `max_output_len` | `256` / `256` | Prompt and generation length caps. |
| `dtype` | `torch.float16` | Compute dtype (`"fp16"`, `"bf16"`, `"fp32"`). |
| `tp_size` / `pp_size` | `1` / `1` | Tensor / pipeline parallel degree for inference. |
| `do_sample`, `temperature`, `top_k`, `top_p` | `False`, `1.0`, `50`, `1.0` | Sampling controls. |
| `use_cuda_graph` | `False` | Capture CUDA graphs to cut kernel-launch overhead. |
| `use_spec_dec` | `False` | Enable speculative decoding. |
| `enable_streamingllm` | `False` | StreamingLLM attention for very long conversations. |

The complete `InferenceConfig` field list is in the [Admin Guide](hpcaitech__ColossalAI__AdminGuide.md).

---

## 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`:

```bash
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

```python
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)
```

```bash
colossalai run --nproc_per_node 8 train.py
```

### Combine tensor + pipeline parallelism for a transformer

```python
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

```python
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

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

### Run with Docker

```bash
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](hpcaitech__ColossalAI__AdminGuide.md) 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`

| Command | Purpose |
|---|---|
| `colossalai run [OPTIONS] SCRIPT [ARGS]` | Launch distributed training on one or more nodes. |
| `colossalai check -i` | Print an installation and environment compatibility report. |

#### `colossalai run` options

| Option | Type | Default | Description |
|---|---|---|---|
| `-H, -host, --host` | str | `None` | Comma-separated hostnames to launch on (`host1,host2`). |
| `--hostfile` | str | `None` | Path to a hostfile (one hostname per line). |
| `--include` | str | `None` | Only use these hosts (with `--hostfile`). Mutually exclusive with `--exclude`. |
| `--exclude` | str | `None` | Do not use these hosts (with `--hostfile`). |
| `--num_nodes` | int | `-1` | Total worker nodes to use (with `--hostfile`). |
| `--nproc_per_node` | int | `None` | Number of GPUs (processes) per node. |
| `--master_port` | int | `29500` | Port for `torch.distributed` communication. |
| `--master_addr` | str | `127.0.0.1` | IP of node 0. Inferred via `hostname -I` if unset. |
| `--extra_launch_args` | str | `None` | Extra torch launcher args, e.g. `arg1=1,arg2=2` → `--arg1=1 --arg2=2`. |
| `--ssh-port` | int | `None` | SSH port for multi-node connections. |
| `-m` | str | `None` | Run a library module as a script (terminates the option list). |

### Public Python API

| Symbol | Description |
|---|---|
| `colossalai.launch(rank, world_size, host, port, backend="nccl", local_rank=None, seed=1024, verbose=True)` | Initialize the distributed environment manually. |
| `colossalai.launch_from_torch(backend="nccl", seed=1024, verbose=True)` | Initialize from torchrun/`colossalai run` environment variables. |
| `colossalai.launch_from_slurm(host, port, backend="nccl", seed=1024, verbose=True)` | Initialize from SLURM environment variables. |
| `colossalai.launch_from_openmpi(host, port, backend="nccl", seed=1024, verbose=True)` | Initialize from OpenMPI environment variables. |
| `colossalai.booster.Booster(device=None, mixed_precision=None, plugin=None)` | High-level training wrapper. |
| `colossalai.booster.plugin.*` | `TorchDDPPlugin`, `TorchFSDPPlugin`, `LowLevelZeroPlugin`, `GeminiPlugin`, `HybridParallelPlugin`, `MoeHybridParallelPlugin`. |
| `colossalai.nn.optimizer.*` | `HybridAdam`, `CPUAdam`, `FusedAdam`, `FusedLAMB`, `FusedSGD`, `Lamb`, `Lars`, `Adafactor`, `CAME`, GaLore + distributed variants. |
| `colossalai.inference.InferenceEngine` / `InferenceConfig` | Accelerated inference. |

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

---

_Generated June 2026 from hpcaitech/ColossalAI at 4f9953b. AI-generated documentation. The library is checked by sampling (three repositories so far); this guide was not individually reviewed. Verify against the project before relying on it. Unofficial; not affiliated with or reviewed by the project._