ColossalAI

环境要求

• PyTorch >= 2.2
• Python >= 3.7
• CUDA >= 11.0
• NVIDIA GPU 计算能力 >= 7.0（V100/RTX20 及以上）
• Linux 操作系统（暂不支持其他系统）

安装与启动步骤

1. 1确认环境满足要求

   安装前先核对 PyTorch、Python、CUDA 版本和 GPU 计算能力，并确认使用 Linux 系统，否则可能安装失败。

2. 2从 PyPI 安装

   最简安装方式。默认安装过程不编译 PyTorch 扩展，CUDA kernel 会在实际用到时于运行时构建。

   pip install colossalai

   复制

3. 3安装并编译扩展

   如需在安装阶段就构建 PyTorch/CUDA 扩展（使用 fused optimizer 时必需），加上 BUILD_EXT=1 前缀。

   BUILD_EXT=1 pip install colossalai

   复制

4. 4安装夜间构建版

   官方每周向 PyPI 发布 nightly 版，可提前使用主分支上尚未正式发布的新特性和修复。

   pip install colossalai-nightly

   复制

5. 5从源码安装

   进入已获取的 ColossalAI 源码目录后执行本地安装；默认同样不编译 CUDA/C++ kernel。

   cd /your/path/ColossalAI
   pip install .

   复制

6. 6源码安装并编译

   源码方式下若需要启用 CUDA kernel 融合，用 BUILD_EXT=1 安装；CUDA 10.2 用户还需手动下载 cub 库放到对应目录。

   BUILD_EXT=1 pip install .

   复制

关键配置

如何确认成功

README 未提供专门的验证命令，安装过程无报错即可；可访问官网文档站查看入门示例运行情况。

常见问题

注意事项

• 安装仅支持 Linux，且必须有计算能力 7.0 及以上的 NVIDIA GPU
• 默认不在安装期编译 CUDA/C++ kernel，首次运行相关功能时构建会额外耗时
• 使用 fused optimizer 时必须预先以 BUILD_EXT=1 方式安装
• 也可从 DockerHub 的 hpcaitech/colossalai 镜像直接使用（README 未给出具体命令，请到 DockerHub 页面查看）