如何通过os.environ["CUDA_VISIBLE_DEVICES"]让多个MIG分区GPU可见

问题说明 #

使用NVIDIA A100的MIG分区GPU时，尝试通过os.environ["CUDA_VISIBLE_DEVICES"]指定多个MIG设备UUID，出现以下问题：

• 单个UUID设置有效，但无法同时加载多个设备
• 逗号分隔多个UUID时无报错，但仅识别1个设备
• 元组/列表赋值会触发类型错误
• 常规物理GPU编号（如0,1）设置无效

GPU信息（nvidia-smi -L输出）：

GPU 1: NVIDIA A100-PCIE-40GB (UUID: GPU-b654bde8-a9d1-d27a-91eb-00000000000a) 
MIG 3g.20gb     Device  0: (UUID: MIG-e4a69dad-640c-5006-b7f6-00000000000c) 
MIG 1g.5gb      Device  1: (UUID: MIG-df1904ce-d118-5cc6-8f05-000000000007) 
MIG 1g.5gb      Device  2: (UUID: MIG-1b6f718c-a2db-59d5-a83d-00000000000a) 
MIG 1g.5gb      Device  3: (UUID: MIG-9882d1bb-3062-5d15-b0d6-000000000009) 
MIG 1g.5gb      Device  4: (UUID: MIG-198d257f-725f-529c-ac47-000000000004)

解决方案 #

1. 修正环境变量格式（核心修复） #

问题出在逗号分隔UUID时额外添加了空格，CUDA环境变量解析器会将带空格的整串视为单个设备标识，导致仅识别第一个UUID。正确的格式是直接用逗号连接UUID，无空格：

import os
# 必须在导入PyTorch/TensorFlow等CUDA依赖库前执行
os.environ["CUDA_VISIBLE_DEVICES"] = "MIG-1b6f718c-a2db-59d5-a83d-00000000000a,MIG-9882d1bb-3062-5d15-b0d6-000000000009,MIG-198d257f-725f-529c-ac47-000000000004"

from torch.cuda import device_count
print('Number of Devices: ', device_count())

执行后应该能正确识别3个MIG设备。

2. 关键注意事项 #

• 设置时机优先：环境变量必须在加载CUDA驱动或任何依赖CUDA的库之前设置，否则修改不会生效。
• 权限验证：确保当前用户对指定的MIG设备有访问权限，可通过nvidia-smi查看设备是否被占用或权限受限。
• 编号替代方案：若UUID使用麻烦，可尝试用MIG设备的本地编号（如0,1,2）代替UUID，但需确认编号对应目标设备（MIG设备编号与物理GPU编号独立）。

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