gpu - 如何为每个任务设置 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 个任务，这是否是预期的行为？我们遗漏或误解了什么？

• salloc/srun 参数？
• slurm.conf 或 gres.conf 设置？

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

操作系统:CentOS 7

Slurm 版本:16.05.6

我们是否被迫使用 wrapper based methods为此？

slurm 版本(14 到 16)在 gpus 的分配方式上是否存在差异？

谢谢!

引用: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