编写 PyTorch DDP torchrun 启动脚本
本页适用于已经使用 torchrun、torch.distributed 或 PyTorch DDP 的训练代码。重点是把创建页中的 Worker 数量 和每个 Worker 的 GPU 数量，转换成 torchrun 参数，并让 启动命令 的结束状态正确反映训练结果。
从 Worker 数量到 torchrun 参数
创建 DDP 任务时，AIStudio 会先启动您在创建页中填写的 Worker。随后，同一份启动命令会在每个 Worker 中执行，并由 torchrun 在 Worker 内启动 PyTorch 训练进程。
例如创建 2 个 Worker，每个 Worker 4 张 GPU：
- AIStudio 启动 2 个 Worker，并为每个 Worker 注入不同的编号。
- 启动脚本在每个 Worker 中执行同一条
torchrun命令。 torchrun --nproc_per_node=4 --nnodes=2会启动总计 8 个 PyTorch 训练进程。
这个命令中的 4 和 2 来自本例的固定规格：每个 Worker 4 张 GPU，Worker 数量为 2。固定规格的脚本可以直接写入这些值；如果希望同一份脚本在不同 Worker 数量或不同 GPU 规格下复用，请像下文示例一样读取平台变量和当前 Worker 的 GPU 数量。
接下来要做的是把平台注入的 Worker 信息交给 torchrun：Worker 数量用于 --nnodes，Worker 编号用于 --node_rank；每个 Worker 内启动多少训练进程，则由脚本中的 --nproc_per_node 决定。
读取平台注入的 DDP 变量
创建页选择 PyTorch DDP 后，平台会为不同 Worker 注入不同的 RANK，并注入相同的 MASTER_ADDR、MASTER_PORT 和 WORLD_SIZE。启动脚本读取这些变量后，平台会自动区分每个 Worker，不用为每个 Worker 手动写一条不同的命令。
| 平台注入变量 | 在本文中的含义 | 传给 torchrun 的参数 | 注意 |
|---|---|---|---|
MASTER_ADDR | 主 Worker 地址 | --master_addr | 所有 Worker 相同 |
MASTER_PORT | 主 Worker 端口 | --master_port | 默认 29500 |
WORLD_SIZE | Worker 数量 | --nnodes | 不是 PyTorch 进程总数 |
RANK | 当前 Worker 编号 | --node_rank | 不是 PyTorch 全局进程 rank |
警告
平台注入的 WORLD_SIZE 表示 Worker 数量。PyTorch 训练进程中的 WORLD_SIZE 表示训练进程总数。为避免混淆，启动脚本中请立即写成 NNODES=${WORLD_SIZE}，后续只使用 NNODES。不要把平台注入的 WORLD_SIZE 重新导出给训练进程。
在创建页选择 PyTorch DDP
在创建训练任务时，先按训练代码选择 Worker 规格 和 Worker 数量，再在 分布式框架 中选择 PyTorch DDP。完整创建步骤请参考使用包年包月资源提交训练任务。
如果脚本默认每个节点 GPU 数量一致，不要随意开启自定义卡数。多数 DDP 脚本会按相同的 nproc_per_node、batch size 和 rank 映射启动；Worker 之间 GPU 数量不一致时，脚本通常需要额外适配。
如果您以前使用物理机或 hostfile
如果您以前手动登录多台机器训练，可以这样理解：
| 物理机训练习惯 | AIStudio 中对应方式 |
|---|---|
| 手动准备机器列表或 hostfile | 创建页填写 Worker 数量 |
| 手动指定哪台机器是 rank 0 | 平台通过 RANK=0 和 MASTER_ADDR 指定 |
| 每台机器运行一条略有不同的命令 | 每个 Worker 运行同一条启动命令，通过 RANK 自动区分 |
| 手动计算节点数 | 使用平台注入的 WORLD_SIZE，并在脚本中保存为 NNODES |
| 手动决定每台机器启动几个进程 | 在脚本中设置 --nproc_per_node |
编写 launch_ddp.sh
把 torchrun 参数封装到脚本中，便于在开发机预检和正式任务中复用。以下示例会自动读取平台变量，并打印本次分布式启动形状：
#!/usr/bin/env bash
set -euo pipefail
TRAIN_SCRIPT=${TRAIN_SCRIPT:-/mnt/public/project-a/train.py}
OUTPUT_DIR=${OUTPUT_DIR:-/mnt/public/project-a/outputs/ddp-smoke}
mkdir -p "$OUTPUT_DIR"
GPUS_PER_NODE=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l | tr -d ' ')
NNODES=${WORLD_SIZE:?platform WORLD_SIZE is required}
NODE_RANK=${RANK:?platform RANK is required}
MASTER_ADDR=${MASTER_ADDR:?platform MASTER_ADDR is required}
MASTER_PORT=${MASTER_PORT:-29500}
TOTAL_PROCS=$((GPUS_PER_NODE * NNODES))
echo "DDP launch: nnodes=${NNODES} node_rank=${NODE_RANK} gpus_per_node=${GPUS_PER_NODE} total_procs=${TOTAL_PROCS} master=${MASTER_ADDR}:${MASTER_PORT}"
torchrun \
--nproc_per_node="${GPUS_PER_NODE}" \
--nnodes="${NNODES}" \
--node_rank="${NODE_RANK}" \
--master_addr="${MASTER_ADDR}" \
--master_port="${MASTER_PORT}" \
"$TRAIN_SCRIPT" \
--output_dir "$OUTPUT_DIR"对于 2 个 Worker、每个 Worker 4 张 GPU 的任务：
- 平台注入
WORLD_SIZE=2。 - 每个 Worker 检测到
GPUS_PER_NODE=4。 torchrun在每个 Worker 上启动 4 个训练进程，总训练进程数为 8。
填写启动命令
在创建页的 启动命令 中调用脚本。启动命令 应尽量短，把复杂逻辑放在脚本里：
set -euo pipefail
export TRAIN_SCRIPT=/mnt/public/project-a/train.py
export OUTPUT_DIR=/mnt/public/project-a/outputs/${RUN_ID:-manual-run}
mkdir -p "$OUTPUT_DIR"
bash /mnt/public/project-a/launch_ddp.sh > >(tee -i "$OUTPUT_DIR/launch_ddp.log") 2>&1如果需要保存日志，优先使用上面的 Process Substitution 写法。不要把训练命令简单写成 bash launch_ddp.sh | tee log，否则日志转存可能改变平台看到的最终结果。详见编写可靠启动命令。
验证 DDP 已经正常启动
第一次提交前，建议先用小数据或少量 step 验证分布式启动链路。日志中至少应能确认：
- 每个 Worker 都打印了
DDP launch，并且node_rank不同。 nnodes等于创建页中的 Worker 数量。gpus_per_node等于单个 Worker 的 GPU 数量。- 训练脚本进入第一个 batch，而不是卡在初始化或通信阶段。
- checkpoint、TensorBoard 日志和训练结果写入共享存储路径。
如果需要确认 torchrun 已经正确设置 PyTorch 进程变量，可以在训练脚本入口临时打印：
import os
print(
"torchrun env:",
"RANK=", os.environ.get("RANK"),
"LOCAL_RANK=", os.environ.get("LOCAL_RANK"),
"WORLD_SIZE=", os.environ.get("WORLD_SIZE"),
)对于 2 个 Worker、每个 Worker 4 张 GPU，训练进程中应能看到 WORLD_SIZE=8。
如果训练脚本会由 rank 0 准备共享数据或保存 checkpoint，请确认其他 rank 会等待共享状态准备完成。更多提交前检查请参考准备代码、数据与输出目录和在开发机中试跑训练脚本。
定位 DDP 启动失败
torchrun参数缺失或为空：检查创建页是否选择 PyTorch DDP，以及脚本是否读取MASTER_ADDR、WORLD_SIZE和RANK。- 只有部分 Worker 进入训练：检查各 Worker 日志中的
node_rank是否完整，Worker 数量是否与nnodes一致。 - 初始化后长时间卡住：检查
MASTER_ADDR、MASTER_PORT是否传入torchrun，训练网和 RDMA 配置 是否符合任务规格。 - 第一个 batch 前失败：检查镜像依赖、数据路径、环境变量和 data collator 是否已在开发机预检。
- 任务显示成功但日志里训练失败：检查 启动命令 是否隐藏了训练脚本退出码，尤其是
tee、上传结果、失败后sleep等收尾逻辑。
运行中需要进一步排查时，可在任务详情页登录 Worker 的 Web Terminal，或使用 atlctl 命令行调试工具下发检查命令。