AMD 显卡在 PyTorch 中 ROCm 版本选择与 torch 安装匹配验证

AMD 显卡跑 PyTorch 的关键并不只是把 torch 装上，而是装到和显卡驱动、ROCm 运行时匹配的那一个版本。ROCm 是 PyTorch 在 AMD 硬件上的底层计算库，PyTorch 官方的 pip 轮子对 ROCm 版本有固定绑定，选错版本最典型的表现是能导入 torch，但一调用 GPU 就报错，或者干脆识别不到设备。这篇内容围绕版本选择和安装后验证给出可操作的步骤。

AMD 显卡使用 PyTorch 时，需要先确认显卡在 ROCm 支持范围内，再选择 PyTorch 官方提供的对应 ROCm 版本轮子。安装后使用 torch.version.hip 和 torch.cuda.is_available() 验证匹配状态；系统驱动版本需要满足 ROCm 运行时要求，否则安装成功也可能无法调用 GPU。

一、先确认显卡是否在 ROCm 支持范围内

不是所有 AMD 显卡都能直接跑 ROCm。消费级显卡中，基于 CDNA 架构的数据中心卡（如 Instinct 系列）一般支持较好；Radeon 系列部分型号可以支持，但旧显卡或部分移动版芯片可能不在支持列表内。这一步没有替代办法，只能去 AMD ROCm 官方文档查当前支持列表。操作上，先确认显卡型号和代号，再比对支持列表。

适用场景：准备买卡或准备装环境前。
操作动作：查看系统识别到的 GPU 型号（Linux 下可执行 lspci | grep -i amd），再到 ROCm 文档对比。
验证方式：如果显卡在列表内，继续下一步；不在列表内，需要考虑更换硬件或使用 CPU 版本的 PyTorch。
风险边界：列表会随 ROCm 版本变化，旧版本支持的显卡型号可能被后续版本放弃，也可能新增型号，需要以你准备安装的 ROCm 版本对应支持列表为准。

二、选择与 PyTorch 轮子匹配的 ROCm 版本

PyTorch 官方发布的 ROCm 版本轮子，会在安装命令里明确标注对应的 ROCm 版本，通常形如 pip install torch `--index-url` https://download.pytorch.org/whl/rocm5.6。这里的 rocm5.6 是 ROCm 运行时版本，不是显卡驱动版本。选择时需要注意：

先确认你系统里安装的 ROCm 驱动版本是否能支撑这个运行时。可以执行 rocm-smi 查看当前 ROCm 驱动版本。
再确认 PyTorch 官方当前提供哪些 ROCm 版本的轮子。通常官方会同时维护几个较新的 ROCm 版本，较旧的版本可能不再提供。
如果系统驱动版本高于轮子要求的版本，一般兼容；如果低于要求，则可能报找不到 HIP 设备之类的错误。

安装时不要手动混合不同来源的 torch 和 torchvision，建议直接使用官方 index-url 安装，避免依赖冲突。例如：

pip install torch torchvision torchaudio \
  `--index-url` https://download.pytorch.org/whl/rocm5.6

这里的 5.6 仅作示意，实际版本号以你确认到的官方可用版本为准。安装完成后，不要立刻跑大型模型，先做最小验证。

三、安装后的匹配验证

验证的目的是确认三件事：PyTorch 能识别到 ROCm 库、显卡能被枚举出来、张量可以放到 GPU 上并且计算成功。建议按顺序执行：

验证驱动和 ROCm 运行时：

rocm-smi `--showdriverversion`

验证 PyTorch 版本信息和绑定的 ROCm 版本：

python -c "import torch; print('torch:', torch.__version__); print('hip:', torch.version.hip)"

验证 GPU 是否可用并做一次简单计算：

python -c "import torch; print('cuda available:', torch.cuda.is_available()); print('device name:', torch.cuda.get_device_name(0)); a = torch.tensor([1.0]).cuda(); print(a + 1)"

在 ROCm 版的 PyTorch 中，torch.cuda.is_available() 通常返回 True，因为底层计算图走的是 HIP。如果这里返回 False，说明 PyTorch 没有找到可用的 HIP 设备，需要检查驱动版本和轮子匹配关系。最后一行如果输出 tensor([2.]) 之类的结果，说明计算链路已经打通。

如果上述某一环节失败，不要急着重装整个系统，先按顺序排查：显卡驱动是否被系统正确加载 → ROCm 驱动版本是否满足要求 → PyTorch 轮子的 ROCm 版本是否与驱动兼容。通常重新安装对应版本的驱动或换一个 PyTorch 轮子能解决大多数问题。

四、常见问题与排查方向

问题1：PyTorch 官方没有我想要的 ROCm 版本轮子怎么办？

尽量使用官方提供的版本组合。如果目标 ROCm 版本没有对应轮子，可以先看当前驱动支持的 ROCm 版本，再选择接近的官方轮子，或者升级驱动到支持所需 ROCm 的版本。不建议自己从源码编译，除非你熟悉 ROCm 的构建过程。

问题2：使用 pip 安装成功后，import torch 不报错，但运行报设备错误

这通常是驱动版本过低或安装的轮子 ROCm 版本与驱动不匹配导致的。先运行 rocm-smi 确认驱动版本，再与 PyTorch 轮子要求的 ROCm 版本对比。注意，系统里可能同时存在多个 ROCm 版本，需要确认 PATH 和 LD_LIBRARY_PATH 指向的是正确的那一个。

问题3：用 conda 安装的 PyTorch ROCm 版是否更简单？

conda 会自动处理一些依赖，但同样存在版本匹配问题。建议安装后也执行上面的最小验证步骤。

整体上，AMD 显卡跑 PyTorch 的流程可以收敛为“查支持列表、选官方匹配版本、安装后最小验证”三步。每一步都有明确的检查和验证方法，按这个流程走，能减少大部分因版本错配带来的异常。