pytorch - SLURM 和 Pytorch 如何一起处理多节点多 GPU 训练

Question

我想问一下，在使用 Slurm 管理工作负载的集群中使用多节点多 GPU 进行训练时，梯度如何聚合。代码是使用 Pytorch 编写的。

例如，当train.py在 Slurm 中启动脚本以在 4 个节点上训练模型时，每个节点有 4 个 GPU，如下所示，该srun命令究竟做了什么？

srun python train.py

它写在srun用于在由 Slurm 管理的集群上运行并行作业的文档中。但是，我想知道它执行了哪些步骤。我想它将脚本复制到 16 个 GPU，每个 GPU 将同时运行脚本。那是对的吗？

如果我torch.distributed.init_process_group()在脚本中使用处理多 GPU 训练，Slurm 将如何处理从每个 GPU 和 Pytorch 收集的梯度？我假设有一个主端口（它是由 Slurm 分配的节点中的一个 GPU 设备）来收集梯度。由于这个torch.distributed包，4 个节点中的所有 16 个 GPU 都会将梯度传输到这个主端口？Slurm 有助于分配资源，但它在梯度聚合中没有任何作用？

我也不确定是否应该srun像上面那样启动脚本，还是应该torch.distributed.launch在我的命令中指定如下。我想确保正确收集渐变。

# NGPU equals to number of GPUs/node
export NGPU=4
srun python -m torch.distributed.launch --nproc_per_node=$NGPU train.py

预先感谢您的帮助！

标签： pytorchdistributed-computingdistributed-systemslurm