ROCm简单入门 - 使用AMD显卡加速PyTorch

ROCm 在如今已经成为继 CUDA 之后，第二大 GPU 并行计算平台，就 PyTorch 而言，PyTorch 的 ROCm 版本在 Python 应用程序接口层面使用了相同的语义所以从现有的代码迁移到 ROCm 版本的 PyTorch 几乎不需要进行任何修改。尽管 ROCm 可能相比 CUDA 存在一定的性能损失，但 AMD GPU 以相对较低的硬件价格使得 AMD+ROCm 的搭配成为人工智能方面不二的性价比之选

本文 MWX 将简单介绍如何在 AMD GPU 上使用 ROCm 加速 PyTorch，并补充部分官方教程中遗漏的部分细节，让我们开始吧~

硬件/系统配置

MWX 使用的是 AMD Radeon RX 7800 XT + AMD Ryzen R5 9600X + 32GB DDR5 的配置，该配置仅供参考，具体硬件需求请参考AMD 的官方文档

系统方面，AMD 官方支持 Ubuntu,Red Hat Enterprise Linux(RHEL),SUSE Linux Enterprise Server(SLES) 三大主流商业 Linux 发行版，与其颇有关系的发行版如 Linux Mint,Rocky Linux,OpenSUSE 等大概率也可以正常使用，但 AMD 官方看上去更希望用户使用 Ubuntu（不少文档只提供 Ubuntu 版本），MWX 因此在此为了避免潜在的问题也选择了基于 Ubuntu 22.04 LTS 的 Zorin OS 17.2（主要是长得好看ヾ(≧▽≦*)o）

*注:ROCm 暂不支持 Windows，如需在 Windows 平台上使用 ROCm，需借助 WSL2，这部分请直接参考AMD 官方文档

ROCm 安装

ROCm 的安装实际上非常简单，参考 AMD 的官方文档
在 Ubuntu 上，直接执行以下命令即可完成安装：

sudo apt update
sudo apt install "linux-headers-$(uname -r)" "linux-modules-extra-$(uname -r)"
sudo usermod -a -G render,video $LOGNAME # 将当前用户添加至 render 和 video 组以便无需 root 权限即可访问 AMD GPU
wget https://repo.radeon.com/amdgpu-install/6.2.2/ubuntu/jammy/amdgpu-install_6.2.60202-1_all.deb #jammy 为 Ubuntu 22.04的代号，对于 Ubuntu 24.04 及其衍生版，请将 jammy 替换为 noble
sudo apt install ./amdgpu-install_6.2.60202-1_all.deb
sudo apt update
sudo apt install amdgpu-dkms rocm

在此之后，还需做一些额外的配置

配置 ld

sudo tee --append /etc/ld.so.conf.d/rocm.conf <<EOF
/opt/rocm/lib
/opt/rocm/lib64
EOF
sudo ldconfig

将 ROCm 的可执行文件添加至 PATH

Plan A:使用 update-alternatives
大多数 Linux 发行版都有 update-alternatives 工具。它有助于管理命令或程序的多个版本。有关 update-alternatives 的更多信息，请参阅 Linux man 文档。
使用以下指令完成配置：
1. 列出所有被支持的 ROCm 命令：
```
update-alternatives --list rocm
```
1. 如果安装了多个 ROCm 版本，update-alternatives 会选择使用最新版本。如需指定想要使用的 ROCm 版本，请使用此命令：
```
update-alternatives --config rocm
```
Plan B:使用 environment-modules
environment-modules 工具简化了 shell 初始化。它允许你使用模块文件修改会话环境。更多信息，请参阅 Environment Modules 文档。
使用以下指令完成配置：
1. 列出可用的 ROCm 版本：
```
module avail
```
2.如果安装了多个 ROCm 版本，使用以下命令选择所需的版本
```
module load rocm/<version>
```
Plan C:手动配置
ROCm 模块文件位于 /opt/rocm-/lib/rocmmod 目录下，
如果以上方法均无法满足需求，可手动将 ROCm 的可执行文件添加至 PATH
如，在 .bashrc 中添加以下内容：
```
export PATH=$PATH:/opt/rocm-6.2.2/bin
```

验证内核驱动程序，ROCm，软件包安装状态

dkms status
rocminfo
clinfo
apt list --installed #这一步可能会列出大量已安装的软件包

重启以确保 ROCm 配置生效

reboot

PyTorch 安装

AMD 官方推荐使用 docker 镜像，以方便管理，可参考AMD 的官方文档
在这里由于 MWX 懒得装 Docker，直接选择 pip 安装
这里可以直接按照 PyTorch 官网指引，执行以下命令安装

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.2

如果存在网络问题，可以考虑使用下载工具下载对应的 whl 文件，再使用 pip 安装
如果不出意外，PyTorch 就成功安装啦~
我们可以简单验证一下，在 Python 中执行以下指令验证 PyTorch 是否成功安装：

import torch
x = torch.rand(5, 3)
print(x)

输出应该与下文类似：

tensor([[0.3380, 0.3845, 0.3217],
        [0.8337, 0.9050, 0.2650],
        [0.2979, 0.7141, 0.9069],
        [0.1449, 0.1132, 0.1375],
        [0.4675, 0.3947, 0.1426]])

在 Python 中执行以下指令验证 ROCm 是否正常工作：

import torch
torch.cuda.is_available()

如果得到了 True，那么恭喜你，至此大功告成
但如果很不幸，ROCm 不可用，可以继续往下看
执行以下指令，查看日志，并尝试找出可疑的输出，并善用搜索

export AMD_LOG_LEVEL=7
python -c "import torch;print(torch.cuda.is_available())"

值得一提的是，如果rocm-smi等工具无异常，有很大可能由于用户不在 render 组内，执行以下命令重新添加用户至 render 及 video 组

sudo usermod -a -G render,video $LOGNAME

完成后请重启系统

reboot