如何通过UUID指定程序使用的CUDA设备

CUDA本身原生支持用UUID指定设备，完全适配MIG切分后的虚拟GPU场景，以下是两种常用方案：

方案1：无需修改代码，通过环境变量指定 #

这是最常用的方案，对所有CUDA兼容的程序都生效：

• 第一步先查询所有设备的UUID，执行命令：
  nvidia-smi -L
  输出结果中每个设备/MIG实例末尾括号内的字符串就是对应UUID，MIG实例的UUID通常以MIG-为前缀。
• 第二步启动程序时给CUDA_VISIBLE_DEVICES环境变量传入目标UUID即可，示例：
  CUDA_VISIBLE_DEVICES=MIG-GPU-b83a21c7-3d2f-4e1a-8b9c-0d1e2f3a4b5c/1/0 python your_train_script.py
  如果需要同时指定多个设备，多个UUID之间用英文逗号分隔即可。

方案2：代码内动态指定设备 #

如果需要在程序运行过程中动态选择对应UUID的设备，可以用对应框架的API实现：

• C/C++原生CUDA场景：调用cuDeviceGetByUUID()接口传入目标UUID获取设备编号，再调用cudaSetDevice()绑定设备即可。
• PyTorch场景：遍历所有可用设备匹配UUID后切换，示例代码：

import torch
target_uuid = "MIG-GPU-b83a21c7-3d2f-4e1a-8b9c-0d1e2f3a4b5c/1/0"
for device_idx in range(torch.cuda.device_count()):
    if str(torch.cuda.get_device_uuid(device_idx)) == target_uuid:
        torch.cuda.set_device(device_idx)
        break

• TensorFlow场景：通过tf.config.list_physical_devices('GPU')获取所有GPU设备的属性，匹配UUID后调用tf.config.set_visible_devices指定可见设备即可。

注意：CUDA_VISIBLE_DEVICES的优先级高于代码内的设备指定，如果环境变量已经配置了可见设备，代码只能在环境变量限定的设备范围内选择。如果传入的UUID不存在，CUDA会直接忽略该条目，不会回退到按设备序号匹配。

内容的提问来源于stack exchange，提问作者Daniel