MJUN Tech Note

GPUを複数搭載した計算機でGPUの順番を設定する

マルチ GPU を搭載したマシンで,nvidia-smiで見える GPU ID と, PyTorch や Tensorflow,CUDA_VISIBLE_DEVICESで指定する GPU ID が異なる現象についてのメモ.

結論から言うと,nvidia-smiでは PCI BUS 順に並ぶが, CUDA ではデフォルトでは早い順(FASTEST_FIRST)に GPU の ID が振られることが原因だった.

A6000 x 2,RTX 3090 x 2 が刺さっているマシンを想定し, nvidia-smiで並ぶ GPU の順番が以下のようだったとする.

0: RTX 3090
1: A6000
2: A6000
3: RTX 3090

すると,PyTorch 側では以下のように並ぶ.

0: A6000  # (nvidia-smiでは1)
1: A6000  # (nvidia-smiでは2)
2: RTX 3090  # (nvidia-smiでは0)
3: RTX 3090  # (nvidia-smiでは3)

PyTorch 側でもnvidia-smi(PCI バス順)の順番で GPU の ID を振りたい時は以下のように環境変数を設定する.

export CUDA_DEVICE_ORDER="PCI_BUS_ID"
# または,
CUDA_DEVICE_ORDER="PCI_BUS_ID" CUDA_VISIBLE_DEVICES="0" python train.py

これで,上記のマシンでは,nvidia-smiの順番となり,RTX 3090 が ID=0 となる.

コード上で変更するには,以下のようにする(Python)

import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 3"

ちなみに,PyTorch では,以下のコードで GPU の情報を表示できる

import torch

# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"

for i in range(torch.cuda.device_count()):
    info = torch.cuda.get_device_properties(i)
    print(f"CUDA:{i} {info.name}, {info.total_memory / 1024 ** 2}MB")

環境変数の設定は,torch.cuda 周りを呼び出す前に行わければならないのに注意.

ここまで書いてきて,常にPCI_BUS_ID順にすればいいじゃんと思うかもしれないが, Multi-GPU で異なる GPU を混ぜて推論するときには, GPU ID=0 で損失計算などを行うので,GPU ID=0 の GPU の計算が早いほうが嬉しい. そのため,早い順に GPU が並ぶのも利点が一応ある.

参考

1. Introduction — CUDA C Programming Guide
The programming guide to the CUDA model and interface.
1. Introduction — CUDA C Programming Guide favicon https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#env-vars