Colossal-AI 命令行工具完全指南：环境自检与分布式训练一键启动
【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI
Colossal-AI 在提供 Python 训练接口之外，还内置了一套命令行工具（CLI），帮助用户完成「安装是否正确的环境体检」与「单机/多机分布式训练的进程一键拉起」两类高频操作。本文以官方文档 docs/source/zh-Hans/basics/command_line_tool.md 为骨架，结合仓库内 colossalai/cli 的源码实现，逐条拆解colossalai check -i与colossalai run的全部参数、输出含义与底层启动逻辑。读完本文，你将能够独立完成 Colossal-AI 环境的兼容性体检，并用一条命令在单机或多机上启动任意分布式训练脚本。
一、命令行工具概览
Colossal-AI 的 CLI 由 Click 框架实现，入口定义在仓库根目录的 setup.py 的entry_points中，将colossalai命令指向colossalai.cli:cli。命令组的注册代码位于 colossalai/cli/cli.py：
import click from .check import check from .launcher import run @click.group() def cli(): pass cli.add_command(run) cli.add_command(check)官方文档将命令行工具的功能归纳为三类：检查 Colossal-AI 是否安装正确、启动分布式训练、以及张量并行基准测试。需要说明的是，就当前仓库源码而言，colossalai/cli 目录中实际实现并注册的子命令只有check与run（对应colossalai/cli/check/与colossalai/cli/launcher/两个模块），文档中列出的张量并行基准测试属于 CLI 的设计能力范围，本文重点展开当前可直接使用的两条命令。
二、安装检查：colossalai check -i
在跑任何分布式训练之前，第一步永远是确认环境是否健康。命令组中注册的check子命令定义于 colossalai/cli/check/init.py，其用法如下：
colossalai check -i # 等价写法 colossalai check --installation-i/--installation是一个布尔开关；若执行colossalai check而不带任何选项，则会输出No option is given提示。
当开启-i后，程序会调用 check_installation.py 中的check_installation()，依次采集当前环境的 PyTorch / CUDA / Colossal-AI 版本、CUDA Extension 的预编译状态，并对它们做交叉兼容性比对，最终打印一份三段的安装报告。
2.1 报告输出的三段内容
以click.echo打印的报告被分为三个区块，含义如下：
Environment（环境信息）
Colossal-AI version：读取colossalai.__version__，即安装时由 version.txt 生成的版本号（当前仓库为0.5.0）。PyTorch version：读取torch.__version__，截取前三位段（如2.1.0）。System CUDA version：通过$CUDA_HOME/bin/nvcc -V探测宿主机安装的 CUDA 工具链版本。CUDA version required by PyTorch：读取torch.version.cuda，即当前 PyTorch 编译时所依赖的 CUDA 版本。
CUDA Extensions AOT Compilation（CUDA 扩展预编译信息）
Found AOT CUDA Extension：是否找到 AOT（ahead-of-time，预先编译）构建的 CUDA 扩展。PyTorch version used for AOT compilation/CUDA version used for AOT compilation：从 Colossal-AI 的版本字符串中解析出来。当安装时设置了环境变量BUILD_EXT=1，扩展会在安装阶段被预编译进colossalai._C，此时版本号形如X.X.X+torchX.XXcuXX.X，可供解析；否则这两项显示为 N/A。
Compatibility（兼容性比对）
PyTorch version match：当前 PyTorch 版本与 AOT 编译所用 PyTorch 是否一致。System and PyTorch CUDA version match：宿主机 CUDA 与 PyTorch 所需 CUDA 是否兼容。System and Colossal-AI CUDA version match：宿主机 CUDA 与 Colossal-AI AOT 编译所用 CUDA 是否兼容。
2.2 结果符号与兼容性判定规则
为了让报告更易读，源码通过to_click_output把布尔值映射为符号：True → ✓，False → x，None → N/A。其中三种取值分别意味着：
- N/A：无法进行该项比对。例如未设置
CUDA_HOME时无法得知系统 CUDA 版本；未开启 AOT 编译时没有「参考版本」可供比对。源码注释明确建议：若 System CUDA version 为 N/A，可通过设置CUDA_HOME环境变量让检测程序定位 CUDA 安装路径。 - ✓：比对通过。
- x：比对不通过，说明环境中存在版本不匹配。
版本兼容性并非简单要求完全相等。查看_is_compatible的实现可以发现其规则：将版本拆分为[major, minor, patch]后逐段比较——major 与 minor 必须完全一致，而 patch 版本（第三段）不一致仍视为兼容；若仅有两段版本号，会补一个x占位。也就是说，PyTorch 的2.1.0与2.1.5属于兼容，而2.1.0与2.2.0则被判定为不兼容。
2.3 AOT 与 JIT：理解两种扩展构建方式
报告特意区分了 AOT 与 JIT 两种 CUDA 扩展的构建方式，这是理解报告结果的关键：
- AOT（预编译）：在
pip install时设置环境变量BUILD_EXT=1，内核在安装阶段编译并固化到colossalai._C中，运行时无需再编译。代价是安装时使用的 PyTorch / CUDA 版本必须与运行环境一致，这正是 Compatibility 区块存在的意义。 - JIT（即时编译）：未开启 AOT 时，CUDA 内核会在程序首次运行期间自动编译到缓存目录
~/.cache/colossalai/torch_extensions。报告对此特别安抚用户：If AOT compilation is not enabled, stay calm as the CUDA kernels can still be built during runtime（若未启用 AOT 编译也无需惊慌，CUDA 内核依然可以在运行时构建）。
2.4 检查结果的排错指引
将报告与源码逻辑对照，可得到如下排错路径：
| 报告现象 | 可能原因 | 处理建议 |
|---|---|---|
| System CUDA version 为 N/A | CUDA_HOME未设置或nvcc不在$CUDA_HOME/bin | 在 shell 中export CUDA_HOME=/usr/local/cuda（以实际路径为准）后重试 |
| PyTorch CUDA version 为 N/A | 安装的是 CPU 版或不带 CUDA 的 PyTorch | 按torch.version.cuda的说明重新安装 CUDA 兼容版 PyTorch（官方渠道见https://pytorch.org/get-started/locally/） |
| 某行 Compatibility 为 x | 宿主机/运行时版本与编译期版本 major 或 minor 不一致 | 统一 PyTorch、CUDA 与 Colossal-AI 预编译时的版本组合 |
| 三项比对全部 N/A | 未开启 AOT 编译 | 属正常状态，扩展将在运行时以 JIT 方式构建 |
三、分布式启动器：colossalai run
CLI 的第二个核心能力是分布式进程启动。run子命令定义于 colossalai/cli/launcher/init.py，它封装了 PyTorch 官方的分布式启动器（torchrun/torch.distributed.launch），其最大价值在于：PyTorch 原生启动器在多机场景下需要在每个节点上分别执行命令，而colossalai run只需在任意一台节点上调用一次，即可由内部逻辑通过 SSH 把命令分发到所有节点并回收执行结果。
3.1 训练代码侧的前置要求
由于该启动器本质上是 PyTorch 启动器的封装，训练脚本中不应使用手动传参的colossalai.launch，而应使用从环境变量读取分布式信息的colossalai.launch_from_torch。具体写法可参考 docs/source/zh-Hans/basics/launch_colossalai.md：
import colossalai # rank / world_size / host / port 等由启动器写入环境变量 colossalai.launch_from_torch() # ... 加载模型、配置优化器、开始训练3.2 完整参数表
下表整理了 colossalai/cli/launcher/init.py 中run命令注册的全部选项及默认值，可作为写命令时的速查手册：
| 参数 | 默认值 | 作用 |
|---|---|---|
-H, --host | None | 主机名列表，格式为<host1>,<host2>，适合节点数较少的场景 |
--hostfile | None | 主机清单文件路径，文件中每行一个主机名，适合节点较多的集群 |
--include | None | 仅使用 hostfile 中列出的部分主机，格式同--host；仅与--hostfile搭配有效 |
--exclude | None | 排除 hostfile 中的部分主机；与--include互斥，仅与--hostfile搭配有效 |
--num_nodes | -1 | 实际参与任务的主机数（从 hostfile 中截取前 N 个）；仅与--hostfile搭配有效 |
--nproc_per_node | None（必填） | 每台节点上启动的进程数（通常等于单机 GPU 数） |
--master_port | 29500 | PyTorch 分布式通信所用的端口 |
--master_addr | 127.0.0.1 | 0 号节点的 IP；多机时若未显式指定，代码会自动改写为首个节点的 hostname |
--extra_launch_args | None | 透传给底层 torch 分布式启动器的额外参数，格式为arg1=1,arg2=2，会被转换为--arg1=1 --arg2=2 |
--ssh-port | None | SSH 连接端口（多机场景需要各节点统一） |
-m | None | 以模块方式运行（python -m语义），对应的值须是模块名而非.py文件 |
user_script | — | 位置参数：用户训练脚本，如train.py |
user_args | — | 可变长位置参数：传递给训练脚本的其余参数 |
参数校验逻辑同样值得注意：--nproc_per_node缺失时直接报错退出；user_script若不以.py结尾或-m值以.py结尾都会触发错误提示；--hostfile与--host同时给出时会被拒绝，--include与--exclude也互斥。
3.3 单节点启动
最常见的用法是在当前节点启动多卡训练，只需指定进程数（GPU 数）即可，默认走29500端口：
# 在当前节点启动 4 进程（4 卡）训练，默认端口 29500 colossalai run --nproc_per_node 4 train.py # 使用 4 卡并将通信端口改为 29505 colossalai run --nproc_per_node 4 --master_port 29505 train.py # 把额外参数传给 torch 分布式启动器，例如启用 --standalone colossalai run --nproc_per_node 4 --extra_launch_args standalone train.py单机场景下无主机清单时，run.py 中的launch_multi_processes会自动把127.0.0.1加入主机列表，仅调用本地运行分支，因此无需关心--master_addr。
3.4 多节点启动的三种方式
多机训练才是colossalai run相比原生 torch 启动器的主要优势场景，提供三种控制主机范围的方式：
方式一：直接列出主机（--host），适合节点数较少的情况：
colossalai run --host host1,host2 --master_addr host1 --nproc_per_node 4 train.py方式二：使用 hostfile（--hostfile），适合节点较多、需要统一维护清单的场景。hostfile 的解析实现在 fetch_hostfile 中：逐行读取、跳过空行、每行一个主机名、不允许出现重复主机。可配合集群调度脚本动态生成，例如应用示例 Colossal-LLaMA 的 hostfile.example 即采用这种一行一节点的格式：
colossalai run --hostfile <hostfile路径> --master_addr host1 --nproc_per_node 4 train.py方式三：在 hostfile 基础上做过滤或截取，应对「集群很大但本次只占用一部分」的需求：
# 只用 hostfile 中的 host1、host2 colossalai run --hostfile <hostfile路径> --master_addr host1 --include host1,host2 --nproc_per_node 4 train.py # 排除 hostfile 中的 host2 colossalai run --hostfile <hostfile路径> --master_addr host1 --exclude host2 --nproc_per_node 4 train.py # 截取 hostfile 中的前 2 个节点 colossalai run --hostfile <hostfile路径> --num_nodes 2 --nproc_per_node 4 train.py上述逻辑在 parse_device_filter 中实现：先校验--include/--exclude互斥及主机名确实存在于 hostfile，再对主机池做增删；若--num_nodes为正整数，则只保留主机池中的前 N 个节点。
3.5 多机运行的底层流程
从源码梳理，colossalai run一次调用的完整生命周期如下：
run命令解析参数并做合法性校验（colossalai/cli/launcher/init.py）；- 根据
--hostfile/--host决定主机池：有 hostfile 则读取解析、过滤、截取；有--host则解析逗号分隔的主机列表；两者皆无则回落到本机127.0.0.1（launch_multi_processes）； - 多节点时若
--master_addr仍是默认的127.0.0.1，自动改写为节点列表首个主机的 hostname； - 收集当前环境变量（跳过含换行的变量），通过
MultiNodeRunner（multinode_runner.py）对每台主机建立 SSH 连接并下发启动命令（runner.send）； - 统一回收各节点的状态消息，打印
====== Training on All Nodes ======汇总；任一节点失败则sys.exit(1)，全部成功则sys.exit(0)，使启动器本身表现为一个可正常返回退出码的进程。
3.6 底层对 PyTorch 启动器的版本适配
get_launch_command 展示了生成命令时的版本分支逻辑，可帮助理解不同 torch 版本下的行为差异：
- torch < 1.9：使用
python -m torch.distributed.launch，并显式传入--master_addr、--master_port、--nnodes、--node_rank、--nproc_per_node； - torch 1.9：改用
python -m torch.distributed.run； - torch > 1.9：直接调用
torchrun； - torch < 2.0：若使用
-m以模块方式运行，会抛出Torch version < 2.0 does not support running as module。
同时，--master_addr/--master_port会被并入默认的 rendezvous（rdzv）参数；--extra_launch_args中若覆盖了这两个键，则以用户传入值为准。整条命令在每台节点上通过runner.send下发，最终形态类似：
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr=host1 --master_port=29500 train.py四、小结：CLI 的适用场景
| 场景 | 推荐做法 |
|---|---|
| 更换机器/容器/驱动后验证环境 | colossalai check -i，重点看 Compatibility 区块是否全 ✓ |
| 单机单卡 | 直接用python train.py或在代码内colossalai.launch即可 |
| 单机多卡 | colossalai run --nproc_per_node <GPU数> train.py |
| 小型多机（节点可枚举） | colossalai run --host h1,h2 --master_addr h1 --nproc_per_node 4 train.py |
| 大规模集群 | 维护 hostfile，配合--include/--exclude/--num_nodes动态圈定本次任务占用的节点 |
关于命令行启动器与colossalai.launch、colossalai.launch_from_torch的更多细节（包括 SLURM、OpenMPI 等其他启动方式的对比），可进一步阅读同一教程体系下的 启动 Colossal-AI；分布式训练中 host、port、rank、world_size 等基础概念则见 分布式训练 与 Colossal-AI 总览。
上述全部命令行行为均可在仓库源码 colossalai/cli 中逐行核对：安装检查逻辑见 check_installation.py，启动器参数与流程见 launcher/init.py 与 launcher/run.py，命令注册与入口见 cli.py 与 setup.py。学会这两条命令，你就掌握了 Colossal-AI 从「环境体检」到「任务拉起」的最小运维闭环。
【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
