尧图网页设计 
Colossal-AI 安装指南PyPI / 源码安装全流程与 BUILD_EXT 内核扩展机制解析【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI本文是基于 Colossal-AI 官方中文安装文档docs/source/zh-Hans/get_started/installation.md编写的一站式安装实战指南。文章完整覆盖环境要求、PyPI 安装、源码安装两条主路径并结合当前仓库源码深入剖析BUILD_EXT1背后AOT 预编译 / JIT 运行时编译两种 CUDA 内核构建机制、版本自检命令与常见编译错误定位方法。读完本文你能够在自己的 Linux GPU 集群上正确装好 Colossal-AI并清楚知道装出来的版本预编译了哪些算子、何时会触发运行时编译、遇到版本不匹配如何排查。环境要求官方安装文档给出的最低环境门槛如下环境项最低要求说明操作系统Linux当前只支持 LinuxWindows 会被安装脚本直接拒绝Python 3.7见下方关于 PyTorch 版本约束的补充说明PyTorch 2.1满足张量并行、混合并行等核心功能所需的算子接口CUDA 11.0编译 CUDA 扩展时的工具链要求NVIDIA GPUCompute Capability 7.0即 Volta 及以上V1007.0/ RTX 20 系列7.5以及更新的 A1008.0、RTX 30 系列8.6、H1009.0等关于版本门槛有两点需要提醒Python 与 PyTorch 的组合要互相匹配。安装文档的门槛是Python 3.7但当前仓库主分支的 requirements/requirements.txt 已将 PyTorch 约束在torch2.2.0,2.5.1而新版 PyTorch 对 Python 版本有自身要求实际装机时应以你选择的 PyTorch 官方版本对应的 Python 支持范围为准。建议用 CUDA 兼容的 PyTorch 预编译包。CUDA 扩展编译时见下文会校验系统 CUDA 版本与PyTorch 内置 CUDA 版本是否一致因此先按自己驱动版本安装对应 cu 版本的 PyTorch再装 Colossal-AI是避免踩坑的关键顺序。如果你在安装过程中遇到问题可以向本项目的 Issue 区反馈附上环境信息与报错日志更易定位。安装方式总览先理解BUILD_EXT这个开关Colossal-AI 内部包含一批高性能 CUDA/C 算子如 fused optimizer、LayerNorm、MoE、FlashAttention、softmax、CPU Adam 等。这些算子有两种构建时机官方文档用BUILD_EXT环境变量来切换AOTAhead-of-Time预编译在pip install阶段把扩展编译好并打进安装产物。命令中显式设置BUILD_EXT1即开启JITJust-in-Time运行时编译不做预编译等代码第一次真正用到某个算子时在运行时用torch.utils.cpp_extension.load()现场编译并缓存.so之后复用。从 setup.py 源码可以看到这一逻辑的落地方式BUILD_EXT int(os.environ.get(BUILD_EXT, 0)) 1 if BUILD_EXT: if not TORCH_AVAILABLE: raise ModuleNotFoundError(...) # BUILD_EXT1 时必须已装好 PyTorch from extensions import ALL_EXTENSIONS for ext_cls in ALL_EXTENSIONS: ext ext_cls() if ext.support_aot and ext.is_available(): ext.assert_compatible() ext_modules.append(ext.build_aot()) else: ext_modules []也就是说开BUILD_EXT1的前提是当前 Python 环境里已经能 import torch并且扩展类只有在当前机器硬件满足条件is_available()且工具链兼容assert_compatible()时才会进入 AOT 编译列表。若一个可用内核都没有安装会直接报错Could not find any kernel compatible with the current environment。官方安装文档的原文提示也印证了这两种模式的分工不添加BUILD_EXT1时PyTorch 扩展会在运行时自动安装即 JIT 编译。因此两条安装路径各有用武之地想开箱即用、首次运行不等待编译 →BUILD_EXT1只想在纯 Python 层面使用、或构建机没有 GPU/不想预编译 → 走 JIT首次使用某算子时自动编译。从 PyPI 安装快速上手路径PyPI 上已发布预编译好的发行版直接安装pip install colossalai注该方式目前只支持 Linux。如需在安装阶段同时完成 PyTorch 扩展的 AOT 预编译避免运行时边用边编加上环境变量BUILD_EXT1 pip install colossalai需要注意当安装的是 PyPI 官方 wheel非源码本地构建时BUILD_EXT1能否真正触发预编译取决于该 wheel 的构建策略官方更推荐对扩展有强需求的用户走下面的源码安装。装完后建议执行 安装后自检 一节中的colossalai check -i报告中会明确显示当前安装产物是否包含 AOT 编译的 CUDA 扩展。从源码安装跟随主分支开发源码安装适用于需要体验主分支最新特性、需要针对自己的 GPU 架构重新编译内核、或想参与开发调试的场景。文档强调此章节与版本库主分支保持一致。基本安装流程git clone https://github.com/hpcaitech/ColossalAI.git cd ColossalAI # 安装依赖 pip install -r requirements/requirements.txt # 安装 colossalai并预编译 CUDA 扩展 BUILD_EXT1 pip install .其中 requirements/requirements.txt 列出了运行期与编译期依赖的完整清单numpy、tqdm、psutil、ninja、torch2.2.0,2.5.1、transformers、peft、ray、fastapi等pip install -r会先把这些装齐为后续的扩展编译与分布式训练做好准备。关闭 CUDA 内核预编译纯 JIT 模式如果你的环境不打算使用 CUDA 内核融合功能融合优化器如 FusedAdam/LAMB 会强制要求 CUDA 扩展此时 JIT 也会在运行到该算子时自动触发编译可以不加BUILD_EXT1pip install .在这种模式下setup.py中的ext_modules为空安装只做纯 Python 打包内核算子由 colossalai/kernel/kernel_loader.py 中的KernelLoader在运行时按需加载——先尝试 import 预编译模块colossalai._Cimport 失败则现场 JIT 编译详见 运行时加载机制。针对旧版 CUDA 10.2 的 cub 手动补齐步骤如果你的环境仍在使用 CUDA 10.2文档说明仍可从源码安装但需要手动下载 cub 库并复制到对应目录# 克隆仓库 git clone https://github.com/hpcaitech/ColossalAI.git cd ColossalAI # 下载 cub 库1.8.0 wget https://github.com/NVIDIA/cub/archive/refs/tags/1.8.0.zip unzip 1.8.0.zip cp -r cub-1.8.0/cub/ colossalai/kernel/cuda_native/csrc/kernels/include/ # 安装 BUILD_EXT1 pip install .需要特别说明适用前提CUDA 10.2 低于本文开头CUDA 11.0的通用门槛该段指令面向的是历史上较老版本的 Colossal-AI。从当前仓库主分支的目录结构看内核源码已统一收拢到 extensions/csrc 与 colossalai/kernel/extensions/csrc 等路径下文档中colossalai/kernel/cuda_native/csrc/kernels/include/这样的旧路径在当前代码布局中已不存在。如果你使用新版本源码请直接执行标准的BUILD_EXT1 pip install .无需手动放置 cub仅当你确实在维护 CUDA 10.2 时代的老版本代码时才需要参考这一段。源码级原理BUILD_EXT1到底在构建什么被纳入 AOT 编译的扩展清单在BUILD_EXT1时setup.py从extensions包导入ALL_EXTENSIONS并逐个构建。当前主分支 extensions/init.py 中注册的扩展类为CpuAdamArmExtension/CpuAdamX86ExtensionCPU Adam区分 ARM 与 x86LayerNormCudaExtensionMoeCudaExtensionMoE 相关内核FusedOptimizerCudaExtension融合优化器内核InferenceOpsCudaExtension推理算子ScaledMaskedSoftmaxCudaExtension/ScaledUpperTriangleMaskedSoftmaxCudaExtensionFlashAttentionDaoCudaExtension/FlashAttentionSdpaCudaExtension/FlashAttentionNpuExtension构建后的扩展模块统一落在colossalai._C这个预编译模块命名空间下对应仓库中的 colossalai/_C/init.py并可通过importlib.import_module在运行时被KernelLoader反查加载相关逻辑见 extensions/cpp_extension.py 中的_CppExtension.import_op()。AOT vs JIT运行时加载机制与缓存目录无论安装时是否预编译运行期统一入口都是 colossalai/kernel/kernel_loader.py 中的KernelLoader。其load()流程是实例化该 loader 注册表里的全部扩展类逐一调用is_available()过滤出当前机器可用的内核并调用assert_compatible()做兼容性校验多个可用内核时按priority降序取最高优先级者调用_Extension.load()——对应 extensions/cpp_extension.py 中先import_op()尝试导入colossalai._C.xxx预编译模块ImportError/ModuleNotFoundError时回退到build_jit()现场编译。JIT 编译产物并非散落在安装目录而是统一缓存在用户目录下。根据 extensions/base_extension.py 中的get_jit_extension_folder_path()缓存路径格式为~/.cache/colossalai/torch_extensions/torch主.次_设备名-设备版本-hash其中 hash 由colossalai.__file__路径经 SHA-256 生成缓存目录按 torch 版本 设备型号区分。因此同一次编译结果在后续运行、重启进程后可直接复用日志中会打印Loading the JIT-built ... kernel during runtime now或Compiling the JIT ... kernel during runtime now提示当前是加载缓存还是新编译。理解了这套机制也就不难明白官方文档那句话的含义不设BUILD_EXT1扩展会在运行时自动构建只是第一次触达某个算子的启动耗时会更长。安装后的环境自检使用colossalai check -i一键体检安装包通过setup.py的entry_points注册了colossalai命令路由定义在 colossalai/cli/cli.pyrun/check两个子命令。其中check子命令colossalai/cli/check/init.py支持-i/--installation参数会输出一份完整的安装体检报告colossalai check -i报告由 colossalai/cli/check/check_installation.py 生成主要字段含义报告区块字段含义与判读EnvironmentColossal-AI version当前 colossalai 版本号version.txt记录的仓库版本为0.5.0若安装时带了 CUDA 扩展版本号形如X.X.XtorchX.XXcuXX.XEnvironmentPyTorch version当前环境的 torch 版本EnvironmentSystem CUDA version由$CUDA_HOME/bin/nvcc -V解析出的系统 CUDA 版本显示 N/A 说明 CUDA_HOME 未配置EnvironmentCUDA version required by PyTorchtorch.version.cuda给出的 PyTorch 内置 CUDA 版本N/A 说明装的是无 CUDA 的 CPU 版 torchCUDA Extensions AOT CompilationFound AOT CUDA Extension当前安装产物是否带 AOT 预编译扩展CUDA Extensions AOT CompilationPyTorch / CUDA version used for AOT compilation预编译时使用的版本可从版本号后缀解析用于核对与当前运行环境的差异Compatibility三项兼容性对比PyTorch 版本是否与 AOT 编译版本一致系统 CUDA 与 PyTorch 内置 CUDA 是否同主版本系统 CUDA 与 AOT 编译所用 CUDA 是否匹配官方实现中_is_compatible()只要求主版本号major与次版本号minor严格一致patch 级差异被忽略这与编译校验逻辑保持一致。命令行快速冒烟体检之外还可以直接确认版本与导入是否正常python -c import colossalai; print(colossalai.__version__) colossalai run --help # 查看官方多机/单机启动器用法run子命令由 colossalai/cli/launcher/run.py 提供可用于后续启动分布式训练脚本前的环境联调。编译期的硬性检查与常见问题定位CUDA 扩展的可用性与兼容性检查集中在两个扩展基类里extensions/cuda_extension.py 的_CudaExtension.is_available()/assert_compatible()extensions/utils.py 提供的check_system_pytorch_cuda_match()/check_pytorch_version()/set_cuda_arch_list()。对照源码安装或首次运行时最容易遇到的几类问题如下1. 找不到CUDA_HOME/nvccassert_compatible()首先检查CUDA_HOME来自torch.utils.cpp_extension.CUDA_HOME为空时直接抛错。请先安装与驱动匹配的 CUDA Toolkit并正确导出export CUDA_HOME/usr/local/cuda export PATH$CUDA_HOME/bin:$PATH注意 extensions/utils.py 里对nvcc -V解析失败、CUDA_HOME指向错误等情况都给出了明确报错文案可按提示逐一核对。2. 系统 CUDA 与 PyTorch 内置 CUDA 主版本不一致check_system_pytorch_cuda_match()会比较两处 CUDA 版本主版本major不一致 → 直接抛异常例如 PyTorch 是 cu118系统是 CUDA 12.x会拒绝编译仅次版本minor不一致 → 发出警告并继续因为 API 兼容允许编译继续进行但源码注释建议出现使用问题时改用完全匹配的 CUDA 版本重新编译。因此先按 PyTorch 的 cu 版本装好对应 CUDA Toolkit再装 Colossal-AI是标准顺序。_CudaExtension.assert_compatible()还会要求 PyTorch 至少为 1.10常量MIN_PYTORCH_VERSION_MINOR 10低于此版本会报错。3. 在无 GPU 的构建机上做交叉编译is_available()判断 CUDA 可用性的条件是torch.cuda.is_available()或设置了环境变量FORCE_CUDA。若构建机没有 GPU 却要预编译内核可用FORCE_CUDA1 BUILD_EXT1 pip install .强行进入 AOT 流程此时 extensions/utils.py 的set_cuda_arch_list()会按 CUDA 主版本自动填充TORCH_CUDA_ARCH_LIST默认覆盖 Pascal 6.0/6.1/6.2、Volta 7.0、Turing 7.5CUDA 11.x 额外加入 Ampere 8.0/8.6以便一次产出多架构可用的二进制。若只想为单一架构例如仅 A100编译可自行导出export TORCH_CUDA_ARCH_LIST8.04. Windows 安装被拒绝setup.py开头即对sys.platform win32抛错并提示改用 WSLWindows Subsystem for Linux这正是文档中只支持 Linux的落地实现。Windows 用户请先在 WSL 中配置 Linux CUDA 环境后再按上文任一方式安装。安装后去哪看运行示例装好并通过colossalai check -i后建议按阅读顺序继续了解框架能力使用 run_demo 指南 跑通第一个基于 Colossal-AI 的分布式训练 demo参考 阅读路线图 按主题系统学习并行策略、优化器与推理能力需要容器化部署时可复用仓库内置的 docker/Dockerfile 作为 Docker 镜像构建起点保证 CUDA/PyTorch 环境与 Colossal-AI 扩展编译条件一致。输出文章【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Colossal-AI 安装指南：PyPI / 源码安装全流程与 BUILD_EXT 内核扩展机制解析
Colossal-AI 安装指南PyPI / 源码安装全流程与 BUILD_EXT 内核扩展机制解析【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI本文是基于 Colossal-AI 官方中文安装文档docs/source/zh-Hans/get_started/installation.md编写的一站式安装实战指南。文章完整覆盖环境要求、PyPI 安装、源码安装两条主路径并结合当前仓库源码深入剖析BUILD_EXT1背后AOT 预编译 / JIT 运行时编译两种 CUDA 内核构建机制、版本自检命令与常见编译错误定位方法。读完本文你能够在自己的 Linux GPU 集群上正确装好 Colossal-AI并清楚知道装出来的版本预编译了哪些算子、何时会触发运行时编译、遇到版本不匹配如何排查。环境要求官方安装文档给出的最低环境门槛如下环境项最低要求说明操作系统Linux当前只支持 LinuxWindows 会被安装脚本直接拒绝Python 3.7见下方关于 PyTorch 版本约束的补充说明PyTorch 2.1满足张量并行、混合并行等核心功能所需的算子接口CUDA 11.0编译 CUDA 扩展时的工具链要求NVIDIA GPUCompute Capability 7.0即 Volta 及以上V1007.0/ RTX 20 系列7.5以及更新的 A1008.0、RTX 30 系列8.6、H1009.0等关于版本门槛有两点需要提醒Python 与 PyTorch 的组合要互相匹配。安装文档的门槛是Python 3.7但当前仓库主分支的 requirements/requirements.txt 已将 PyTorch 约束在torch2.2.0,2.5.1而新版 PyTorch 对 Python 版本有自身要求实际装机时应以你选择的 PyTorch 官方版本对应的 Python 支持范围为准。建议用 CUDA 兼容的 PyTorch 预编译包。CUDA 扩展编译时见下文会校验系统 CUDA 版本与PyTorch 内置 CUDA 版本是否一致因此先按自己驱动版本安装对应 cu 版本的 PyTorch再装 Colossal-AI是避免踩坑的关键顺序。如果你在安装过程中遇到问题可以向本项目的 Issue 区反馈附上环境信息与报错日志更易定位。安装方式总览先理解BUILD_EXT这个开关Colossal-AI 内部包含一批高性能 CUDA/C 算子如 fused optimizer、LayerNorm、MoE、FlashAttention、softmax、CPU Adam 等。这些算子有两种构建时机官方文档用BUILD_EXT环境变量来切换AOTAhead-of-Time预编译在pip install阶段把扩展编译好并打进安装产物。命令中显式设置BUILD_EXT1即开启JITJust-in-Time运行时编译不做预编译等代码第一次真正用到某个算子时在运行时用torch.utils.cpp_extension.load()现场编译并缓存.so之后复用。从 setup.py 源码可以看到这一逻辑的落地方式BUILD_EXT int(os.environ.get(BUILD_EXT, 0)) 1 if BUILD_EXT: if not TORCH_AVAILABLE: raise ModuleNotFoundError(...) # BUILD_EXT1 时必须已装好 PyTorch from extensions import ALL_EXTENSIONS for ext_cls in ALL_EXTENSIONS: ext ext_cls() if ext.support_aot and ext.is_available(): ext.assert_compatible() ext_modules.append(ext.build_aot()) else: ext_modules []也就是说开BUILD_EXT1的前提是当前 Python 环境里已经能 import torch并且扩展类只有在当前机器硬件满足条件is_available()且工具链兼容assert_compatible()时才会进入 AOT 编译列表。若一个可用内核都没有安装会直接报错Could not find any kernel compatible with the current environment。官方安装文档的原文提示也印证了这两种模式的分工不添加BUILD_EXT1时PyTorch 扩展会在运行时自动安装即 JIT 编译。因此两条安装路径各有用武之地想开箱即用、首次运行不等待编译 →BUILD_EXT1只想在纯 Python 层面使用、或构建机没有 GPU/不想预编译 → 走 JIT首次使用某算子时自动编译。从 PyPI 安装快速上手路径PyPI 上已发布预编译好的发行版直接安装pip install colossalai注该方式目前只支持 Linux。如需在安装阶段同时完成 PyTorch 扩展的 AOT 预编译避免运行时边用边编加上环境变量BUILD_EXT1 pip install colossalai需要注意当安装的是 PyPI 官方 wheel非源码本地构建时BUILD_EXT1能否真正触发预编译取决于该 wheel 的构建策略官方更推荐对扩展有强需求的用户走下面的源码安装。装完后建议执行 安装后自检 一节中的colossalai check -i报告中会明确显示当前安装产物是否包含 AOT 编译的 CUDA 扩展。从源码安装跟随主分支开发源码安装适用于需要体验主分支最新特性、需要针对自己的 GPU 架构重新编译内核、或想参与开发调试的场景。文档强调此章节与版本库主分支保持一致。基本安装流程git clone https://github.com/hpcaitech/ColossalAI.git cd ColossalAI # 安装依赖 pip install -r requirements/requirements.txt # 安装 colossalai并预编译 CUDA 扩展 BUILD_EXT1 pip install .其中 requirements/requirements.txt 列出了运行期与编译期依赖的完整清单numpy、tqdm、psutil、ninja、torch2.2.0,2.5.1、transformers、peft、ray、fastapi等pip install -r会先把这些装齐为后续的扩展编译与分布式训练做好准备。关闭 CUDA 内核预编译纯 JIT 模式如果你的环境不打算使用 CUDA 内核融合功能融合优化器如 FusedAdam/LAMB 会强制要求 CUDA 扩展此时 JIT 也会在运行到该算子时自动触发编译可以不加BUILD_EXT1pip install .在这种模式下setup.py中的ext_modules为空安装只做纯 Python 打包内核算子由 colossalai/kernel/kernel_loader.py 中的KernelLoader在运行时按需加载——先尝试 import 预编译模块colossalai._Cimport 失败则现场 JIT 编译详见 运行时加载机制。针对旧版 CUDA 10.2 的 cub 手动补齐步骤如果你的环境仍在使用 CUDA 10.2文档说明仍可从源码安装但需要手动下载 cub 库并复制到对应目录# 克隆仓库 git clone https://github.com/hpcaitech/ColossalAI.git cd ColossalAI # 下载 cub 库1.8.0 wget https://github.com/NVIDIA/cub/archive/refs/tags/1.8.0.zip unzip 1.8.0.zip cp -r cub-1.8.0/cub/ colossalai/kernel/cuda_native/csrc/kernels/include/ # 安装 BUILD_EXT1 pip install .需要特别说明适用前提CUDA 10.2 低于本文开头CUDA 11.0的通用门槛该段指令面向的是历史上较老版本的 Colossal-AI。从当前仓库主分支的目录结构看内核源码已统一收拢到 extensions/csrc 与 colossalai/kernel/extensions/csrc 等路径下文档中colossalai/kernel/cuda_native/csrc/kernels/include/这样的旧路径在当前代码布局中已不存在。如果你使用新版本源码请直接执行标准的BUILD_EXT1 pip install .无需手动放置 cub仅当你确实在维护 CUDA 10.2 时代的老版本代码时才需要参考这一段。源码级原理BUILD_EXT1到底在构建什么被纳入 AOT 编译的扩展清单在BUILD_EXT1时setup.py从extensions包导入ALL_EXTENSIONS并逐个构建。当前主分支 extensions/init.py 中注册的扩展类为CpuAdamArmExtension/CpuAdamX86ExtensionCPU Adam区分 ARM 与 x86LayerNormCudaExtensionMoeCudaExtensionMoE 相关内核FusedOptimizerCudaExtension融合优化器内核InferenceOpsCudaExtension推理算子ScaledMaskedSoftmaxCudaExtension/ScaledUpperTriangleMaskedSoftmaxCudaExtensionFlashAttentionDaoCudaExtension/FlashAttentionSdpaCudaExtension/FlashAttentionNpuExtension构建后的扩展模块统一落在colossalai._C这个预编译模块命名空间下对应仓库中的 colossalai/_C/init.py并可通过importlib.import_module在运行时被KernelLoader反查加载相关逻辑见 extensions/cpp_extension.py 中的_CppExtension.import_op()。AOT vs JIT运行时加载机制与缓存目录无论安装时是否预编译运行期统一入口都是 colossalai/kernel/kernel_loader.py 中的KernelLoader。其load()流程是实例化该 loader 注册表里的全部扩展类逐一调用is_available()过滤出当前机器可用的内核并调用assert_compatible()做兼容性校验多个可用内核时按priority降序取最高优先级者调用_Extension.load()——对应 extensions/cpp_extension.py 中先import_op()尝试导入colossalai._C.xxx预编译模块ImportError/ModuleNotFoundError时回退到build_jit()现场编译。JIT 编译产物并非散落在安装目录而是统一缓存在用户目录下。根据 extensions/base_extension.py 中的get_jit_extension_folder_path()缓存路径格式为~/.cache/colossalai/torch_extensions/torch主.次_设备名-设备版本-hash其中 hash 由colossalai.__file__路径经 SHA-256 生成缓存目录按 torch 版本 设备型号区分。因此同一次编译结果在后续运行、重启进程后可直接复用日志中会打印Loading the JIT-built ... kernel during runtime now或Compiling the JIT ... kernel during runtime now提示当前是加载缓存还是新编译。理解了这套机制也就不难明白官方文档那句话的含义不设BUILD_EXT1扩展会在运行时自动构建只是第一次触达某个算子的启动耗时会更长。安装后的环境自检使用colossalai check -i一键体检安装包通过setup.py的entry_points注册了colossalai命令路由定义在 colossalai/cli/cli.pyrun/check两个子命令。其中check子命令colossalai/cli/check/init.py支持-i/--installation参数会输出一份完整的安装体检报告colossalai check -i报告由 colossalai/cli/check/check_installation.py 生成主要字段含义报告区块字段含义与判读EnvironmentColossal-AI version当前 colossalai 版本号version.txt记录的仓库版本为0.5.0若安装时带了 CUDA 扩展版本号形如X.X.XtorchX.XXcuXX.XEnvironmentPyTorch version当前环境的 torch 版本EnvironmentSystem CUDA version由$CUDA_HOME/bin/nvcc -V解析出的系统 CUDA 版本显示 N/A 说明 CUDA_HOME 未配置EnvironmentCUDA version required by PyTorchtorch.version.cuda给出的 PyTorch 内置 CUDA 版本N/A 说明装的是无 CUDA 的 CPU 版 torchCUDA Extensions AOT CompilationFound AOT CUDA Extension当前安装产物是否带 AOT 预编译扩展CUDA Extensions AOT CompilationPyTorch / CUDA version used for AOT compilation预编译时使用的版本可从版本号后缀解析用于核对与当前运行环境的差异Compatibility三项兼容性对比PyTorch 版本是否与 AOT 编译版本一致系统 CUDA 与 PyTorch 内置 CUDA 是否同主版本系统 CUDA 与 AOT 编译所用 CUDA 是否匹配官方实现中_is_compatible()只要求主版本号major与次版本号minor严格一致patch 级差异被忽略这与编译校验逻辑保持一致。命令行快速冒烟体检之外还可以直接确认版本与导入是否正常python -c import colossalai; print(colossalai.__version__) colossalai run --help # 查看官方多机/单机启动器用法run子命令由 colossalai/cli/launcher/run.py 提供可用于后续启动分布式训练脚本前的环境联调。编译期的硬性检查与常见问题定位CUDA 扩展的可用性与兼容性检查集中在两个扩展基类里extensions/cuda_extension.py 的_CudaExtension.is_available()/assert_compatible()extensions/utils.py 提供的check_system_pytorch_cuda_match()/check_pytorch_version()/set_cuda_arch_list()。对照源码安装或首次运行时最容易遇到的几类问题如下1. 找不到CUDA_HOME/nvccassert_compatible()首先检查CUDA_HOME来自torch.utils.cpp_extension.CUDA_HOME为空时直接抛错。请先安装与驱动匹配的 CUDA Toolkit并正确导出export CUDA_HOME/usr/local/cuda export PATH$CUDA_HOME/bin:$PATH注意 extensions/utils.py 里对nvcc -V解析失败、CUDA_HOME指向错误等情况都给出了明确报错文案可按提示逐一核对。2. 系统 CUDA 与 PyTorch 内置 CUDA 主版本不一致check_system_pytorch_cuda_match()会比较两处 CUDA 版本主版本major不一致 → 直接抛异常例如 PyTorch 是 cu118系统是 CUDA 12.x会拒绝编译仅次版本minor不一致 → 发出警告并继续因为 API 兼容允许编译继续进行但源码注释建议出现使用问题时改用完全匹配的 CUDA 版本重新编译。因此先按 PyTorch 的 cu 版本装好对应 CUDA Toolkit再装 Colossal-AI是标准顺序。_CudaExtension.assert_compatible()还会要求 PyTorch 至少为 1.10常量MIN_PYTORCH_VERSION_MINOR 10低于此版本会报错。3. 在无 GPU 的构建机上做交叉编译is_available()判断 CUDA 可用性的条件是torch.cuda.is_available()或设置了环境变量FORCE_CUDA。若构建机没有 GPU 却要预编译内核可用FORCE_CUDA1 BUILD_EXT1 pip install .强行进入 AOT 流程此时 extensions/utils.py 的set_cuda_arch_list()会按 CUDA 主版本自动填充TORCH_CUDA_ARCH_LIST默认覆盖 Pascal 6.0/6.1/6.2、Volta 7.0、Turing 7.5CUDA 11.x 额外加入 Ampere 8.0/8.6以便一次产出多架构可用的二进制。若只想为单一架构例如仅 A100编译可自行导出export TORCH_CUDA_ARCH_LIST8.04. Windows 安装被拒绝setup.py开头即对sys.platform win32抛错并提示改用 WSLWindows Subsystem for Linux这正是文档中只支持 Linux的落地实现。Windows 用户请先在 WSL 中配置 Linux CUDA 环境后再按上文任一方式安装。安装后去哪看运行示例装好并通过colossalai check -i后建议按阅读顺序继续了解框架能力使用 run_demo 指南 跑通第一个基于 Colossal-AI 的分布式训练 demo参考 阅读路线图 按主题系统学习并行策略、优化器与推理能力需要容器化部署时可复用仓库内置的 docker/Dockerfile 作为 Docker 镜像构建起点保证 CUDA/PyTorch 环境与 Colossal-AI 扩展编译条件一致。输出文章【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考