如何为每个任务设置1个gpu的slurm / salloc但让工作使用多个gpus？

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上等待从其他任务中解脱出来，就像目前的情况一样。

或者这是预期的行为，即使我们有4个任务可用/免费（总共4个）超过1个gpu？我们缺少什么或误解了什么？

总结我们希望能够使用slurm和mpi这样每个等级/任务使用1个gpu，但是这个工作可以在4个gpus中传播任务/等级。目前看来我们仅限于设备0。我们还希望避免因使用mpi而在salloc / sbatch中进行多次srun提交。

操作系统：CentOS 7

Slurm版本：16.05.6

我们被迫使用wrapper based methods吗？

在如何分配gpus方面，与slurm版本（14到16）存在差异吗？

谢谢！

参考：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

Answer 2

此功能计划于19.05。有关详细信息，请参阅https://bugs.schedmd.com/show_bug.cgi?id=4979。

请注意，如果您的工作没有请求该节点上的所有GPU，那么建议的'srun bash ...'解决方案将会中断，因为另一个进程可能控制GPU0。