在初次接触 Slurm 集群时，我发现其启动分布式训练的方式与传统的多机集群存在一定差异，尤其在进行多机分布式训练时更为明显。为此，本文总结了我在 Slurm 集群上尝试分布式训练过程中积累的一些经验，希望对有类似需求的读者有所帮助。

在传统集群上进行分布式训练

在传统集群中，启动多机分布式训练的主要流程包括以下几个步骤：

保证各节点环境一致（例如配置相同的 Conda 环境）
配置节点间的正常通信（如设置 SSH 免密登录）
合理分发各节点所需的训练数据
分别在每台机器上手动启动训练脚本

由于本文主要聚焦于在 Slurm 集群环境下进行分布式训练的实践经验，而传统集群的分布式训练已有较多成熟的教程（例如这篇文章），此处不再赘述其详细流程。

Slurm 集群与传统集群的主要差别

下面的内容主要来自 ChatGPT 老师

相比传统集群，Slurm 集群在分布式训练的配置与启动流程上具有更高的自动化程度和更好的资源管理能力，主要体现在以下几个方面：

共享存储环境：在 Slurm 集群中，集群节点通常挂载有统一的共享文件系统，例如 NFS 或 Lustre。这样无需手动同步各节点的环境或数据，只需在共享目录中配置好训练代码、环境依赖和数据，各节点即可访问同一份资源，避免了手动分发数据和复制环境的繁琐操作。
自动执行训练脚本 ：通过编写 Slurm 作业脚本，可以使用 srun 或 sbatch 等命令统一提交训练任务。Slurm 会自动在分配到的节点上并行启动脚本，而无需用户登录每台机器手动运行，大大提高了使用效率和可重复性。
节点间通信自动配置 ：Slurm 会在作业启动时自动完成节点间的主机发现和通信配置，因此用户通常无需手动设置 SSH 免密登录。这为启动多节点分布式训练提供了便利，也降低了配置错误的风险。

总结一下 G 老师的话，就是经过合理配置，后续只需使用命令一键提交训练任务即可，节省了大量繁琐的重复工作

使用 `torchrun` 启动分布式训练

经过多年的发展，分布式训练的启动方式已日趋多样化，例如常见的有 torchrun、DeepSpeed 和 Hugging Face 的 accelerate 等工具。然而，无论采用哪种方式，至少都需要明确以下基本的分布式配置参数：

所使用的节点数
每个节点的 GPU 数量
当前节点在多节点中的编号（Node Rank）
主节点的 IP 地址以及用于通信的端口号

因此本文使用 PyTorch 原生的启动方式 torchrun 来作为例子进行介绍。

以一个由 2 个节点（node_0，node_1），每个节点 8 卡组成的集群为例，假设 node_0 作为主节点，在传统集群上，我们首先需要在 node_0 上执行：

torchrun \
    --master_addr="192.168.1.1" \
    --master_port="29500" \
    --nnodes=2 \
    --nproc_per_node=8 \
    --node_rank=0 \
    train.py

其中:

master_addr 为主节点的 ip 地址
master_port 为通信端口
nnodes 为训练任务所需的节点数
nproc_per_node 为每个节点的进程数，一般每个 GPU 对应一个进程
node_rank 为当前节点的编号

对于 node_1 节点，只需将 node_rank 设置为 1 即可，其余参数保持不变。

使用 Slurm 提交训练任务

我们已经知道，在传统集群中，我们通常需要手动登录每台机器，并执行上述命令启动训练任务。而在 Slurm 集群中，我们只需使用 srun 或 sbatch 等命令提交训练任务，Slurm 会自动在分配到的节点上启动训练任务。因此，为了在 Slurm 集群启动分布式训练，我们只需要知道如何获取上述参数，即可类似地使用 torchrun 启动训练任务。

Slurm 提交任务的流程

为了更好地理解 Slurm 是如何进行分布式训练的，这里简要介绍一下 Slurm 提交任务的流程：

使用 srun 或 sbatch 等命令提交训练任务, 指定所需计算资源
对于 sbatch 提交的任务，我们需要在 sbatch 脚本中使用 srun 将训练任务分发到各个节点上
在训练任务脚本中，使用 Slurm 提供的环境变量获取 torchrun 所需的参数，并启动训练任务

使用 `sbatch` 提交任务

根据前面的流程介绍，我们需要三个脚本：

提交任务的脚本：submit-job.sh
分发任务的脚本：job.sh
启动训练的脚本：train.sh

提交任务的脚本

下面是一个submit-job.sh的示例，使用时按需修改Config部分即可：

#!/usr/bin/bash
# ------------------------ Config ------------------------ #
job_name="example-job"        # 任务名称
partition="example-partition" # 分区名称
nnodes=2                      # 所需节点数
gpus_per_node=8               # 每个节点所需的 GPU 数量
cpus_per_gpu=16               # 每个 GPU 所需的 CPU 数量
quotatype="reserved"          # 任务类型, 如 `reserved`, `spot` 等
output_dir="slurm-outputs"    # 输出目录
# ------------------------ Setup ------------------------ #
timestamp=$(date +%Y%m%d_%H%M%S)
output_dir=${output_dir}/${timestamp}
export TIMESTAMP=${timestamp}
export OUTPUT_DIR=${output_dir}
mkdir -p ${output_dir}
# ------------------------ Submit ------------------------ #
sbatch \
    --job-name=${job_name} \
    --partition=${partition} \
    --ntasks-per-node=1 \
    --nodes=${nnodes} \
    --gres=gpu:${gpus_per_node} \
    --cpus-per-task=$((cpus_per_gpu * gpus_per_node)) \
    --output=${output_dir}/slurm-%j.log \
    --quotatype=${quotatype} \
    job.sh

分发任务的脚本

下面是一个job.sh的示例：

#!/usr/bin/bash
# ------------------ Get Number of GPUs ------------------ #
n_gpus=$(( $(echo "$SLURM_JOB_GPUS" | tr -cd ',' | wc -c) + 1 )) # 根据环境变量获取 GPU 数量
# ------------------- Setup Environment ------------------ #
export MASTER_ADDR=$SLURMD_NODENAME          # 获取主节点地址
export MASTER_PORT=$((RANDOM % 101 + 20000)) # 随机生成通信端口，防止端口冲突，这里使用 20000-20100 之间的随机数，可按需修改
export NNODES=$SLURM_JOB_NUM_NODES           # 获取节点数
export NPROC_PER_NODE=${n_gpus}              # 获取每个节点的 GPU 数量
# ------------------------- Main ------------------------- #
srun bash train.sh # 使用 `srun` 将启动训练的命令分发到各个节点上，即每个节点上都会执行 `bash train.sh`

启动训练的脚本

下面是一个train.sh的示例：

#!/usr/bin/bash
export PYTHONPATH=$(pwd):$PYTHONPATH
export NODE_RANK=${NODE_RANK:-$SLURM_NODEID} # 获取当前节点的编号
export OMP_NUM_THREADS=$NPROC_PER_NODE

torchrun \
    --master_addr=$MASTER_ADDR \
    --master_port=$MASTER_PORT \
    --nnodes=$NNODES \
    --nproc_per_node=$NPROC_PER_NODE \
    --node_rank=$NODE_RANK \
    train.py

分布式训练任务示例

下面是一个train.py的示例，为了简单起见，这里使用 all_reduce 操作来测试分布式通信是否正常：

import os
import torch
import torch.distributed as dist


def get_local_rank():
    gpus_per_node = int(os.environ["NPROC_PER_NODE"])
    rank = dist.get_rank()
    local_rank = rank % gpus_per_node
    return local_rank


def setup():
    dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo")


def cleanup():
    if dist.is_initialized():
        dist.destroy_process_group()


def train():
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = torch.device(f"cuda:{get_local_rank()}" if torch.cuda.is_available() else "cpu")
    # Initialize tensor
    tensor = torch.tensor([rank], dtype=torch.int32, device=device)
    print(f"[rank{rank}] tensor: {tensor}\n", end="")
    # Reduce tensor
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    # Calculate expected value
    expected = torch.tensor([sum(range(world_size))], dtype=torch.int32, device=device)
    # Compare reduced and expected values
    print(f"[rank{rank}] expected: {expected}, \n"
          f"[rank{rank}] reduced:  {tensor}\n", end="")
    if torch.allclose(tensor, expected):
        print(f"[rank{rank}] Success ✓\n", end="")
    else:
        print(f"[rank{rank}] Failed ✗\n", end="")
    dist.barrier()


def main():
    setup()
    train()
    if dist.get_rank() == 0:
        print("\n=================================\n"
              "Distributed test completed. \n"
              "If all processes show 'Success', \n"
              "the test is successful. \n"
              "=================================\n")
    cleanup()


if __name__ == "__main__":
    main()

有了上述所有脚本，我们只需使用 bash submit_job.sh 命令即可启动分布式训练任务。

使用 `srun` 提交任务

使用 srun 提交任务的流程与使用 sbatch 类似，只有小部分区别：

需要将 master_port 提前在提交任务的脚本确定
将 SLURM_JOB_GPUS 改为 SLURM_STEP_GPUS
由于 srun 提交任务时，会直接将任务分发到各个节点，因此需要使用 scontrol show hostnames 获取主节点地址

train.sh 和 train.py 与使用 sbatch 时相同，这里不再赘述。

提交任务的脚本

下面是一个srun-submit-job.sh的示例，同样地，使用时按需修改Config部分即可：

#!/usr/bin/bash
# ------------------------ Config ------------------------ #
job_name="example-job"
partition="example-partition"
nnodes=2
gpus_per_node=8
cpus_per_gpu=16
quotatype="reserved"
output_dir="slurm-outputs"
# ------------------------ Setup ------------------------ #
timestamp=$(date +%Y%m%d_%H%M%S)
output_dir=${output_dir}/${timestamp}
export TIMESTAMP=${timestamp}
export OUTPUT_DIR=${output_dir}
mkdir -p ${output_dir}
export MASTER_PORT=$((RANDOM % 101 + 20000)) # 提前确定通信端口
# ------------------------ Submit ------------------------ #
srun \
    --job-name=${job_name} \
    --partition=${partition} \
    --ntasks-per-node=1 \
    --nodes=${nnodes} \
    --gres=gpu:${gpus_per_node} \
    --cpus-per-task=$((cpus_per_gpu * gpus_per_node)) \
    --quotatype=${quotatype} \
    bash srun-job.sh

分发任务的脚本

下面是一个srun-job.sh的示例：

#!/usr/bin/bash
# ------------------ Get Number of GPUs ------------------ #
n_gpus=$(( $(echo "$SLURM_STEP_GPUS" | tr -cd ',' | wc -c) + 1 ))
# ------------------- Setup Environment ------------------ #
export MASTER_ADDR=$(scontrol show hostnames ${SLURM_JOB_NODELIST} | head -n 1)
export NNODES=$SLURM_JOB_NUM_NODES
export NPROC_PER_NODE=${n_gpus}
# ------------------------- Main ------------------------- #
bash train.sh

有了上述所有脚本，我们只需使用 bash srun-submit-job.sh 命令即可启动分布式训练任务。

结语

本文介绍了在 Slurm 集群上使用 sbatch 和 srun 提交分布式训练任务的流程，并给出了相应的示例脚本。

需要注意的是，本文所给出的示例脚本仅在本人所在的集群上进行过测试，其他集群可能需要根据实际情况进行调整，本文主要为实现原理提供参考。

WesKwong

做人最紧要嘅就系开心

在 Slurm 集群进行分布式训练

在传统集群上进行分布式训练

Slurm 集群与传统集群的主要差别

使用 `torchrun` 启动分布式训练

使用 Slurm 提交训练任务

Slurm 提交任务的流程

使用 `sbatch` 提交任务

使用 `srun` 提交任务

结语

添加新评论

在 Slurm 集群进行分布式训练

在传统集群上进行分布式训练

Slurm 集群与传统集群的主要差别

使用 torchrun 启动分布式训练

使用 Slurm 提交训练任务

Slurm 提交任务的流程

使用 sbatch 提交任务

使用 srun 提交任务

结语

添加新评论

使用 `torchrun` 启动分布式训练

使用 `sbatch` 提交任务

使用 `srun` 提交任务