如何為每個任務設置 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）。 不是所有任務都得到設備=0。

這樣每個任務就不會像當前那樣等待設備 0 從其他任務中釋放出來。

或者即使我們有超過 1 個可用/免費的 GPU（總共 4 個）用於 4 個任務，這是否是預期的行為？ 我們錯過了什么或誤解了什么？

總結我們希望能夠使用 slurm 和 mpi，這樣每個 rank/task 使用 1 gpu，但作業可以在 4 gpus 之間傳播任務/ranks。 目前看來我們僅限於設備 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

Answer 2

此功能計划在 19.05 推出。 有關詳細信息，請參閱https://bugs.schedmd.com/show_bug.cgi?id=4979 。

請注意，如果您的作業沒有請求該節點上的所有 GPU，建議的“srun bash...”解決方案將中斷，因為另一個進程可能控制 GPU0。

Answer 3

要為每個任務完成一個 GPU，您需要使用srun命令的--gpu-bind開關。 例如，如果我有三個節點，每個節點有 8 個 GPU，並且我希望每個節點運行 8 個任務，每個任務都綁定到一個唯一的 GPU，則以下命令可以解決問題：

srun -p gfx908_120 -n 24 -G gfx908_120:24 --gpu-bind=single:1  -l bash -c 'echo $(hostname):$ROCR_VISIBLE_DEVICES'