hpcaitech/ColossalAI · 上手攻略

• 仓库：hpcaitech/ColossalAI
• 链接：https://github.com/hpcaitech/ColossalAI
• 分类：ai（llm-infra）
• 作者：Jay
• 更新：2026-07-11

是什么

ColossalAI 是 HPC-AI Tech 开源的分布式深度学习训练框架，核心目标是让大模型训练更便宜、更快、更普惠。它整合了多维并行（数据并行、流水线并行、张量并行）、ZeRO 优化、异构训练、FP8 低精度等技术，支持从单卡到大规模 GPU 集群的各类场景。

ColossalAI 最著名的成就是支撑了开源视频生成项目 Open-Sora 的训练——让研究者用远低于闭源方案的成本训练 Sora 类视频模型。2024 年 9 月，母公司 HPC-AI Tech 完成了 5000 万美元 A 轮融资。

解决什么问题

• 训练成本高：FP8 混合精度、ZeRO 显存优化让 GPU 利用率大幅提升，官方称可节省 30%+ 训练成本
• 大模型并行困难：3D 并行（数据 + 流水线 + 张量）开箱即用，配置文件中声明即可
• 视频生成训练门槛高：Open-Sora 生态让普通团队也能训练自己的视频生成模型
• 多框架兼容：支持 PyTorch、PyTorch Lightning，还可通过 lightning-colossalai 插件无缝接入 Lightning 生态

快速安装

# 基础安装（不含 PyTorch 扩展编译）
pip install colossalai

# 安装时同时编译 PyTorch C++/CUDA 扩展（推荐）
BUILD_EXT=1 pip install colossalai

# 从源码安装
git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI
pip install -r requirements/requirements.txt
BUILD_EXT=1 pip install .

⚠️ 依赖：需要 CUDA 11.6+、PyTorch ≥1.10。编译扩展需要 nvcc（CUDA 编译器）在 PATH 中。若使用 CUDA 10.2，需要手动下载 CUB 库（详见官方文档）。

核心用法

1. 单 GPU 训练（最低门槛）

import torch
import colossalai

# 初始化分布式环境（单 GPU 时 auto 设为主动检测）
colossalai.launch(backend='nccl')

# 定义模型（以 GPT-2 为例）
model = GPT2LMModel()

# 定义优化器
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# 构造 engine
engine, optimizer, _, _, _ = colossalai.initialize(
    model, optimizer, None, None, None
)

# 训练循环
for data in dataloader:
    engine.zero_grad()
    output = engine(data)
    loss = output.loss
    engine.backward(loss)
    engine.step()

2. 多维并行配置

在 config.py 或 YAML 中声明并行策略：

from colossalai.zero.zero_ctx import ZeroInitContext

# ZeRO-2 配置（显存优化）
zero_config = dict(
    model_params_tag='zero2_optim',  # 优化器状态分片
)

# 3D 并行配置示例
parallel_config = dict(
    tensor_parallel_size=2,    # 张量并行
    pipeline_parallel_size=2,  # 流水线并行
    data_parallel_size=4,     # 数据并行
)

colossalai.launch(config=parallel_config, ...)

⚠️ 注意：ColossalAI 并行配置较为复杂，建议从单卡开始跑通流程，再逐步增加并行维度。张量并行需要模型支持切分（不是所有层都能 TP 切分）。

3. 使用 PyTorch Lightning 集成

pip install lightning-colossalai

import lightning as L
from lightning.pytorch.strategies import ColossalAIStrategy

strategy = ColossalAIStrategy(
    tensor_parallel_size=2,
    zero_optimization='zero2'
)

trainer = L.Trainer(
    strategy=strategy,
    max_epochs=3,
    devices=4,
)

4. 使用预训练模型（示例：LLaMA）

from colossalai.booster import Booster
from colossalai.booster.plugin import TorchDDPPlugin, LowLevelZeroPlugin

# 选择插件
plugin = LowLevelZeroPlugin(stage=2)

booster = Booster(plugin=plugin)

model, optimizer, dataloader, criterion = booster.enable(
    model=model,
    optimizer=optimizer,
    dataloader=dataloader,
    criterion=criterion,
)

5. 视频生成（Open-Sora）

ColossalAI 官方提供 Open-Sora 训练脚本，核心思想是：

git clone https://github.com/hpcaitech/Open-Sora.git
cd Open-Sora
pip install -r requirements.txt
BUILD_EXT=1 pip install colossalai

# 运行训练（需多卡）
torchrun --nproc_per_node=8 scripts/train.py --config configs/xxx.py

典型适用场景

• 大模型预训练：LLaMA、GPT、Bloom 等模型的高效分布式预训练
• 视频生成：Open-Sora 生态，支持单次生成 16 秒 720p HD 视频
• 多模态训练：图像、语言混合训练
• RLHF / SFT：完整的 RLHF 流程支持（ColossalChat 方案）
• 跨租户云端训练：官方还提供 HPC-AI Cloud 托管环境，免去环境配置烦恼

坑与注意

1. 编译失败：默认 pip install colossalai 不编译 CUDA 扩展，很多优化功能不可用。强烈建议使用 BUILD_EXT=1 pip install colossalai
2. CUDA 10.2 特殊处理：需要手动下载 CUB 库并复制到指定路径，官方文档有详细步骤
3. ZeRO 与流水线并行：ColossalAI 中 ZeRO 与流水线并行组合有时不稳定，建议分开测试后再组合
4. 文档质量参差：核心概念文档较清晰，但某些新功能（如 FP8）的文档更新不及时，遇到问题多看 GitHub Issues
5. 多节点需要 MPI/NCCL：多节点训练需要配置好 NCCL（GPU 间通信），对于新用户有一定门槛
6. 版本更新较快：ColossalAI 仍在活跃开发中，API 变化较大（v1.x 相比 v0.x 有较大改动），生产环境建议锁定版本

与同类对比

结论：如果你要做视频生成（Open-Sora 路线），ColossalAI 是最顺滑的选择。如果主要是训练语言大模型，DeepSpeed 生态更成熟、文档更完善。两者 ZeRO 技术互通，可以视为互补而非互斥。

一句话推荐结论

ColossalAI 是大模型训练领域的「多面手」，尤其在视频生成和异构训练场景上有独特优势——如果你在训练 Sora 类模型或需要多维并行，ColossalAI 的开箱即用配置值得一试。