如何为每个任务 1 个 GPU 设置 slurm/salloc 但让作业使用多个 GPU？

Question

我们正在寻求一些有关 slurm salloc GPU 分配的建议。目前，给出：

% salloc -n 4 -c 2 -gres=gpu:1
% srun env | grep CUDA   
CUDA_VISIBLE_DEVICES=0
CUDA_VISIBLE_DEVICES=0
CUDA_VISIBLE_DEVICES=0
CUDA_VISIBLE_DEVICES=0

然而，我们希望使用的不仅仅是设备 0。
有没有办法用 srun/mpirun 指定 salloc 以获得以下内容？

CUDA_VISIBLE_DEVICES=0
CUDA_VISIBLE_DEVICES=1
CUDA_VISIBLE_DEVICES=2
CUDA_VISIBLE_DEVICES=3

这是期望的，以便每个任务获得 1 个 GPU，但总体 GPU 使用量分布在 4 个可用设备中（请参阅下面的 gres.conf）。不是所有任务都获得 device=0 的地方。

这样，每个任务就不会像当前的情况一样等待设备 0 释放其他任务。

或者即使我们有超过 1 个可用/空闲的 GPU（总共 4 个）用于 4 个任务，这是否是预期的行为？我们遗漏或误解了什么？

总结我们希望能够使用 slurm 和 mpi，以便每个等级/任务使用 1 个 GPU，但作业可以在 4 个 GPU 之间分散任务/等级。目前看来我们仅限于设备 0。我们还希望避免由于 mpi 使用而在 salloc/sbatch 中多次提交 srun。

操作系统：CentOS 7

Slurm 版本：16.05.6

我们是否被迫为此使用基于包装的方法？

slurm 版本（14 到 16）在 GPU 分配方式上是否有差异？

谢谢你！

参考：gres.conf

Name=gpu File=/dev/nvidia0
Name=gpu File=/dev/nvidia1
Name=gpu File=/dev/nvidia2
Name=gpu File=/dev/nvidia3

Answer 1

首先，尝试请求四个 GPU

% salloc -n 4 -c 2 -gres=gpu:4

对于--gres=gpu:1，所有任务仅看到一个 GPU 是预期的行为。对于--gres=gpu:4，输出将是

CUDA_VISIBLE_DEVICES=0,1,2,3
CUDA_VISIBLE_DEVICES=0,1,2,3
CUDA_VISIBLE_DEVICES=0,1,2,3
CUDA_VISIBLE_DEVICES=0,1,2,3

为了得到你想要的，你可以使用包装脚本，或者修改你的 srun 命令，如下所示：

srun bash -c 'CUDA_VISIBLE_DEVICES=$SLURM_PROCID env' | grep CUDA

那么你会得到

CUDA_VISIBLE_DEVICES=0
CUDA_VISIBLE_DEVICES=1
CUDA_VISIBLE_DEVICES=2
CUDA_VISIBLE_DEVICES=3