在节点和GPU之间分布Pytorch进程

Question

我有两个节点，每个节点都有两个GPU。我想要运行4个进程，每个进程使用一个GPU。我使用了nccl（如果这有关系的话）。

这是我尝试过的Slurm脚本。我试了不同的设置组合。偶尔能按我想要的那样工作，但大多数时候，它会在一个节点上创建4个进程，并把2个进程分配给1个GPU。这会导致程序变慢，还可能出现内存不足的问题，导致all_gather失败。

我该如何正确分配进程呢？

#!/bin/bash
#SBATCH -J jobname
#SBATCH -N 2
#SBATCH --cpus-per-task=10

# version 1
#SBATCH --ntasks-per-node=2
#SBATCH --gres=gpu:2
#SBATCH --gpu-bind=none

# version 2
#SBATCH --ntasks-per-node=2
#SBATCH --gres=gpu:2

# version 3
#SBATCH --ntasks=4
#SBATCH --ntasks-per-node=2
#SBATCH --gres=gpu:2

# version 4
#SBATCH --ntasks=4
#SBATCH --ntasks-per-node=2
#SBATCH --gres=gpu:2
#SBATCH --gpus-per-task=1

# # version 5
#SBATCH --ntasks=4
#SBATCH --ntasks-per-node=2
#SBATCH --gpus-per-task=1

module load miniconda3
eval "$(conda shell.bash hook)"
conda activate gpu-env

nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST) )
nodes_array=($nodes)
head_node=${nodes_array[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)

echo Node IP: $head_node_ip
export LOGLEVEL=INFO

export NCCL_P2P_LEVEL=NVL

srun  torchrun --nnodes 2 --nproc_per_node 2 --rdzv_id $RANDOM --rdzv_backend c10d --rdzv_endpoint $head_node_ip:29678 mypythonscript.py  

在python脚本中：

dist.init_process_group(backend="nccl")
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

日志：

[W socket.cpp:464] [c10d] The server socket has failed to listen on [::]:29678 (errno: 98 - Address already in use).
[2024-03-31 15:46:06,691] torch.distributed.elastic.agent.server.local_elastic_agent: [INFO] log directory set to: /tmp/torchelastic_f6ldgsym/4556_xxbhwnb4
[2024-03-31 15:46:06,691] torch.distributed.elastic.agent.server.api: [INFO] [default] starting workers for entrypoint: python
[2024-03-31 15:46:06,691] torch.distributed.elastic.agent.server.api: [INFO] [default] Rendezvous'ing worker group
[W socket.cpp:464] [c10d] The server socket has failed to bind to 0.0.0.0:29678 (errno: 98 - Address already in use).
[E socket.cpp:500] [c10d] The server socket has failed to listen on any local network address.

我不确定这是否相关，因为在成功的情况下，我也看到了这些信息。

更新：我之前是按照pytorch的教程使用torchrun，按照这个教程，它就能正常工作。