Windows WSL2中使用 AMD 显卡的 GPU (ROCm) 加速运行 PyTorch 大模型

环境准备

• Windows 11 WSL2 (Ubuntu 22.04.5)
• AMD 显卡驱动 24.6.1及以上版本（本文以7900XTX为案例）
• 科学上网（自行研究）

首先是根据 AMD官网进行安装ROCm

需要注意的是这里只能安装该版本的 amdgpu-install

sudo apt update
wget https://repo.radeon.com/amdgpu-install/6.2.3/ubuntu/jammy/amdgpu-install_6.2.60203-1_all.deb
sudo apt install ./amdgpu-install_6.2.60203-1_all.deb

进行ROCm安装, 如果上面安装的不是这个版本的amdgpu-install,这里则会安装ROCm失败

amdgpu-install -y --usecase=wsl,rocm --no-dkms

检测是否安装成功

#一般显示显卡的信息则是安装成功，如下图
rocminfo

安装PyTorch

接下来安装适合ROCm的 PyTorch

#安装python
sudo apt install python3-pip -y
pip3 install