WSL2中配置AMD GPU运行DALL-E Playground故障排查

数据老陈

2026-8-30

问题环境说明

硬件配置：AMD Radeon RX 6600 XT 显卡、AMD Ryzen 5 3600XT 处理器、32GB 内存
宿主机系统：Windows 11 Pro 21H1（OS 内部版本 22000.675）
虚拟化层：WSL2，内核版本 5.10.16.3-microsoft-standard-WSL2
子系统：Ubuntu 20.04 LTS
部署目标：本地部署 DALL-E Playground

现存故障现象

项目前后端已成功部署，但仅能调用CPU运行，服务启动时输出警告日志：
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
已执行的GPU适配操作：
- 按PyTorch官方指引执行pip3 install torch torchvision --extra-index-url https://download.pytorch.org/whl/rocm4.5.2，安装适配ROCm 4.5.2版本的PyTorch，官方示例代码验证PyTorch安装状态正常
- 参考网络教程安装rock-dkms包，同时安装了全套ROCm 5.1.0版本组件（含rocm-dkms 5.1.0、rocminfo、rocm-opencl等），以及libdrm-amdgpu1、libosdgpu3.4.0依赖
ROCm检测异常结果：
1. 执行/opt/rocm/bin/rocminfo返回报错：ROCk module is NOT loaded, possibly no GPU devices
2. 执行/opt/rocm/opencl/bin/clinfo可识别到AMD Accelerated Parallel Processing OpenCL平台，但可用设备数量为0
其他GPU检测结果：执行glxinfo -B可识别到经D3D12映射的AMD Radeon RX 6600 XT设备，显示硬件加速已开启（显存识别数值存在偏差），glmark2基准测试可正常调用GPU运行
待确认问题：无法定位故障根因为WSL侧AMD GPU透传配置错误，还是DALL-E Playground程序本身适配问题，已知ZLUDA方案可解决同类问题但尚未测试。

排查与解决思路

当前glxinfo、glmark2可正常识别GPU并运行，说明WSL基础GPU半虚拟化透传功能正常，故障全部来自软件配置层面，不需要调整WSL或Windows的GPU透传基础开关，按以下步骤排查即可：

第一步：修正WSL2环境下ROCm的错误安装逻辑
WSL2使用微软定制的半虚拟化内核，不需要安装任何dkms类内核模块包，rocm-dkms、rock-dkms是为原生Linux系统准备的，在WSL中安装不仅无效，还会和WSL自带的GPU透传驱动冲突。首先执行命令彻底卸载错误安装的dkms包：
sudo apt remove --purge rocm-dkms rock-dkms
其次当前使用的WSL内核版本为5.10.16.3，版本过低不支持ROCm用户态组件透传，需要在宿主机以管理员权限打开PowerShell，执行wsl --update将WSL内核升级到5.10.43.2及以上版本，升级完成后执行wsl --shutdown重启WSL实例生效。
第二步：统一ROCm组件版本，消除版本冲突
当前安装的PyTorch适配ROCm 4.5.2，但系统级ROCm组件为5.1.0版本，版本不匹配是设备识别失败的核心原因之一。二选一统一版本即可：
- 方案A：将系统ROCm组件降级到4.5.2版本，卸载现有5.1.0版本所有ROCm包后安装4.5.2版本的rocm-libs、rocm-opencl、rocminfo用户态组件，不要安装任何dkms相关包
- 方案B：将PyTorch升级到适配ROCm 5.1.0的版本，替换现有4.5.2版本的PyTorch、torchvision安装包
  安装完成后将ROCm路径写入shell环境变量，编辑~/.bashrc追加以下内容：
```
export PATH=$PATH:/opt/rocm/bin:/opt/rocm/profiler/bin:/opt/rocm/opencl/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/rocm/lib:/opt/rocm/opencl/lib
```
执行source ~/.bashrc生效后，重新运行/opt/rocm/bin/rocminfo，正常情况下即可识别到RX 6600 XT设备。
第三步：验证深度学习框架GPU可用性
rocminfo识别正常后，进入Python交互环境执行以下命令验证PyTorch GPU调用能力：
```
import torch
print(torch.cuda.is_available())
```
ROCm版PyTorch会将AMD GPU映射为CUDA兼容设备，返回True即为PyTorch侧GPU调用正常。
注意启动日志中的CPU fallback警告来自TensorFlow/JAX组件，当前仅安装了ROCm版PyTorch，未安装对应ROCm版本的TensorFlow/JAX，需要安装和ROCm版本匹配的TensorFlow、JAX包，否则这两个组件仍会自动回退到CPU运行。
第四步：程序侧配置兜底
框架层GPU验证通过后，启动DALL-E Playground前手动添加环境变量，避免自动检测失效：
```
export JAX_PLATFORMS=rocm
export XLA_PYTHON_CLIENT_PREALLOCATE=false
```
再启动后端服务，观察日志是否还有GPU缺失的警告。
第五步：ZLUDA方案兜底
如果以上原生ROCm配置完成后仍存在设备识别问题，可直接采用ZLUDA方案：不需要安装完整ROCm组件，安装ZLUDA后直接使用标准CUDA版本的PyTorch、TensorFlow、JAX即可，ZLUDA会自动将CUDA调用转换为AMD GPU可执行的指令，在WSL2环境下对RX 6000系显卡的兼容性普遍高于原生ROCm，且不需要修改项目代码。

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