总目录 › ColossalAI › 第 2 章
Chapter 02

安装与依赖：含 CUDA extension 编译

📌 commit 4f9953be335e · v0.5.0 PyPI / 源码 / Docker / NPU 四条路

ColossalAI 比纯 pip 安装的框架挑剔一些：自带 CUDA 扩展（fused kernel）、对 PyTorch / CUDA 版本有要求、 Docker 镜像是企业部署的标准方式。本章给你 4 条装机路径 + 一份故障 checklist。

2.1装前要承认的事

2.2关键依赖（来自 requirements/requirements.txt）

依赖版本约束说明
torch >=2.2.0,<=2.5.1 这个上下界钉死，新版 PyTorch 可能 break
transformers ==4.51.3 精确钉死
peft >=0.7.1,<=0.13.2 —
bitsandbytes >=0.39.0 QLoRA / 量化用
ray — ColossalChat / 推理用
diffusers ==0.29.0 Open-Sora / 文生图任务
galore_torch — GaLore 优化器
safetensors — checkpoint 格式
ninja — 编译 CUDA 扩展

对 Python 版本要求是 >= 3.7，实际推荐 3.10 / 3.11。

2.3四条安装路径

路线 A：PyPI（推荐第一次用）

# 1. 先装好对应 PyTorch
pip install torch==2.5.1 torchvision \
    --index-url https://download.pytorch.org/whl/cu124

# 2. 装 ColossalAI（默认不编译 CUDA 扩展）
pip install colossalai

# 3. （可选）想第一次跑就编完 CUDA 扩展
BUILD_EXT=1 pip install colossalai

不带 BUILD_EXT=1 时安装会很快（几分钟），但第一次跑训练会卡 5-15 分钟编译。 要么提前编译（生产场景），要么开发期容忍这个延迟。

路线 B：源码 + editable install

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

# 装上完整 requirements
pip install -r requirements/requirements.txt
pip install -r requirements/requirements-test.txt    # 跑测试时用

路线 C：nightly 版本

pip install colossalai-nightly

每周自动从 main 分支构建。用主仓最新特性（如 v0.5+ 的某些 plugin 更新）就靠它。

路线 D：Docker

# 拉官方镜像
docker pull hpcaitech/colossalai:latest

# 启动
docker run --gpus all -it --rm --shm-size=8g \
  -v $(pwd):/workspace \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  hpcaitech/colossalai:latest bash

# 镜像里已经预装：
# - CUDA 12.4 + cuDNN
# - PyTorch 2.4
# - ColossalAI 最新稳定版
# - flash-attn / TransformerEngine 等可选扩展

仓库 docker/ 目录提供 Dockerfile，企业部署时按需修改。

2.4BUILD_EXT 在编译什么

ColossalAI 自带的 CUDA 扩展在 extensions/ 目录，包括：

扩展用途
fused Adam fused optimizer step
fused LayerNorm NLP 模型用
fused softmax attention 用
scaled upper triang masked softmaxcausal attention
multi tensor apply 梯度裁剪等批量 op
fused RMSNorm LLaMA 系
autochunk / fp8 / inference kernel各类专用

JIT 编译机制：第一次调用时检查目标 op 是否已经编译；没编译就用 ninja 现编一份，缓存到 ~/.cache/colossalai/。

2.5验证安装

# 1. import 能成功
python -c "import colossalai; print(colossalai.__version__)"

# 2. 看主要 API 能 import
python -c "
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin
print('Plugins OK')
"

# 3. 跑 colossalai CLI
colossalai check -i

# 4. （耗时）编译并 import 一个 CUDA op
python -c "
import torch
from colossalai.kernel.kernel_loader import FusedAdamOpLoader
op = FusedAdamOpLoader().load()
print('Fused Adam loaded')
"

第 3 条 colossalai check -i 会打印环境完整报告（PyTorch / CUDA / NCCL / extensions 状态），出 issue 时第一手数据。

2.6Linux / macOS / Windows 支持矩阵

OS支持说明
Ubuntu 20.04+ / RHEL 8+ ✅ 一等公民 所有 example 都在这测过
WSL2 (Ubuntu) ✅ 能用 性能略损于裸机
macOS ❌ 无 NVIDIA GPU，开发期可装但不能训
Windows native ❌ CUDA 扩展编译路径不通

2.7NPU / ROCm 支持

v0.5+ 加入 Ascend NPU 适配（看 colossalai/accelerator/）：

from colossalai.accelerator import set_accelerator
set_accelerator("npu")    # 用昇腾

需要单独装 CANN + torch_npu。AMD ROCm 暂未官方支持。

2.8常见安装错误

现象原因修复
ImportError: cannot import name 'XXX' from 'transformers' transformers 版本不是 4.51.3 pip install transformers==4.51.3
第一次跑训练卡很久 JIT 编译 CUDA 扩展中 正常 5-15 分钟；下次启动会快
JIT 编译报 nvcc not found 没装 CUDA toolkit 装 cuda-toolkit-12-4；或 export CUDA_HOME=/usr/local/cuda
RuntimeError: GPU compute capability < 7.0 GPU 太旧 P100 / K80 等不支持，换 V100 / A100 / RTX20+
装 nightly 后 import 报错 nightly 偶发 bug 回退到 stable：pip install colossalai==0.5.0
RuntimeError: ninja is required 缺 ninja pip install ninja
多 GPU 启动报 NCCL 错 没设 CUDA_DEVICE_MAX_CONNECTIONS=1 等 看官方 examples 里的环境变量

2.9装完后第一件事

跑仓库自带 examples/language/llama/finetune.py 的 demo（单卡能跑通）：

cd ColossalAI/examples/language/llama
colossalai run --nproc_per_node 1 finetune.py \
  --model_type meta-llama/Llama-3.2-1B \
  --plugin gemini \
  --batch_size 2 \
  --max_length 512

看到 loss 在降 = 安装成功。这种最小例子也是把 JIT 编译触发完的好时机，之后的训练就不会卡。

2.10这章带走的