使用 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