AMD ROCm 实战：PyTorch 迁移与性能优化指南

1. 从“显卡预算砍半”说起：ROCm 到底解决了谁的痛点

如果你最近在配深度学习工作站，大概率经历过这样的心理斗争：一张 24GB 显存的卡，价格能顶一台不错的整机；而同样显存容量的另一家方案，价格可能只要一半甚至更低。问题在于，过去几年里，“便宜大显存”这条路在 PyTorch 生态里几乎走不通——驱动、算子库、框架适配层层设卡，最后往往以“还是加钱买省心”收场。

AMD ROCm 这几年的变化，恰恰是在拆这道墙。它的核心承诺可以用一句话概括：让 PyTorch 代码在 AMD GPU 上跑起来，尽量不改代码。注意这里的措辞，是“尽量不改”，而不是“完全不用改”。标题里说“一行不用改”，在绝大多数标准训练和推理场景下是成立的，但前提是你把环境这一层做对了。环境没做对，改一百行代码也救不回来。

这篇文章面向三类人：一是手里已经有 AMD 显卡、想拿来跑 PyTorch 的个人开发者；二是正在做整机选型、纠结要不要为生态溢价买单的团队技术负责人；三是单纯想搞清楚 ROCm 和 CUDA 到底差在哪、值不值得迁移的工程师。我会把环境搭建、代码适配、性能实测、踩坑排查这几件事讲透，让你看完能直接动手，而不是停留在“听说 AMD 也能跑”的层面。

先给一个结论性的判断：ROCm 现在的成熟度，已经跨过了“能跑”的门槛，进入了“好用”和“挑场景”并存的阶段。标准 CNN、Transformer 训练，主流推理框架，大部分情况下体验接近 CUDA；但涉及冷门算子、自定义 CUDA 扩展、特定版本组合时，仍然需要你有一定的排查能力。这篇文章的价值，就是帮你把“挑场景”这部分提前摸清楚。

2. ROCm 与 CUDA 的对应关系：先搞懂映射，再谈迁移

2.1 软件栈的层级对照

很多人第一次接触 ROCm 会被一堆缩写搞晕：HIP、rocBLAS、MIOpen、RCCL……其实只要把它和 CUDA 生态做一一映射，立刻就清楚了。

这张表是理解 ROCm 的钥匙。你会发现，HIP 的定位就是“GPU 编程的可移植层”。它提供了一套语法接近 CUDA 的 API，同时配套一个hipify工具，能把 CUDA 源码自动转成 HIP 源码。PyTorch 官方对 AMD 的支持，本质上就是维护了一套 HIP 版本的算子实现，在编译期把 CUDA 调用映射到 HIP 调用。

所以当你在 AMD 卡上import torch并调用torch.cuda.is_available()时，返回 True 并不矛盾——PyTorch 保留了torch.cuda这个命名空间作为“GPU 抽象层”，底层实际走的是 HIP。这也是为什么大量现有代码不用改就能跑：框架层已经把差异吃掉了。

2.2 为什么“代码不用改”能成立

关键在于 PyTorch 的设备抽象做得足够干净。你的代码里写的通常是这几类：

• device = torch.device("cuda")
• model.to(device)
• tensor.cuda()
• torch.cuda.amp混合精度

这些调用在 ROCm 版本的 PyTorch 里全部被保留，底层转发到 HIP。也就是说，你不需要把cuda改成hip，改了反而可能报错。这一点和很多人直觉相反，值得单独强调。

真正需要改代码的场景只有两类：一是你写了自定义 CUDA 扩展（.cu文件），需要用hipify转换；二是你用到了某些尚未在 ROCm 上实现的算子，需要找替代实现。前者是编译期的事，后者是算子覆盖度的事，和“日常写模型代码”基本无关。

2.3 版本对应关系是迁移的第一道坎

ROCm 和 PyTorch 的版本绑定比 CUDA 更紧。CUDA 那边，你装个 12.x 的驱动，往往能跑好几个 PyTorch 版本；ROCm 这边，PyTorch 官方 wheel 是针对特定 ROCm 版本编译的，装错了就是各种undefined symbol或者直接 import 失败。

常见的对应关系大致是这样：ROCm 5.7 对应 PyTorch 2.0/2.1，ROCm 6.0 对应 PyTorch 2.2/2.3，ROCm 6.1/6.2 对应更新的版本。具体以 PyTorch 官网的安装矩阵为准，不要凭记忆装。我见过太多人卡在这一步，折腾一整天以为是显卡问题，其实是 wheel 和 ROCm 版本错配。

提示：安装前先跑rocm-smi确认驱动和 ROCm 版本，再对照 PyTorch 官网选择对应的安装命令。这一步花五分钟，能省你五小时。

3. 环境搭建实操：从裸机到torch.cuda.is_available()返回 True

3.1 系统选择：为什么我更推荐 Ubuntu 22.04

ROCm 官方支持的系统列表里，Ubuntu 22.04 LTS 是体验最顺的。原因很实际：内核版本、驱动包、依赖库的兼容性测试最充分。如果你用 CentOS 或者较新的非 LTS 发行版，很可能遇到内核模块编译失败、amdgpu驱动加载异常这类问题。

Windows 用户要注意，ROCm 在 Windows 上的支持一直落后于 Linux。虽然近两年有进展，但生产级训练仍然强烈建议走 Linux。如果你只有 Windows 机器，双系统或者 WSL2 是可行路径，但 WSL2 下的 ROCm 支持有额外限制，性能也有损耗，适合尝鲜不适合压榨性能。

至于 macOS，AMD 显卡在 Mac 上的 ROCm 支持基本可以忽略，不在本文讨论范围。

3.2 驱动与 ROCm 安装的完整链路

安装分两大块：内核驱动（amdgpu）和用户态运行时（ROCm）。顺序不能乱。

第一步，确认你的显卡在支持列表里。lspci | grep -i amd能列出显卡型号。如果这条命令没反应，说明pciutils没装，先sudo apt install pciutils。注意，很多人在这一步卡住是因为命令拼写或者权限问题，不是硬件问题。

第二步，安装内核驱动。官方推荐用amdgpu-install脚本，它会处理 DKMS 模块编译。命令大致是：

sudo apt update sudo apt install wget gnupg2 wget https://repo.radeon.com/amdgpu-install/<version>/ubuntu/jammy/amdgpu-install_<version>_all.deb sudo apt install ./amdgpu-install_<version>_all.deb sudo amdgpu-install --usecase=rocm

这里的<version>要替换成你实际下载的版本号。--usecase=rocm表示只装 ROCm 相关组件，不装图形驱动那套，避免和桌面环境冲突。

第三步，把当前用户加入render和video组，否则rocm-smi会报权限错误：

sudo usermod -aG render,video $USER

改完组要重新登录才生效。这一步极其容易被忽略，然后你会看到rocm-smi提示 “No GPU found”，误以为驱动没装好。

第四步，验证。rocm-smi能正常输出显卡温度、显存占用、功耗，说明驱动层通了。rocminfo能看到 agent 信息，说明运行时也通了。

3.3 PyTorch 安装：选对 wheel 比什么都重要

到这一步，安装 PyTorch 反而简单了，前提是你选对了命令。以 ROCm 6.x 为例，官方给的安装命令形如：

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

注意--index-url指向的是 ROCm 专用源，不是默认的 PyPI。如果你直接pip install torch，装到的是 CPU 版本或者 CUDA 版本，torch.cuda.is_available()必然返回 False。

装完验证三连：

import torch print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0)) print(torch.version.hip)

第一行返回 True，第二行打印出你的 AMD 显卡型号，第三行显示 HIP 版本号，就说明整条链路通了。如果第一行是 False，八成是 wheel 装错了；如果第二行报错，可能是设备权限问题；如果第三行是 None，说明装的是非 ROCm 版本。

3.4 一个容易被忽略的细节：显存分配策略

AMD 显卡在 Linux 下的显存管理和 NVIDIA 有差异。默认情况下，amdgpu驱动可能会把一部分系统内存当作显存池，这在跑大模型时是好事，但在监控显存占用时会让你困惑——rocm-smi显示的显存数字和torch.cuda.memory_allocated()对不上。

如果你需要精确控制显存，可以设置环境变量HSA_OVERRIDE_GFX_VERSION来适配某些不在官方列表里的显卡型号。这个变量在社区里被广泛使用，尤其是 RDNA 架构的消费级卡。用法是：

export HSA_OVERRIDE_GFX_VERSION=11.0.0

具体填什么值取决于你的显卡架构，RDNA3 通常是 11.0.0，RDNA2 是 10.3.0。填错会导致算子编译失败或者结果异常，所以只在官方不支持你的卡时才用，并且要验证输出正确性。

4. 代码迁移的真实工作量：哪些要改，哪些别动

4.1 标准训练脚本：基本零改动

我拿一个标准的 ResNet 训练脚本做过测试，从 CUDA 机器迁到 ROCm 机器，唯一需要动的是安装命令，代码本身一行没改。torch.device("cuda")、DataParallel、autocast、GradScaler全部正常工作。

这里的关键认知是：PyTorch 的cuda命名空间在 ROCm 版本里是一个兼容层。你写cuda，它认；你写hip，它反而不一定认。所以迁移的第一原则是——别自作聪明去改设备字符串。

4.2 自定义 CUDA 扩展：必须走 hipify

如果你在项目里写了.cu文件，比如自定义的算子或者 CUDA kernel，那就绕不开hipify了。流程是：

hipify-perl your_kernel.cu > your_kernel.hip

然后修改编译脚本，把nvcc换成hipcc，把 CUDA 头文件路径换成 HIP 的。PyTorch 的torch.utils.cpp_extension在 ROCm 环境下会自动识别 HIP 源文件，但前提是你的setup.py里正确声明了源文件后缀。

实测下来，简单的 element-wise kernel 转换后基本能直接跑，复杂的涉及 warp-level 原语的 kernel 可能需要手动调整。__shfl_sync这类函数在 HIP 里有对应实现，但参数语义有细微差别，需要逐个核对。

4.3 混合精度：amp能用，但要注意 dtype 支持

torch.cuda.amp在 ROCm 上可用，autocast和GradScaler都能正常工作。但要注意，某些显卡对 bf16 的支持不完整。如果你的卡是较老的架构，bf16 可能走的是软件模拟路径，速度反而比 fp16 慢。

判断方法很简单，跑一个小的 benchmark，分别测 fp16 和 bf16 的吞吐。如果 bf16 明显慢，就老老实实用 fp16。这个结论因卡而异，没有统一答案，必须实测。

4.4 多卡通信：RCCL 的坑比 NCCL 多一些

单卡场景下 ROCm 体验已经相当顺滑，多卡才是真正的分水岭。RCCL 作为 NCCL 的对应实现，功能覆盖没问题，但在拓扑识别、P2P 带宽、特定互联方式上，稳定性和性能调优空间比 NCCL 差一截。

如果你要做多卡训练，建议先用rccl-tests跑一遍带宽测试，确认卡间通信正常。常见问题是 P2P 没启用，导致通信走 PCIe 绕行，速度暴跌。可以通过环境变量NCCL_P2P_DISABLE=0强制启用，但前提是硬件和驱动支持。

注意：多卡场景下，ROCm 对主板 PCIe 拓扑比较敏感。如果卡插在不同的 PCIe root complex 下，通信性能可能不达预期。选主板时尽量让多张卡挂在同一个 CPU 的 PCIe 通道下。

5. 性能实测：省下的钱和付出的时间怎么算账

5.1 单卡训练吞吐对比

我拿一个中等规模的 Transformer 模型做了对比测试，batch size 和序列长度固定，分别在一张高端 NVIDIA 卡和一张同价位区间的 AMD 卡上跑。结论是：在标准注意力结构上，两者的吞吐差距在可接受范围内，AMD 卡的性价比优势明显。

具体数字因模型和精度而异，但趋势是一致的：计算密集型的算子，AMD 卡能发挥出硬件理论算力；访存密集型的算子，软件栈的优化程度影响更大，差距可能拉开。

这里要提醒的是，不要只看峰值算力参数。纸面 TFLOPS 和实际训练吞吐之间隔着算子库优化、显存带宽利用率、kernel launch 开销等一堆因素。选卡时看实测 benchmark，比看规格表靠谱得多。

5.2 推理场景：ROCm 的优势区间

推理是 ROCm 更容易出成绩的场景。原因在于推理的算子模式相对固定，主流模型的结构都被 MIOpen 覆盖得比较好。加上推理对多卡通信依赖低，RCCL 的短板影响小。

如果你主要做推理部署，AMD 方案的性价比会非常突出。尤其是大显存需求的场景，同样的预算能买到更大的显存，直接决定了你能不能把模型塞进去。

5.3 什么场景下不建议迁移

说了这么多好处，也得说清楚边界。以下场景我建议谨慎：

• 依赖大量冷门自定义 CUDA 算子：hipify 转换成本高，且可能遇到无对应实现的情况。
• 需要特定 NVIDIA 专有库：比如某些闭源的加速库，ROCm 没有对应替代。
• 团队完全没有 Linux 和驱动排查经验：ROCm 的排错门槛目前仍高于 CUDA，遇到问题需要一定的系统能力。
• 对多卡大规模训练有强需求：RCCL 的成熟度还需要时间。

这些不是否定 ROCm，而是帮你判断投入产出比。省下的显卡钱，可能要花一部分在环境调试上，这笔账要提前算。

6. 踩坑排查实录：那些文档里不会写的错误

6.1hipErrorNoBinaryForGpu：最常见的拦路虎

这个错误的字面意思是“没有为当前 GPU 编译的二进制”。根因通常是你的显卡架构不在 PyTorch wheel 预编译的架构列表里。解决办法有两个：一是用HSA_OVERRIDE_GFX_VERSION伪装成受支持的架构；二是从源码编译 PyTorch，把目标架构加进去。

前者快但有风险，后者慢但干净。我的建议是先用前者验证可行性，确认能跑通再考虑要不要源码编译。

6.2 显存泄漏的假象

有段时间我发现训练几个 epoch 后显存持续增长，以为是泄漏。排查后发现是 PyTorch 的缓存分配器行为，和 CUDA 下的表现一致，只是rocm-smi的显示方式让我误判了。用torch.cuda.memory_summary()看真实的分配情况，才发现缓存是正常复用的。

这个坑的教训是：监控工具的口径要统一。用 PyTorch 自己的显存统计，不要混用rocm-smi的数字做判断。

6.3 算子 fallback 导致的性能骤降

某些算子在 ROCm 上没有 GPU 实现，PyTorch 会静默 fallback 到 CPU。这种情况下代码能跑，但速度慢得离谱。排查方法是开TORCH_SHOW_CPP_STACKTRACES=1或者用 profiler 看算子执行设备。

我遇到过一次，某个归一化层走了 CPU，整个训练慢了十倍。定位到之后换成等效的 GPU 实现，问题解决。这类问题不会报错，只会让你觉得“AMD 怎么这么慢”，其实是算子覆盖的问题。

6.4 版本升级后的回归

ROCm 和 PyTorch 都在快速迭代，升级后出现回归是有可能的。我的做法是锁定版本组合，在requirements.txt或者容器镜像里写死，不轻易升级。等社区验证过新组合稳定后再跟进。

容器化是规避这类问题的好办法。ROCm 官方提供了 Docker 镜像，把驱动之外的所有依赖都打包好了，换机器时一致性有保障。

7. 把 ROCm 用顺手的几个长期习惯

第一个习惯是建一个验证脚本。每次换环境或者升级后，跑一遍这个脚本，覆盖is_available、矩阵乘法、卷积、注意力、混合精度、多卡通信这几个关键路径。任何一项失败，立刻定位，不要等到训练跑了一半才发现问题。

第二个习惯是记录版本矩阵。ROCm 版本、PyTorch 版本、显卡型号、内核版本，这四个信息每次出问题都要用到。维护一个表格，团队里共享，能省掉大量重复排查。

第三个习惯是优先用容器。裸机装 ROCm 的变量太多，容器能把大部分不确定性隔离掉。官方镜像加上自己的依赖层，构建一次到处跑。

第四个习惯是关注算子覆盖度而不是版本号。新版本不一定对你的模型更友好，关键看它有没有补上你需要的算子。看 release note 时重点看算子库的更新，而不是笼统的版本号。

我在实际使用中的体会是，ROCm 现在的状态很像几年前的 CUDA——能用，但需要你懂一点底层。这个门槛在快速降低，每过一个版本，需要手动干预的场景就少一批。对于预算敏感、又愿意花点时间折腾的开发者来说，现在入场是个不错的时间点。省下的显卡预算，足够你买好几年的云算力做对照实验了。