AMD显卡跑PyTorch:ROCm安装与CUDA代码无缝迁移实战 1. 从一张显卡账单说起为什么我开始认真折腾 ROCm去年帮一个做计算机视觉的朋友配训练机预算批下来两万出头。按老思路一张能跑主流大模型的卡就得占掉大半预算剩下的钱配 CPU、内存、主板、电源处处捉襟见肘。当时我随口提了一句“要不看看 AMD 的卡”他第一反应是“PyTorch 不是只认 CUDA 吗换 AMD 不是给自己找罪受”这个反应太典型了。过去几年torch.cuda几乎成了 PyTorch 生态里 GPU 加速的代名词device cuda写进代码就像呼吸一样自然。而 AMD 的 ROCm 在很多人印象里还停留在“能跑但难装、文档稀烂、踩坑无数”的阶段。但实际情况是从 ROCm 5.x 到现在的 6.x整个安装链路和 PyTorch 适配已经发生了质变——显卡预算砍半PyTorch 代码一行不改这句话不是营销话术是我自己实测下来的结论。这篇文章想解决的问题很具体如果你手里有一张 AMD 的显卡尤其是 RDNA2/RDNA3 架构的消费级卡或者 Instinct 系列的计算卡想用它来跑 PyTorch 训练和推理到底该怎么落地我会把 ROCm 的安装逻辑、PyTorch 的适配原理、torch.cuda为什么能无缝迁移、实际踩过的坑和排查方法全部拆开讲清楚。不管你是刚搭深度学习环境的新手还是手里管着几台训练机的老手都能从里面找到能直接抄作业的步骤。先给一个最核心的认知ROCm 对 PyTorch 的适配走的是接口兼容路线而不是让你改代码。PyTorch 在编译时会针对不同后端生成对应的算子实现ROCm 版本提供的 PyTorch wheel 包里torch.cuda这个命名空间被保留了下来底层实际调用的是 HIPHeterogeneous-compute Interface for Portability。HIP 是 AMD 的一套类 CUDA 编程接口语法和 CUDA 高度相似PyTorch 官方在构建 ROCm 版本时把 CUDA 相关的调用映射到了 HIP 上。所以你代码里写的torch.cuda.is_available()、.to(cuda)、torch.cuda.amp在 ROCm 环境下依然有效只是背后跑的是 AMD 的运行时。这就是“一行不用改”的技术底座。理解了这一点后面所有的安装和配置本质上都是在做一件事让 PyTorch 找到 ROCm 运行时并让 ROCm 正确识别你的显卡。2. 装之前先搞明白ROCm 到底支持哪些卡很多人卡在第一步不是安装命令敲错了而是手里的卡根本不在 ROCm 的官方支持列表里。ROCm 的硬件支持分几个层级我按实际可用性给你梳理一下。2.1 官方明确支持的显卡型号AMD 官方文档里ROCm 支持的 GPU 主要分两类。一类是Instinct 系列计算卡比如 MI100、MI200 系列MI210、MI250、MI250X、MI300 系列这些是数据中心级的支持最完整FP64、FP32、FP16、BF16 各种精度都齐全适合大规模训练。另一类是Radeon 消费级显卡从 RDNA2 架构开始逐步纳入支持比如 RX 6900 XT、RX 6800 XT、RX 6800、RX 6700 XT 等RDNA3 架构的 RX 7900 XTX、RX 7900 XT、RX 7800 XT、RX 7700 XT 也在支持范围内。RDNA4 的新卡RX 9000 系列在 ROCm 6.3 之后也开始陆续进入支持列表。这里有个关键点消费级卡的支持是“能用”和“好用”之间的区别。官方支持列表里的卡装完 ROCm 和 PyTorch 之后基本能直接跑起来不在列表里的卡可能需要手动指定HSA_OVERRIDE_GFX_VERSION环境变量来“伪装”成相近架构的卡能跑但稳定性看运气。2.2 不在列表里怎么办gfx 版本覆盖的实操我手里有一张 RX 6600 XTgfx 版本是gfx1032而 ROCm 官方支持的 RDNA2 消费卡主要是gfx1030RX 6800/6900 系列。直接装完 ROCm 后rocminfo能看到卡但 PyTorch 里torch.cuda.is_available()返回False。解决办法是设置环境变量export HSA_OVERRIDE_GFX_VERSION10.3.0这个值的含义是让 ROCm 把gfx1032当作gfx1030来处理。10.3.0 对应 gfx103010.3.1 对应 gfx1031以此类推。设置完之后再跑 PyTorchis_available()就变成True了。注意HSA_OVERRIDE_GFX_VERSION不是万能药。它本质上是让运行时用相近架构的指令集来执行大部分算子没问题但涉及特定硬件特性的算子可能会出错或者性能打折。如果你的卡完全不在任何相近架构的覆盖范围内强行覆盖可能直接崩溃。2.3 核显和 APU 的情况Ryzen 7000/8000 系列桌面 APU 里的核显比如 Radeon 780Mgfx1103以及移动端的核显ROCm 的支持比较有限。部分型号可以通过HSA_OVERRIDE_GFX_VERSION11.0.0来尝试但内存带宽和计算单元数量摆在那里跑小模型推理还行训练基本不用考虑。如果你只是想在本地跑个量化后的小模型做推理核显可以凑合真要训练还是得独显。3. 安装 ROCm从系统准备到验证的完整链路安装 ROCm 这件事官方文档给的是“标准流程”但实际执行时会遇到各种系统层面的问题。我按 Ubuntu 22.04 和 24.04 这两个最常用的版本把完整链路拆一遍。3.1 系统准备内核版本和依赖检查ROCm 对内核版本有要求。Ubuntu 22.04 默认的 5.15 内核可以用但建议升级到 6.2 以上对 AMD 显卡的驱动支持更好。Ubuntu 24.04 默认的 6.8 内核没问题。检查当前内核uname -r然后确认系统里没有残留的旧版 ROCm 或者冲突的驱动sudo apt list --installed | grep -i rocm sudo apt list --installed | grep -i amdgpu如果有旧版本先清理干净。这一步很多人会忽略结果新版本装上去和旧版本的文件混在一起出现各种奇怪的报错。3.2 添加软件源和安装核心包AMD 官方提供了 apt 源。以 ROCm 6.2 为例写这篇文章时 6.3 已经发布但 6.2 的生态兼容性更稳添加源的命令sudo apt update sudo apt install -y wget gnupg2 wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg /dev/null echo deb [archamd64 signed-by/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.2 jammy main | sudo tee /etc/apt/sources.list.d/rocm.list注意把jammy换成你的系统代号Ubuntu 24.04 是noble。然后更新并安装sudo apt update sudo apt install -y rocm-hip-sdk rocm-developer-toolsrocm-hip-sdk包含了 HIP 运行时、编译器、库文件rocm-developer-tools包含了rocminfo、rocm-smi这些诊断工具。安装包比较大几个 G下载时间取决于网络。3.3 用户组和权限配置装完之后当前用户需要加入render和video组否则访问 GPU 设备会权限不足sudo usermod -a -G render,video $USER执行完这条命令后必须重新登录或者重启组权限才会生效。我见过有人装完直接跑rocminfo报权限错误折腾半天才发现是没重新登录。3.4 验证 ROCm 是否正常工作重新登录后先跑rocminfo | grep -i gfx如果能看到你的显卡对应的 gfx 版本号说明 ROCm 运行时已经识别到卡了。再跑rocm-smi这个命令会显示显卡的温度、功耗、显存占用等信息。如果两个命令都正常输出ROCm 这一层就通了。实操心得rocminfo输出里如果出现HSA_STATUS_ERROR_OUT_OF_RESOURCES或者类似的错误大概率是内核模块没加载好。可以试着重启或者检查dmesg | grep -i amdgpu看驱动加载日志。4. PyTorch 的 ROCm 版本安装方式和 CUDA 版的区别PyTorch 官方从 1.8 版本开始提供 ROCm 版本的预编译包到现在 2.x 系列ROCm 版本的发布节奏基本跟得上主线。安装方式和 CUDA 版最大的区别在于源地址和包名后缀。4.1 用 pip 安装 ROCm 版 PyTorchCUDA 版 PyTorch 的安装命令通常是pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121ROCm 版把cu121换成rocm6.2pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm6.2这里的rocm6.2对应你安装的 ROCm 版本。如果你装的是 ROCm 6.3就换成rocm6.3。版本号必须匹配否则 PyTorch 找不到对应的运行时库。4.2 验证 PyTorch 是否正确识别 GPU装完之后跑一段最简单的验证代码import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.device_count()) print(torch.cuda.get_device_name(0))如果is_available()返回Trueget_device_name(0)显示的是你的 AMD 显卡型号那就说明整条链路通了。注意这里打印出来的设备名可能是AMD Radeon RX 7900 XTX之类的而不是 NVIDIA 的型号这是正常的。4.3 为什么代码不用改torch.cuda 的兼容层逻辑前面提到过ROCm 版 PyTorch 保留了torch.cuda命名空间。具体来说PyTorch 在编译时如果检测到构建目标是 ROCm会把原本调用 CUDA Runtime API 的地方替换成 HIP API 的调用。HIP API 的函数名和 CUDA 几乎一一对应比如cudaMalloc对应hipMalloccudaMemcpy对应hipMemcpy。PyTorch 的代码里通过宏定义和条件编译在 ROCm 构建时走 HIP 路径。对于用户代码来说你写的tensor.to(cuda)、torch.cuda.synchronize()、torch.cuda.amp.autocast()在 ROCm 环境下都会被正确路由到 HIP 后端。这就是“一行不用改”的实现原理。注意虽然接口兼容但并不是所有 CUDA 特有的功能在 ROCm 上都有对应实现。比如某些依赖 cuDNN 特定版本的功能、NCCL 多卡通信的某些配置在 ROCm 上对应的是 MIOpen 和 RCCL行为可能有细微差异。大部分常见模型和训练流程没问题但涉及冷门算子或者深度定制 CUDA 扩展的场景需要额外验证。5. 实际跑起来从单卡训练到多卡并行的配置装好环境只是第一步真正跑训练的时候还有一些配置项需要调整。5.1 单卡训练的环境变量调优ROCm 环境下有几个环境变量对性能和稳定性影响很大export HSA_OVERRIDE_GFX_VERSION11.0.0 # 仅非官方支持卡需要 export PYTORCH_HIP_ALLOC_CONFexpandable_segments:True export GPU_MAX_HW_QUEUES4PYTORCH_HIP_ALLOC_CONF对应 CUDA 版的PYTORCH_CUDA_ALLOC_CONF用来控制显存分配策略。expandable_segments:True可以减少显存碎片对长序列训练比较有用。GPU_MAX_HW_QUEUES控制硬件队列数量默认值在某些负载下可能偏低适当调高能改善吞吐。5.2 混合精度训练在 ROCm 上的表现PyTorch 的 AMP自动混合精度在 ROCm 上通过torch.cuda.amp依然可用底层对应的是 HIP 的 FP16 支持。RDNA3 和 CDNA 架构的卡对 FP16 和 BF16 都有硬件加速实测下来 AMP 的加速比在 1.5 到 2 倍之间和同级别 NVIDIA 卡的表现接近。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for data, target in dataloader: optimizer.zero_grad() with autocast(): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这段代码在 CUDA 和 ROCm 环境下完全一致不需要任何修改。5.3 多卡训练RCCL 替代 NCCL多卡训练时PyTorch 的分布式模块默认使用 NCCL 做通信后端。在 ROCm 环境下对应的库是 RCCLROCm Communication Collectives Library。PyTorch 的 ROCm 版本会自动把nccl后端映射到 RCCL代码里写的dist.init_process_group(backendnccl)依然有效。不过多卡配置有几个坑要注意。首先是卡之间的拓扑PCIe 带宽和 NVLinkAMD 对应的是 Infinity Fabric的差异会影响通信效率。其次是HIP_VISIBLE_DEVICES环境变量对应 CUDA 的CUDA_VISIBLE_DEVICES用来指定哪些卡对进程可见export HIP_VISIBLE_DEVICES0,1,2,3如果你在代码里用torch.cuda.set_device()指定卡号在 ROCm 环境下也是生效的但建议优先用环境变量来控制避免代码里的硬编码。6. 踩坑实录那些文档里不会写的问题和排查方法这部分是我自己折腾过程中遇到的实际问题按出现频率排序。6.1 常见问题速查表问题现象可能原因排查方法解决方式torch.cuda.is_available()返回 FalseROCm 未正确安装或显卡不支持跑rocminfo看是否识别到卡检查 ROCm 版本和显卡 gfx 版本必要时设HSA_OVERRIDE_GFX_VERSIONrocminfo报权限错误用户不在 render/video 组groups命令查看当前组sudo usermod -a -G render,video $USER后重新登录训练中途显存溢出显存分配策略问题监控rocm-smi显存占用设PYTORCH_HIP_ALLOC_CONFexpandable_segments:True减小 batch size多卡训练卡死RCCL 通信问题检查HIP_VISIBLE_DEVICES设置确认卡间拓扑尝试换 PCIe 插槽或降低通信频率某些算子报 HIP 错误算子未在 ROCm 上实现定位到具体算子用 CPU 实现替代或找社区补丁安装后系统不稳定内核模块冲突dmesggrep -i amdgpu6.2 一个典型的显存溢出排查过程有一次跑一个 ViT 的微调任务batch size 设到 16 就 OOM但同样的模型在 24G 显存的 NVIDIA 卡上能跑到 32。一开始以为是 AMD 卡显存管理效率低后来用rocm-smi监控发现显存占用在训练过程中波动很大峰值远超稳态值。排查下来是 PyTorch 的显存缓存策略问题。ROCm 版的 PyTorch 默认的显存分配器和 CUDA 版行为有差异缓存块释放不够及时。设置PYTORCH_HIP_ALLOC_CONFexpandable_segments:True之后峰值显存降下来了batch size 能跑到 24。这个经验说明ROCm 环境下的显存调优不能照搬 CUDA 的经验需要根据实际监控数据来调整。6.3 关于“显卡被物理移除”的报错有些人在 Windows 和 Linux 双系统切换后会遇到GPU has been physically removed之类的报错。这通常是驱动状态不一致导致的。Linux 下 AMD 显卡用的是内核自带的amdgpu驱动Windows 下用的是 AMD 官方驱动两边切换时如果没正常关机比如 Windows 快速启动没关显卡的固件状态可能残留导致 Linux 侧识别异常。解决办法是在 Windows 里关闭快速启动或者每次切换系统时完整关机再开机。如果已经出现这个问题可以尝试sudo modprobe -r amdgpu sudo modprobe amdgpu重新加载内核模块或者直接重启。7. 性能实测同价位 AMD 卡和 NVIDIA 卡的对比光说“能用”不够还得看“好不好用”。我拿手头的 RX 7900 XTX24G 显存和一张同价位的 NVIDIA 卡做了几组对比测试任务包括 ResNet-50 训练、BERT-base 微调、Stable Diffusion 推理。7.1 训练吞吐对比在 ResNet-50 的 ImageNet 子集上batch size 设为 64FP16 混合精度指标RX 7900 XTX同价位 NVIDIA 卡单步耗时约 0.42s约 0.38s显存占用18.2G16.8G功耗约 320W约 280W差距在 10% 左右考虑到价格差异同显存容量下 AMD 卡通常便宜 30% 到 50%性价比优势明显。7.2 推理场景的表现Stable Diffusion 1.5 生成 512x512 图片20 步采样指标RX 7900 XTX同价位 NVIDIA 卡单张耗时约 2.8s约 2.3s显存峰值6.4G5.9G推理场景下差距略大一些主要是因为 MIOpenAMD 的深度学习库在某些卷积算子的优化上还不如 cuDNN 成熟。但日常使用完全够用而且省下来的预算可以多买一张卡做并行。7.3 什么场景下不建议用 ROCm说了这么多优点也得说清楚边界。以下几种情况我建议还是优先考虑 NVIDIA需要用到特定的 CUDA 扩展库比如某些自定义算子、特定版本的 cuDNN 特性团队协作环境里其他人都在用 CUDA统一环境成本更低需要用到 TensorRT 做极致推理优化依赖某些只在 CUDA 上验证过的冷门模型或框架ROCm 的生态在快速完善但客观说CUDA 积累了十几年的工具链和社区资源短期内不可能完全追平。选择 ROCm 的核心逻辑是性价比和特定场景下的可用性而不是全面替代。8. 把 ROCm 纳入日常开发流的一些经验最后分享几个我在实际使用中总结的习惯能让 ROCm 环境用起来更顺手。第一环境隔离一定要做。ROCm 版本和 PyTorch 版本的对应关系比较敏感建议用 conda 或者 venv 给每个项目建独立环境把 ROCm 相关的环境变量写进激活脚本里。这样切换项目时不会互相干扰。第二监控工具常开。rocm-smi可以加--showpower --showtemp --showmemuse参数持续监控训练时开一个终端窗口盯着能提前发现显存泄漏或者温度异常。第三版本升级要谨慎。ROCm 和 PyTorch 的版本组合有兼容性矩阵升级其中一个之前先查清楚对应的版本要求。我吃过一次亏把 ROCm 从 6.0 升到 6.2结果 PyTorch 还是 6.0 版本的 wheel跑起来各种算子报错回退才恢复正常。第四社区资源要会用。ROCm 的 GitHub issue 区和 PyTorch 的 ROCm 相关讨论区里很多问题已经有现成的解决方案。遇到报错先搜关键词大概率能找到同样踩过坑的人。这套东西折腾下来最大的感受是ROCm 已经过了“能不能用”的阶段进入了“怎么用更好”的阶段。对于预算敏感的个人开发者和中小团队来说它提供了一个实实在在的选项。代码不用改这件事省下的不只是迁移成本更是心理负担——你不需要为了用 AMD 的卡去重新学一套编程模型PyTorch 的那套东西原封不动就能跑。