Colossal-AI 命令行工具完全指南环境自检与分布式训练一键启动【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAIColossal-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.pyimport 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 CompilationCUDA 扩展预编译信息Found AOT CUDA Extension是否找到 AOTahead-of-time预先编译构建的 CUDA 扩展。PyTorch version used for AOT compilation/CUDA version used for AOT compilation从 Colossal-AI 的版本字符串中解析出来。当安装时设置了环境变量BUILD_EXT1扩展会在安装阶段被预编译进colossalai._C此时版本号形如X.X.XtorchX.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 → xNone → 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_EXT1内核在安装阶段编译并固化到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/ACUDA_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 runCLI 的第二个核心能力是分布式进程启动。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.mdimport colossalai # rank / world_size / host / port 等由启动器写入环境变量 colossalai.launch_from_torch() # ... 加载模型、配置优化器、开始训练3.2 完整参数表下表整理了 colossalai/cli/launcher/init.py 中run命令注册的全部选项及默认值可作为写命令时的速查手册参数默认值作用-H, --hostNone主机名列表格式为host1,host2适合节点数较少的场景--hostfileNone主机清单文件路径文件中每行一个主机名适合节点较多的集群--includeNone仅使用 hostfile 中列出的部分主机格式同--host仅与--hostfile搭配有效--excludeNone排除 hostfile 中的部分主机与--include互斥仅与--hostfile搭配有效--num_nodes-1实际参与任务的主机数从 hostfile 中截取前 N 个仅与--hostfile搭配有效--nproc_per_nodeNone必填每台节点上启动的进程数通常等于单机 GPU 数--master_port29500PyTorch 分布式通信所用的端口--master_addr127.0.0.10 号节点的 IP多机时若未显式指定代码会自动改写为首个节点的 hostname--extra_launch_argsNone透传给底层 torch 分布式启动器的额外参数格式为arg11,arg22会被转换为--arg11 --arg22--ssh-portNoneSSH 连接端口多机场景需要各节点统一-mNone以模块方式运行python -m语义对应的值须是模块名而非.py文件user_script—位置参数用户训练脚本如train.pyuser_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.1launch_multi_processes多节点时若--master_addr仍是默认的127.0.0.1自动改写为节点列表首个主机的 hostname收集当前环境变量跳过含换行的变量通过MultiNodeRunnermultinode_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_nodetorch 1.9改用python -m torch.distributed.runtorch 1.9直接调用torchruntorch 2.0若使用-m以模块方式运行会抛出Torch version 2.0 does not support running as module。同时--master_addr/--master_port会被并入默认的 rendezvousrdzv参数--extra_launch_args中若覆盖了这两个键则以用户传入值为准。整条命令在每台节点上通过runner.send下发最终形态类似torchrun --nproc_per_node4 --nnodes2 --node_rank0 --master_addrhost1 --master_port29500 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），仅供参考