如何指定PyTorch脚本使用特定的GPU单元？

Question

我有一个 Python 训练脚本，它利用 CUDA GPU 来训练模型（Kohya Trainer 脚本可用此处）。遇到内存不足错误：

OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 1; 23.65 
GiB total capacity; 144.75 MiB already allocated; 2.81 MiB free; 146.00 MiB 
reserved in total by PyTorch) If reserved memory is >> allocated memory try 
setting max_split_size_mb to avoid fragmentation.  See documentation for Memory 
Management and PYTORCH_CUDA_ALLOC_CONF

经过调查，我发现脚本使用的是 GPU 单元 1，而不是单元 0。单元 1 目前使用率很高，GPU 内存剩余不多，而 GPU 单元 0 仍然有足够的资源。如何指定脚本使用 GPU 单元 0？

即使我改变了：

text_encoder.to("cuda")

至：

text_encoder.to("cuda:0")

脚本仍在使用 GPU 单元 1，如错误消息中所指定。

nvidia-smi

的输出：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.60.11    Driver Version: 525.60.11    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| 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  NVIDIA GeForce ...  Off  | 00000000:81:00.0 Off |                  Off |
| 66%   75C    P2   437W / 450W |   5712MiB / 24564MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce ...  Off  | 00000000:C1:00.0 Off |                  Off |
| 32%   57C    P2   377W / 450W |  23408MiB / 24564MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1947      G   /usr/lib/xorg/Xorg                  4MiB |
|    0   N/A  N/A     30654      C   python                           5704MiB |
|    1   N/A  N/A      1947      G   /usr/lib/xorg/Xorg                  4MiB |
|    1   N/A  N/A     14891      C   python                          23400MiB |
+-----------------------------------------------------------------------------+

---

更新1

同一个笔记本可以看到2个GPU单元：

import torch
for i in range(torch.cuda.device_count()):
    print(torch.cuda.get_device_properties(i))

输出：

_CudaDeviceProperties(name='NVIDIA GeForce RTX 4090', major=8, minor=9, total_memory=24217MB, multi_processor_count=128)
_CudaDeviceProperties(name='NVIDIA GeForce RTX 4090', major=8, minor=9, total_memory=24217MB, multi_processor_count=128)

---

更新2

设置

CUDA_VISIBLE_DEVICES=0

会导致此错误：

运行时错误：CUDA 错误：设备序号无效

Answer 1

我希望这有帮助： 当我为模型分配特定 GPU 时，我发现 Nvidia-Ami 输出中的 GPU 索引与 cuda 索引不匹配。 我的 Tesla P40 出现在 Nvidia-smi 输出中的索引 0 上，但在 pytorch 代码中引用为“cuda:2”）或 CUDA_VISIBLE_DEVICES=2。