一个由GPU多分区导致的Hyper-V虚拟机无法正常使用Pytorch&TensorFlow GPU问题

前言

在Hyper-V虚拟机上，可以正常通过 lspci 查看设备，nvidia-smi 也可以正常输出，但在使用 PyTorch 和 TensorFlow 时，GPU 模式无法正常工作。本文将介绍故障排查的过程和解决方案。

问题排查

某天发现自己在Hyper-V的Ubuntu虚拟机上突然无法正常使用GPU，重装了Python环境、PyTorch和TensorFlow包、NVCC等工具后问题依旧存在。

在StackOverflow中发现了一个解决方法，通过在Python代码中执行以下命令：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0"

或者在shell中设置环境变量：

export CUDA_VISIBLE_DEVICES="0"

然后再运行以下Python代码，即可正常输出：

import torch
print(torch.cuda.is_available())

# 或
import tensorflow as tf
print(tf.config.list_physical_devices("GPU"))

上述命令控制了哪些GPU设备对程序可见。既然通过控制特定GPU设备的方法能使程序正常执行，基本可以排除环境和驱动等问题。

在Linux中输入以下命令，可以查看现有的PCI设备。在只有一个GPU设备的机器上，正常情况下只会输出一条内容，但在我的虚拟机中却得到了两条：

lspci

在宿主机上输入以下命令查看分配给Hyper-V虚拟机的GPU分区，正常情况下也应该看到一段信息：

Get-VMGpuPartitionAdapter

因此，可以发现问题在于GPU分区时错误地重复分配，导致在虚拟机里始终有一个GPU分区无法使用。而当执行PyTorch或TensorFlow时，它们默认使用所有的GPU分区，从而导致报错。

问题修复

删除GPU分区

Remove-VMGpuPartitionAdapter -VMName <VMName>

重新分配GPU分区

Add-VMGpuPartitionAdapter -VMName <vmname>

随后重启虚拟机，再次输入 lspci，即可看到正确的PCI设备数量。

再次运行以下Python代码，即可正常输出：

import torch
print(torch.cuda.is_available())

# 或
import tensorflow as tf
print(tf.config.list_physical_devices("GPU"))