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 生态做一一映射立刻就清楚了。CUDA 生态ROCm 对应作用CUDA RuntimeHIP Runtime主机端 API管理设备与内存cuBLASrocBLAS基础线性代数cuDNNMIOpen深度学习算子库NCCLRCCL多卡通信NVCChipcc编译器驱动nvidia-smirocm-smi设备监控CUDA 源码HIP 源码可移植的 GPU 编程模型这张表是理解 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.1ROCm 6.0 对应 PyTorch 2.2/2.3ROCm 6.1/6.2 对应更新的版本。具体以 PyTorch 官网的安装矩阵为准不要凭记忆装。我见过太多人卡在这一步折腾一整天以为是显卡问题其实是 wheel 和 ROCm 版本错配。提示安装前先跑rocm-smi确认驱动和 ROCm 版本再对照 PyTorch 官网选择对应的安装命令。这一步花五分钟能省你五小时。3. 环境搭建实操从裸机到torch.cuda.is_available()返回 True3.1 系统选择为什么我更推荐 Ubuntu 22.04ROCm 官方支持的系统列表里Ubuntu 22.04 LTS 是体验最顺的。原因很实际内核版本、驱动包、依赖库的兼容性测试最充分。如果你用 CentOS 或者较新的非 LTS 发行版很可能遇到内核模块编译失败、amdgpu驱动加载异常这类问题。Windows 用户要注意ROCm 在 Windows 上的支持一直落后于 Linux。虽然近两年有进展但生产级训练仍然强烈建议走 Linux。如果你只有 Windows 机器双系统或者 WSL2 是可行路径但 WSL2 下的 ROCm 支持有额外限制性能也有损耗适合尝鲜不适合压榨性能。至于 macOSAMD 显卡在 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 --usecaserocm这里的version要替换成你实际下载的版本号。--usecaserocm表示只装 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_VERSION11.0.0具体填什么值取决于你的显卡架构RDNA3 通常是 11.0.0RDNA2 是 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_DISABLE0强制启用但前提是硬件和驱动支持。注意多卡场景下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_STACKTRACES1或者用 profiler 看算子执行设备。我遇到过一次某个归一化层走了 CPU整个训练慢了十倍。定位到之后换成等效的 GPU 实现问题解决。这类问题不会报错只会让你觉得“AMD 怎么这么慢”其实是算子覆盖的问题。6.4 版本升级后的回归ROCm 和 PyTorch 都在快速迭代升级后出现回归是有可能的。我的做法是锁定版本组合在requirements.txt或者容器镜像里写死不轻易升级。等社区验证过新组合稳定后再跟进。容器化是规避这类问题的好办法。ROCm 官方提供了 Docker 镜像把驱动之外的所有依赖都打包好了换机器时一致性有保障。7. 把 ROCm 用顺手的几个长期习惯第一个习惯是建一个验证脚本。每次换环境或者升级后跑一遍这个脚本覆盖is_available、矩阵乘法、卷积、注意力、混合精度、多卡通信这几个关键路径。任何一项失败立刻定位不要等到训练跑了一半才发现问题。第二个习惯是记录版本矩阵。ROCm 版本、PyTorch 版本、显卡型号、内核版本这四个信息每次出问题都要用到。维护一个表格团队里共享能省掉大量重复排查。第三个习惯是优先用容器。裸机装 ROCm 的变量太多容器能把大部分不确定性隔离掉。官方镜像加上自己的依赖层构建一次到处跑。第四个习惯是关注算子覆盖度而不是版本号。新版本不一定对你的模型更友好关键看它有没有补上你需要的算子。看 release note 时重点看算子库的更新而不是笼统的版本号。我在实际使用中的体会是ROCm 现在的状态很像几年前的 CUDA——能用但需要你懂一点底层。这个门槛在快速降低每过一个版本需要手动干预的场景就少一批。对于预算敏感、又愿意花点时间折腾的开发者来说现在入场是个不错的时间点。省下的显卡预算足够你买好几年的云算力做对照实验了。