Pytorch 根据 MIG ID 调用分块后的 GPU

将一块 GPU 分成多块后，当在 pytorch 中调用时，使用 torch.device 只能调用 cuda:0，而这默认是第一张子块；如果使用 cuda:1 等不同序号，则会报错说没有此编号的 cuda。

在 GPU 分块后，子块的编号是根据 MIG id 确定的，所以要调用非第一张子块，则需要根据其 MIG id 指定 cuda。

方法

首先，使用 nvidia-smi -L 命令查看各 GPU 子块的 MIG id：

nvidia-smi -L
GPU 0: NVIDIA A100-SXM4-40GB (UUID: GPU-XXXXXXXXXX)
  MIG 3g.20gb     Device  0: (UUID: MIG-XXXXXXXXXX)
  MIG 3g.20gb     Device  1: (UUID: MIG-XXXXXXXXXX)

后面括号中的 UUID 就是在 pytorch 中使用的编号。

然后，设置环境变量：

os.environ['CUDA_VISIBLE_DEVICES'] = "MIG-XXXXXXXXXX"

这样，pytorch 就会将你指定的这张子块设置为默认 GPU，在之后的使用中，直接指定 torch.device('cuda:0') 就可以正常调用了。