使用SLURM进行多节点训练

本教程介绍了一个框架，关于如何使用在许多超级计算中心可用的SLURM工作负载管理器在多个节点上的多个GPU上进行分布式训练。 代码基于我们的单节点多GPU训练教程。 如果您不熟悉中的分布式训练概念，请先前往那里了解基础知识。

注意

本教程的完整脚本可以在examples/multi_gpu/distributed_sampling_multinode.py找到。 你可以在旁边找到示例*.sbatch文件，并根据你的需求进行调整。

用于管理启动的提交脚本

由于我们现在在多个节点上运行，我们不能再使用我们的__main__入口点并从那里启动进程。 这就是工作负载管理器的作用，它允许我们准备一个特殊的*.sbatch文件。 这个文件是一个标准的bash脚本，其中包含有关如何设置进程和环境的说明。

我们的示例从通常的shebang #!/bin/bash 开始，并带有特殊注释，指示SLURM系统应为我们的训练运行保留哪些资源。 具体配置通常取决于您的站点（以及您的使用限制！）。 以下是一个最小示例，适用于我们可用的相当不受限制的配置：

#!/bin/bash
#SBATCH --job-name=pyg-multinode-tutorial # identifier for the job listings
#SBATCH --output=pyg-multinode.log        # outputfile
#SBATCH --partition=gpucloud              # ADJUST this to your system
#SBATCH -N 2                              # number of nodes you want to use
#SBATCH --ntasks=4                        # number of processes to be run
#SBATCH --gpus-per-task=1                 # every process wants one GPU!
#SBATCH --gpu-bind=none                   # NCCL can't deal with task-binding...

此示例将在两个节点上各创建两个进程，每个进程保留一个GPU。

在接下来的部分中，我们必须为torch.distributed设置一些环境变量，以便正确执行集合过程。 理论上，你也可以在进程中设置这些变量：

export MASTER_PORT=$(expr 10000 + $(echo -n $SLURM_JOBID | tail -c 4))
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
echo "MASTER_ADDR:MASTER_PORT="${MASTER_ADDR}:${MASTER_PORT}

如果您不希望让您的脚本随机打开一个端口并监听传入的连接，您也可以使用共享文件系统上的文件。

现在唯一剩下要添加的就是训练脚本的执行：

srun python distributed_sampling_multinode.py

注意python调用是如何以srun命令为前缀的，因此将启动--ntasks副本。

最后，要将*.sbatch文件本身提交到工作队列中，请在您的shell中使用sbatch工具：

sbatch distributed_sampling_multinode.sbatch

使用配置了pyxis容器的集群

如果你的集群支持NVIDIA开发的pyxis插件，你可以使用一个现成的容器，该容器每月都会更新最新的NVIDIA和内容，更多信息请参见这里。 该容器设置了所有必要的环境变量，你现在可以直接从命令提示符使用srun运行示例：

srun --partition=<partitionname> -N<num_nodes> --ntasks=<number of GPUS in total> --gpus-per-task=1 --gpu-bind=none --container-name=pyg-test --container-image=<image_url> --container-mounts='.:/workspace' python3 distributed_sampling_multinode.py

请注意，--container-mounts='.:/workspace' 将当前文件夹（应包含示例代码）挂载到容器的默认启动文件夹 workspace 中。

如果您希望最终在容器中自定义包而无需访问docker（在公共HPC上非常可能），您可以按照本教程创建自己的镜像。

修改训练脚本

由于SLURM现在负责创建多个进程，并且我们无法共享任何数据（每个进程都将加载完整的数据集！），我们的__main__部分现在必须查询由SLURM或pyxis容器生成的进程设置的环境：

# Get the world size from the WORLD_SIZE variable or directly from SLURM:
world_size = int(os.environ.get('WORLD_SIZE', os.environ.get('SLURM_NTASKS')))
# Likewise for RANK and LOCAL_RANK:
rank = int(os.environ.get('RANK', os.environ.get('SLURM_PROCID')))
local_rank = int(os.environ.get('LOCAL_RANK', os.environ.get('SLURM_LOCALID')))
run(world_size, rank, local_rank)

torch.distributed.init_process_group() 函数现在将从环境中获取 MASTER_ADDR：

def run(world_size: int, rank: int, local_rank: int):
    dist.init_process_group('nccl', world_size=world_size, rank=rank)

我们还必须根据是否要将其用于节点本地目的（如选择GPU）或全局任务（如数据分割）来替换rank的使用

train_idx = data.train_mask.nonzero(as_tuple=False).view(-1)
train_idx = train_idx.split(train_idx.size(0) // world_size)[rank]

虽然我们需要将模型分配给节点本地的GPU，因此使用local_rank：

model = SAGE(dataset.num_features, 256, dataset.num_classes).to(local_rank)
model = DistributedDataParallel(model, device_ids=[local_rank])