使用 SLURM 请求 2 个 GPU 并运行 1 个 python 脚本

Question

我正在尝试分配 2 个 GPU 并在这 2 个 GPU 上运行 1 个 python 脚本。 python 脚本需要变量 $AMBERHOME（通过 amber.sh 脚本获取）和 $CUDA_VISIBLE_DEVICES。 对于我请求的两个 GPU，$CUDA_VISIBLE_DEVICES 变量应该等于 0,1。

目前，我一直在试验这个基本脚本。

#!/bin/bash
#
#BATCH --job-name=test
#SBATCH --output=slurm_info
#SBATCH --nodes=2
#SBATCH --ntasks=2
#SBATCH --time=5:00:00
#SBATCH --partition=gpu-v100

## Prepare Run
source /usr/local/amber20/amber.sh
export CUDA_VISIBLE_DEVICES=0,1

## Perform Run
python calculations.py

当我运行脚本时，我可以看到请求了 2 个 GPU。

JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
11111   GPU         test   jsmith CF       0:02     2     gpu-[1-2]

当我查看 output ('slurm_info') 我看到了，

cpu-bind=MASK - gpu-1, task  0  0 [10111]: mask 0x1 set

当然还有关于失败工作的信息。

通常，当我在本地工作站上运行此脚本时，那里有 2 个 GPU，当在命令行中输入 nvidia-smi 时，我看到...

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.91.03    Driver Version: 460.91.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce RTX 208...  On   | 00000000:00:1E.0 Off |                    0 |
| N/A   29C    P0    24W / 300W |      0MiB / 16160MiB |      0%   E. Process |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  GeForce RTX 208...  On   | 00000000:00:1E.0 Off |                    0 |
| N/A   29C    P0    24W / 300W |      0MiB / 16160MiB |      0%   E. Process |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

但是，当我在集群上将 nvidia-smi 与我以前的批处理脚本一起使用时，我看到以下内容。

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.91.03    Driver Version: 460.91.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce RTX 208...  On   | 00000000:00:1E.0 Off |                    0 |
| N/A   29C    P0    24W / 300W |      0MiB / 16160MiB |      0%   E. Process |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

这让我觉得当 python 脚本运行时，它只会看到一个 GPU。

Answer 1

您正在请求两个节点，而不是两个GPU 。 请求 GPU 的正确语法取决于 Slurm 版本以及集群的设置方式。 但是您通常使用#SBATCH -G 2来请求两个 GPU。

Slurm 通常也会为您处理CUDA_VISIBLE_DEVICES ，所以不需要。 尝试这个：

#!/bin/bash
#
#BATCH --job-name=test
#SBATCH --output=slurm_info
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=2 #change that accordingly
#SBATCH -G 2
#SBATCH --time=5:00:00
#SBATCH --partition=gpu-v100

## Prepare Run
source /usr/local/amber20/amber.sh

## Perform Run
python calculations.py