Slurm Pitfalls
本文使用的 Slurm 版本为 22.05.6
问题背景
在编写分布式训练脚本时，我们通常需要通过环境变量向每个进程传递 RANK、WORLD_SIZE 和 LOCAL_RANK 等信息。本文记录了一个在 Slurm 脚本中设置环境变量时遇到的陷阱。
假设我们有以下两个文件：
run_slurm.sh
#!/bin/bash#SBATCH --job-name=pt-distributed#SBATCH --partition=c003t#SBATCH -N 2#SBATCH -n 8#SBATCH --time=00:10:00#SBATCH --output=test-slurm-%j.out#SBATCH --error=test-slurm-%j.err#SBATCH --comment=bupt_hpc
export RANK=$SLURM_PROCIDexport WORLD_SIZE=$SLURM_NTASKSexport LOCAL_RANK=$SLURM_LOCALID
python test.pytest.py
import os
rank = int(os.environ.get("RANK", -1))world_size = int(os.environ.get("WORLD_SIZE", -1))local_rank = int(os.environ.get("LOCAL_RANK", -1))
print(f"{os.uname().nodename}, rank={rank}, world_size={world_size}, local_rank={local_rank}")我们期望每个进程能获取到正确的 rank 信息，输出应该类似于下图：
问题现象
问题一：缺少 srun 命令
最初版本的脚本中，我们忘记添加 srun 命令，直接执行：
python test.py结果只输出了一行：
cpu1, rank=0, world_size=8, local_rank=0这说明脚本只在主节点上运行了一次 Python，并没有启动 8 个并行进程。Slurm 的 -n 8 只是声明了任务数，必须通过 srun 才能真正启动多个并行进程。
为什么输出是 cpu1？ Slurm 会从分配的节点中选择一个作为”批处理主节点”（batch host），整个 run_slurm.sh 脚本只在这个节点上执行一次。
为什么 rank 和 local_rank 是 0？ 在批处理脚本中，$SLURM_PROCID 和 $SLURM_LOCALID 的值始终为 0，因为批处理脚本本身被视为 rank 0 的任务。只有通过 srun 启动的子任务才会获得各自的 rank 编号。
问题二：环境变量未正确传递
添加 srun 后，输出变成了 8 行，但所有进程的 rank 都是 0：
cpu2, rank=0, world_size=8, local_rank=0cpu1, rank=0, world_size=8, local_rank=0cpu1, rank=0, world_size=8, local_rank=0cpu1, rank=0, world_size=8, local_rank=0cpu1, rank=0, world_size=8, local_rank=0cpu1, rank=0, world_size=8, local_rank=0cpu1, rank=0, world_size=8, local_rank=0cpu1, rank=0, world_size=8, local_rank=0所有进程的 RANK 和 LOCAL_RANK 都是 0！这是因为 Slurm 脚本中的 export 语句只在脚本启动时执行一次，此时 $SLURM_PROCID 等变量还未被设置为每个任务特定的值。
为什么所有进程都能读取到 RANK 等变量，而不是 -1？ 当 srun 启动子任务时，它会将主节点脚本中 export 的环境变量拷贝到所有子任务的环境中。因此，即使子任务运行在其他节点（如 cpu2），它们也能读取到这些变量。
但问题在于：拷贝的是变量值，而不是变量引用。主节点上 $SLURM_PROCID 在脚本执行时的值为 0，所以 RANK 被设置为字符串 "0" 并被拷贝到所有子任务，而不是让每个子任务各自解析 $SLURM_PROCID。
解决方案
第一步：直接使用 Slurm 原生环境变量
修改 test.py，直接读取 Slurm 提供的环境变量：
rank = int(os.environ.get("SLURM_PROCID", -1))world_size = int(os.environ.get("SLURM_NTASKS", -1))local_rank = int(os.environ.get("SLURM_LOCALID", -1))输出变为：
cpu2, rank=7, world_size=8, local_rank=0cpu1, rank=5, world_size=8, local_rank=5cpu1, rank=0, world_size=8, local_rank=0cpu1, rank=2, world_size=8, local_rank=2cpu1, rank=3, world_size=8, local_rank=3cpu1, rank=1, world_size=8, local_rank=1cpu1, rank=4, world_size=8, local_rank=4cpu1, rank=6, world_size=8, local_rank=6现在 rank 值正确了，但任务分布不均匀：cpu1 上运行了 7 个进程，cpu2 上只有 1 个进程。
第二步：指定每个节点的任务数
在 run_slurm.sh 中添加 --ntasks-per-node=4：
#SBATCH --ntasks-per-node=4再次提交任务后，输出终于正常：
cpu1, rank=0, world_size=8, local_rank=0cpu2, rank=4, world_size=8, local_rank=0cpu1, rank=2, world_size=8, local_rank=2cpu1, rank=1, world_size=8, local_rank=1cpu1, rank=3, world_size=8, local_rank=3cpu2, rank=7, world_size=8, local_rank=3cpu2, rank=5, world_size=8, local_rank=1cpu2, rank=6, world_size=8, local_rank=2现在每个节点各运行 4 个进程，rank 分配也正确了。
总结
在 Slurm 脚本中直接使用 export 无法为不同的 task 设置不同的环境变量值。 Slurm 提供的 SLURM_PROCID、SLURM_LOCALID 等变量是每个任务独有的，应该直接在 Python 代码中读取，而不是在 shell 脚本中提前 export。
看完这篇文章，你能否解释为什么在 知乎：在 Slurm 集群上使用 torchrun 进行分布式训练（单机多卡，多机多卡） 中要将 torchrun 单独写在一个脚本中，为什么设置 --ntasks-per-node=1 以及 --cpus-per-task=$((cpus_per_gpu * gpus_per_node))
支持与分享
如果这篇文章对你有帮助，欢迎分享给更多人或赞助支持！