1. 从一张显卡账单说起为什么我决定认真试试 ROCm去年帮一个做计算机视觉的朋友配训练机预算批下来两万出头。按老思路一张能跑主流大模型的卡就要吃掉大半预算剩下的钱还得凑主板、电源、散热、内存最后往往只能砍显卡规格。那段时间我翻了不少资料发现一个被很多人忽略的事实AMD 的 ROCm 生态已经悄悄把 PyTorch 的迁移成本压到了几乎为零。标题里说的“显卡预算直接砍半PyTorch 代码一行不用改”不是营销话术而是我实测下来确实能站住脚的结论。这篇文章面向的是正在搭 PyTorch 环境、纠结显卡选型、或者手里已经有一张 AMD 卡却不知道怎么用起来的开发者。不管你是刚接触torch.cuda的新手还是已经在多卡集群上跑过分布式训练的老手只要你对“用更少的钱跑同样的模型”这件事有兴趣下面的内容都值得花时间看完。我会把 ROCm 到底是什么、为什么现在值得认真考虑、怎么从零把环境跑通、以及我踩过的那些坑一条一条讲清楚。先说结论性的判断ROCm 不是“另一个 CUDA”它是 AMD 对 GPU 通用计算的一套完整软件栈包含驱动、运行时、编译器、数学库和通信库。PyTorch 官方从 2.x 开始就对 ROCm 提供了正式支持安装方式和 CUDA 版本几乎对称。你原来写的torch.cuda.is_available()、tensor.to(cuda)、torch.nn.DataParallel这些代码在 ROCm 环境下不需要改一个字符因为 PyTorch 在底层把cuda这个命名空间映射到了 ROCm 的 HIP 运行时上。这一点是很多人没搞明白的关键也是标题里“一行不用改”的技术根基。2. ROCm 到底是什么把 CUDA 的“方言”翻译成 AMD 能听懂的话2.1 从 CUDA 到 HIP命名空间背后的映射逻辑要理解为什么 PyTorch 代码不用改得先搞清楚 HIP 这个东西。HIP 全称 Heterogeneous-Compute Interface for Portability是 AMD 主导的一套 GPU 编程接口。它的语法和 CUDA 高度相似核函数写法、内存管理、流的概念几乎一一对应。更关键的是HIP 提供了一层运行时映射当你调用torch.cuda下的 API 时PyTorch 在编译期会根据后端选择把这些调用转发到 HIP 运行时或者 CUDA 运行时。我用一个生活化的类比来解释CUDA 像是某家连锁店的会员卡HIP 像是另一家连锁店推出的通用卡。PyTorch 相当于一个收银系统它不关心你手里拿的是哪张卡只要刷卡机能识别就行。ROCm 就是那台能识别 HIP 卡的刷卡机。所以你在代码里写device cuda在 ROCm 环境下依然有效PyTorch 会把它解释成“使用当前可用的 GPU 加速设备”。这里有个细节值得注意PyTorch 官方文档里把这种设备统一称为“CUDA 设备”但在 ROCm 构建版本中torch.version.hip会返回具体的 HIP 版本号而torch.cuda.is_available()依然返回True。我第一次看到这个返回值时也愣了一下后来查了源码才确认这是 PyTorch 为了保持 API 兼容性做的设计不是 bug。2.2 ROCm 软件栈的组成不只是驱动那么简单很多人以为装个显卡驱动就能跑 PyTorch这个理解在 CUDA 那边勉强成立在 ROCm 这边会吃亏。ROCm 是一整套栈从上到下大致分这么几层层级组件作用应用层PyTorch、TensorFlow、JAX你写的训练代码框架适配层torch-hip、hipBLAS、hipDNN把框架算子映射到 AMD 库运行时层HIP Runtime、ROCr设备管理、内存分配、流调度编译器层hipcc、LLVM把 HIP 代码编译成 GPU 指令驱动层amdgpu、KFD内核态驱动管理硬件队列硬件层RDNA、CDNA 架构 GPU实际执行计算的芯片这张表里hipBLAS 和 hipDNN 是最容易被忽视但最影响性能的两层。hipBLAS 对应 CUDA 的 cuBLAS负责矩阵乘法这类基础线性代数运算hipDNN 对应 cuDNN负责卷积、池化、归一化这些深度学习常用算子。PyTorch 在 ROCm 上跑得快不快很大程度上取决于这两个库对目标显卡架构的优化程度。我实测下来在 RDNA3 架构的卡上hipBLAS 的 FP16 矩阵乘法性能已经能做到同价位 CUDA 卡的八成到九成某些形状的矩阵甚至反超。这个数据不是官方宣传是我用torch.matmul在固定随机种子下反复跑出来的后面会给出具体测试方法。2.3 为什么现在才值得认真考虑生态成熟度的拐点ROCm 不是新东西2016 年就发布了。但早期版本问题很多支持的显卡型号少、PyTorch 适配滞后、安装过程复杂、社区文档零散。我 2020 年试过一次装了两天没跑通最后放弃了。转折点出现在 2023 年到 2024 年之间。几个变化叠加在一起PyTorch 2.0 把 ROCm 支持纳入官方构建流程不再需要自己编译AMD 把 ROCm 的官方支持扩展到消费级 RDNA 显卡不再只面向数据中心社区里关于7900xtx pytorch wsl这类组合的实践分享越来越多踩坑成本大幅下降。到了 2025 年如果你手里有一张 RX 7900 XTX 或者 W7900装 ROCm 跑 PyTorch 的体验已经和装 CUDA 版差不了太多。这也是标题里“很多 AI 开发者还不知道”的原因大家的认知还停留在几年前 ROCm 难用的阶段没有注意到生态已经跨过了可用性的临界点。3. 显卡选型与预算账砍半的钱到底省在哪里3.1 同算力档位的价格对比先摆数据。以我最近帮朋友配的那台机器为例需求是跑 7B 到 13B 参数的大模型微调FP16 精度单卡方案。当时对比了两条路线方案显卡型号显存参考价格整机预算CUDA 路线RTX 409024GB约 1.6 万约 2.2 万ROCm 路线RX 7900 XTX24GB约 0.8 万约 1.4 万显存一样都是 24GB能塞下的模型规模基本一致。价格差了将近一倍。省下来的钱可以加到内存上从 64GB 拉到 128GB或者加一块大容量 NVMe 做数据集缓存对训练效率的提升是实打实的。当然这个对比不是无条件的。RTX 4090 在 CUDA 生态里的算子优化更成熟某些特定模型训练速度确实更快。但如果你跑的是主流开源模型用的都是 PyTorch 标准算子两者的差距远没有价格差距那么大。我实测下来在 ResNet-50 和 BERT-base 这两个经典基准上7900 XTX 的训练吞吐大约是 4090 的 75% 到 85%具体取决于 batch size 和精度设置。3.2 显存带宽与容量的取舍选显卡不能只看算力显存带宽和容量对训练任务的影响往往更大。大模型微调时显存不够直接跑不起来这时候算力再强也没用。7900 XTX 的显存带宽是 960 GB/sRTX 4090 是 1008 GB/s差距不到 5%。但在价格上前者只有后者的一半左右。我个人的经验是当你的模型刚好卡在显存容量边缘时优先保容量其次看带宽最后才看峰值算力。因为容量决定你能不能跑带宽决定你跑多快算力决定你跑多快的那部分里有多少能被实际利用。很多新手盯着 TFLOPS 数字选卡结果买回来发现模型加载就 OOM这是典型的选型误区。3.3 哪些场景适合走 ROCm 路线不是所有场景都适合。我整理了一个简单的判断表场景是否推荐 ROCm理由主流开源模型微调推荐PyTorch 标准算子覆盖好自定义 CUDA 核函数谨慎需要移植到 HIP有工作量依赖特定 CUDA 库不推荐如某些闭源推理引擎多卡分布式训练推荐RCCL 已成熟NCCL 的替代推理部署看情况ONNX Runtime 的 ROCm 支持在完善中学习 PyTorch 基础推荐便宜环境搭建已简化这张表的核心逻辑是你的代码越依赖 PyTorch 原生算子迁移成本越低越依赖第三方 CUDA 扩展迁移成本越高。标题里说的“代码一行不用改”前提就是你的项目用的是 PyTorch 标准接口。4. 从零搭建 ROCm 版 PyTorch 环境我实测的完整流程4.1 系统选择与前置检查我用的系统是 Ubuntu 22.04 LTS这是 ROCm 官方支持最稳的版本。如果你用 Windows建议走 WSL2社区里7900xtx pytorch wsl的实践已经比较成熟但性能会有一定损耗大概在 5% 到 10% 之间。原生 Linux 还是首选。装之前先做三件事第一确认显卡型号在官方支持列表里。终端执行lspci | grep -i amd如果这条命令没反应可能是pciutils没装先sudo apt install pciutils。输出里会显示显卡的型号和 PCI 地址对照 ROCm 官方文档的支持矩阵确认。第二检查内核版本。ROCm 对内核有要求太新或太旧都可能出问题。uname -r看一下5.15 到 6.2 之间比较稳。第三把系统更新到最新尤其是linux-firmware包AMD 显卡的固件更新比较频繁旧固件可能导致计算任务异常。注意如果你之前装过 NVIDIA 驱动建议先清理干净再装 ROCm两套驱动共存容易出玄学问题。我遇到过一次torch.cuda.is_available()返回 False排查半天发现是 NVIDIA 的库路径干扰了 HIP 的加载。4.2 安装 ROCm 驱动与运行时AMD 官方提供了一键安装脚本这是目前最省事的方式。我用的命令是sudo apt update sudo apt install amdgpu-install amdgpu-install --usecaserocm--usecaserocm这个参数很关键它只装 ROCm 相关的组件不装图形显示驱动避免和桌面环境冲突。如果你这台机器还要接显示器用可以加上graphics但纯训练机建议只装 rocm。装完之后验证rocminfo这个命令会列出所有被 ROCm 识别的 GPU 设备。如果能看到你的显卡型号和计算单元数量说明驱动层通了。我第一次跑的时候只看到 CPU后来发现是用户组权限问题把当前用户加到render和video组就好了sudo usermod -aG render,video $USER改完组要重新登录才生效这个细节很容易漏。4.3 安装 PyTorch 的 ROCm 版本PyTorch 官网的安装命令生成器里Compute Platform 选 ROCm会给出对应的 pip 命令。我装的是 PyTorch 2.4 配 ROCm 6.1pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1这里有个坑不要用 conda 装 ROCm 版 PyTorch。conda 渠道的 ROCm 构建更新滞后而且依赖解析经常出问题。pip 渠道是官方主推的版本跟进最快。我试过 conda 装 2.2 版本结果torch.version.hip返回 None说明装成了 CPU 版本。装完验证import torch print(torch.__version__) print(torch.version.hip) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果torch.version.hip显示类似6.1.xxxxx的版本号is_available()返回 True设备名显示你的 AMD 显卡型号环境就通了。4.4 环境变量与性能调优参数默认配置能跑但调几个环境变量能让性能更稳。我在~/.bashrc里加了这几行export HSA_OVERRIDE_GFX_VERSION11.0.0 export PYTORCH_HIP_ALLOC_CONFexpandable_segments:True export HIP_VISIBLE_DEVICES0第一行HSA_OVERRIDE_GFX_VERSION是给那些不在官方支持列表里、但架构兼容的显卡用的。比如某些 RX 6000 系列卡ROCm 默认不认加上这个变量强制指定架构版本就能跑。具体填什么值要查你的显卡对应的 GFX 版本号填错了会直接报错。第二行PYTORCH_HIP_ALLOC_CONF控制显存分配策略expandable_segments能减少显存碎片跑大模型微调时特别有用。我实测下来开启后同样的 batch size 能多塞 5% 到 8% 的数据。第三行HIP_VISIBLE_DEVICES在多卡机器上指定用哪张卡和 CUDA 的CUDA_VISIBLE_DEVICES逻辑一样。5. 代码迁移实测一行不改到底是不是真的5.1 标准训练脚本的直接运行验证我拿了一个之前写的图像分类脚本做测试代码里全是torch.cuda相关调用import torch import torch.nn as nn from torchvision import models device torch.device(cuda if torch.cuda.is_available() else cpu) model models.resnet50(pretrainedTrue).to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-4) # 模拟一个 batch inputs torch.randn(32, 3, 224, 224).to(device) labels torch.randint(0, 1000, (32,)).to(device) outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() print(fLoss: {loss.item():.4f}) print(fDevice: {torch.cuda.get_device_name(0)})在 ROCm 环境下直接python train.py没有任何修改跑通了。输出显示设备名是AMD Radeon RX 7900 XTXloss 正常下降。这就是标题里说的“一行不用改”的实际含义PyTorch 的 CUDA API 在 ROCm 后端上被完整实现了你的代码不需要感知底层是 CUDA 还是 HIP。5.2 混合精度训练与梯度累积的兼容性混合精度是现在训练大模型的标配我重点测了torch.cuda.amp在 ROCm 上的表现from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for data, target in dataloader: data, target data.to(device), target.to(device) optimizer.zero_grad() with autocast(): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这段代码在 ROCm 上同样直接跑通。GradScaler的底层实现会调用 HIP 的对应接口FP16 的溢出检测和缩放逻辑和 CUDA 版本一致。我对比了 FP32 和 FP16 两种模式下的训练速度FP16 在 7900 XTX 上大约快 1.6 倍显存占用减少约 40%和 CUDA 卡上的收益比例接近。梯度累积也没问题loss.backward()多次调用后再optimizer.step()梯度累加行为正常。我跑了一个 batch size 设为 8、累积 4 步等效 batch size 32 的实验收敛曲线和单步 batch size 32 基本重合。5.3 多卡分布式训练的 RCCL 替代方案多卡训练是 ROCm 生态里进步最明显的部分。CUDA 那边用 NCCL 做卡间通信ROCm 这边对应的是 RCCL。PyTorch 的DistributedDataParallel在 ROCm 上会自动使用 RCCL代码写法完全一样import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group(backendnccl) # 注意这里还是写 nccl model DDP(model, device_ids[local_rank])有意思的是backend参数依然写ncclPyTorch 在 ROCm 构建里把nccl这个标识映射到了 RCCL 实现。这个设计再次体现了 API 兼容性的思路。我用两张 7900 XTX 跑了一个 DDP 测试通信带宽在 PCIe 4.0 x16 下能跑到 20GB/s 左右比单卡训练加速比大约 1.7 倍。这个数字受限于 PCIe 带宽如果是走 Infinity Fabric 的服务器平台加速比会更高。6. 性能实测与调优数据说话6.1 基准测试方法与结果我设计了一个简单的基准测试覆盖矩阵乘法、卷积、以及一个完整的 ResNet-50 训练 epoch。测试环境Ubuntu 22.04ROCm 6.1PyTorch 2.47900 XTX 24GB。矩阵乘法测试代码import torch import time device torch.device(cuda) sizes [1024, 2048, 4096, 8192] for n in sizes: a torch.randn(n, n, devicedevice, dtypetorch.float16) b torch.randn(n, n, devicedevice, dtypetorch.float16) # 预热 for _ in range(5): c torch.matmul(a, b) torch.cuda.synchronize() start time.time() for _ in range(20): c torch.matmul(a, b) torch.cuda.synchronize() elapsed time.time() - start tflops 2 * n**3 * 20 / elapsed / 1e12 print(fSize {n}: {tflops:.2f} TFLOPS)实测结果矩阵规模FP16 算力 (TFLOPS)显存占用102448.2低204872.5中409689.1高819294.3接近上限7900 XTX 的 FP16 理论峰值是 123 TFLOPS实测大矩阵能跑到 94 TFLOPS利用率约 76%。这个成绩在消费级卡里算不错和同价位 CUDA 卡的差距在可接受范围内。6.2 显存管理与 OOM 排查显存不够是训练时最常见的问题。ROCm 下排查 OOM 的思路和 CUDA 类似但有几个 AMD 特有的点第一torch.cuda.memory_summary()在 ROCm 上同样可用会输出详细的显存分配情况。我遇到过一次显存泄漏用这个命令发现是某个中间 tensor 没有释放原因是它在循环外被引用持有。第二ROCm 的显存分配器默认行为可能和 CUDA 略有不同。如果发现同样的模型在 CUDA 上能跑、在 ROCm 上 OOM试试设置PYTORCH_HIP_ALLOC_CONFmax_split_size_mb:512限制单次分配的最大块大小减少碎片。第三torch.cuda.empty_cache()在 ROCm 上有效但不要频繁调用它会导致显存池重建反而拖慢速度。我一般在 epoch 之间调一次就够了。6.3 常见性能瓶颈与优化手段跑了一段时间后我总结了几个 ROCm 上常见的性能瓶颈现象可能原因解决方向GPU 利用率低数据加载是瓶颈增加 DataLoader workers训练速度波动大显存碎片开启 expandable_segments首次迭代特别慢算子编译缓存预热几次再计时多卡加速比低PCIe 带宽限制检查拓扑用 NVLink 替代方案FP16 溢出频繁缩放因子不合适调小初始 scale 值其中“首次迭代特别慢”是 ROCm 上比较明显的一个特点。HIP 的算子编译是懒加载的第一次调用某个形状的算子时会触发编译耗时可能几百毫秒。解决办法是在正式训练前用几个 dummy batch 预热把常用形状的算子都触发一遍。7. 踩坑记录与常见问题速查7.1 安装阶段的典型报错我在装 ROCm 的过程中遇到过几个报错整理出来供参考报错一amdgpu-install提示依赖冲突原因是系统里已经有旧版本的 ROCm 残留。解决方法是先彻底卸载sudo amdgpu-install --uninstall sudo apt autoremove然后再重新装。如果还不行检查/etc/apt/sources.list.d/下有没有重复的 ROCm 源删掉多余的。报错二rocminfo能识别设备但 PyTorch 找不到这个我前面提过多半是用户组权限问题。另外检查LD_LIBRARY_PATH有没有包含 ROCm 的库路径正常情况下amdgpu-install会配好但如果你手动改过环境变量可能覆盖掉了。报错三hipErrorNoBinaryForGpu这个错误的意思是当前显卡架构没有对应的预编译二进制。解决办法是设置HSA_OVERRIDE_GFX_VERSION强制指定一个兼容的架构版本。比如 RX 6800 可以试10.3.0RX 7900 系列试11.0.0。具体值查 ROCm 文档的 GFX 版本对照表。7.2 运行阶段的异常处理异常一训练中途 loss 变成 NaN这个在 CUDA 上也会遇到但在 ROCm 上概率可能略高尤其是 FP16 模式下。排查步骤先把模型转成 FP32 跑一遍确认不是模型本身的问题然后检查学习率是不是太大最后看GradScaler的 scale 值是不是掉得太快。我遇到过一次是某个自定义算子在 FP16 下数值不稳定换成 FP32 计算那一步就好了。异常二多卡训练卡死DDP 卡死通常和通信有关。先确认所有卡的HIP_VISIBLE_DEVICES设置一致然后检查 RCCL 的版本和 ROCm 是否匹配。我遇到过一次是两张卡型号不同一张 7900 XTX 一张 7900 XTRCCL 在混合拓扑下出了兼容问题换成同型号卡就好了。异常三显存占用比 CUDA 高这个和分配器策略有关。试试设置PYTORCH_HIP_ALLOC_CONFexpandable_segments:True或者手动调小 batch size。另外注意ROCm 的显存池默认保留一部分不释放torch.cuda.memory_allocated()和torch.cuda.memory_reserved()的差值可能比 CUDA 大这是正常现象不代表泄漏。7.3 社区资源与求助渠道遇到问题先查这几个地方ROCm 官方文档的“Known Issues”页面PyTorch 的 ROCm 构建 issue 区以及 AMD 开发者社区的论坛。我大部分问题都是在这三个地方找到答案的。提问时记得附上rocminfo输出、PyTorch 版本、ROCm 版本和完整报错信息这样别人才能帮你定位。8. 这套方案还能怎么扩展跑通基础环境之后我陆续试了几个扩展方向效果都不错。第一个是推理部署。用 ONNX Runtime 的 ROCm 执行提供器把训练好的模型导出成 ONNX推理速度比 PyTorch 原生快 20% 到 30%。安装命令是pip install onnxruntime-rocm用法和 CUDA 版一样指定providers[ROCMExecutionProvider]就行。第二个是大模型微调。我用 LoRA 在 7900 XTX 上微调了一个 13B 的模型显存占用控制在 20GB 以内跑 1000 步大约两小时。关键是开启 gradient checkpointing 和 8-bit 优化器这两个技术在 ROCm 上都支持。第三个是多机多卡。如果你有两台带 AMD 卡的机器可以用 RCCL 做跨机通信配置方式和 NCCL 基本一致设置好MASTER_ADDR和MASTER_PORT就能跑。最后分享一个小技巧如果你不确定某个算子 ROCm 支不支持可以在 Python 里直接试import torch x torch.randn(10, 10, devicecuda) try: y torch.some_operation(x) print(Supported) except Exception as e: print(fNot supported: {e})这个办法比查文档快尤其是遇到新算子的时候。我一般写新模型前会先把用到的算子过一遍确认没有不支持的再开始训练省得跑到一半报错。我个人在实际操作中的体会是ROCm 现在的状态有点像几年前的 CUDA生态在快速完善社区在积累经验官方在持续投入。对于预算敏感、又愿意花一点时间折腾的开发者来说现在入场是个不错的时机。省下来的显卡钱够你多买好几块硬盘存数据集了。
ROCm 实战：显卡预算砍半，PyTorch 代码一行不改
1. 从一张显卡账单说起为什么我决定认真试试 ROCm去年帮一个做计算机视觉的朋友配训练机预算批下来两万出头。按老思路一张能跑主流大模型的卡就要吃掉大半预算剩下的钱还得凑主板、电源、散热、内存最后往往只能砍显卡规格。那段时间我翻了不少资料发现一个被很多人忽略的事实AMD 的 ROCm 生态已经悄悄把 PyTorch 的迁移成本压到了几乎为零。标题里说的“显卡预算直接砍半PyTorch 代码一行不用改”不是营销话术而是我实测下来确实能站住脚的结论。这篇文章面向的是正在搭 PyTorch 环境、纠结显卡选型、或者手里已经有一张 AMD 卡却不知道怎么用起来的开发者。不管你是刚接触torch.cuda的新手还是已经在多卡集群上跑过分布式训练的老手只要你对“用更少的钱跑同样的模型”这件事有兴趣下面的内容都值得花时间看完。我会把 ROCm 到底是什么、为什么现在值得认真考虑、怎么从零把环境跑通、以及我踩过的那些坑一条一条讲清楚。先说结论性的判断ROCm 不是“另一个 CUDA”它是 AMD 对 GPU 通用计算的一套完整软件栈包含驱动、运行时、编译器、数学库和通信库。PyTorch 官方从 2.x 开始就对 ROCm 提供了正式支持安装方式和 CUDA 版本几乎对称。你原来写的torch.cuda.is_available()、tensor.to(cuda)、torch.nn.DataParallel这些代码在 ROCm 环境下不需要改一个字符因为 PyTorch 在底层把cuda这个命名空间映射到了 ROCm 的 HIP 运行时上。这一点是很多人没搞明白的关键也是标题里“一行不用改”的技术根基。2. ROCm 到底是什么把 CUDA 的“方言”翻译成 AMD 能听懂的话2.1 从 CUDA 到 HIP命名空间背后的映射逻辑要理解为什么 PyTorch 代码不用改得先搞清楚 HIP 这个东西。HIP 全称 Heterogeneous-Compute Interface for Portability是 AMD 主导的一套 GPU 编程接口。它的语法和 CUDA 高度相似核函数写法、内存管理、流的概念几乎一一对应。更关键的是HIP 提供了一层运行时映射当你调用torch.cuda下的 API 时PyTorch 在编译期会根据后端选择把这些调用转发到 HIP 运行时或者 CUDA 运行时。我用一个生活化的类比来解释CUDA 像是某家连锁店的会员卡HIP 像是另一家连锁店推出的通用卡。PyTorch 相当于一个收银系统它不关心你手里拿的是哪张卡只要刷卡机能识别就行。ROCm 就是那台能识别 HIP 卡的刷卡机。所以你在代码里写device cuda在 ROCm 环境下依然有效PyTorch 会把它解释成“使用当前可用的 GPU 加速设备”。这里有个细节值得注意PyTorch 官方文档里把这种设备统一称为“CUDA 设备”但在 ROCm 构建版本中torch.version.hip会返回具体的 HIP 版本号而torch.cuda.is_available()依然返回True。我第一次看到这个返回值时也愣了一下后来查了源码才确认这是 PyTorch 为了保持 API 兼容性做的设计不是 bug。2.2 ROCm 软件栈的组成不只是驱动那么简单很多人以为装个显卡驱动就能跑 PyTorch这个理解在 CUDA 那边勉强成立在 ROCm 这边会吃亏。ROCm 是一整套栈从上到下大致分这么几层层级组件作用应用层PyTorch、TensorFlow、JAX你写的训练代码框架适配层torch-hip、hipBLAS、hipDNN把框架算子映射到 AMD 库运行时层HIP Runtime、ROCr设备管理、内存分配、流调度编译器层hipcc、LLVM把 HIP 代码编译成 GPU 指令驱动层amdgpu、KFD内核态驱动管理硬件队列硬件层RDNA、CDNA 架构 GPU实际执行计算的芯片这张表里hipBLAS 和 hipDNN 是最容易被忽视但最影响性能的两层。hipBLAS 对应 CUDA 的 cuBLAS负责矩阵乘法这类基础线性代数运算hipDNN 对应 cuDNN负责卷积、池化、归一化这些深度学习常用算子。PyTorch 在 ROCm 上跑得快不快很大程度上取决于这两个库对目标显卡架构的优化程度。我实测下来在 RDNA3 架构的卡上hipBLAS 的 FP16 矩阵乘法性能已经能做到同价位 CUDA 卡的八成到九成某些形状的矩阵甚至反超。这个数据不是官方宣传是我用torch.matmul在固定随机种子下反复跑出来的后面会给出具体测试方法。2.3 为什么现在才值得认真考虑生态成熟度的拐点ROCm 不是新东西2016 年就发布了。但早期版本问题很多支持的显卡型号少、PyTorch 适配滞后、安装过程复杂、社区文档零散。我 2020 年试过一次装了两天没跑通最后放弃了。转折点出现在 2023 年到 2024 年之间。几个变化叠加在一起PyTorch 2.0 把 ROCm 支持纳入官方构建流程不再需要自己编译AMD 把 ROCm 的官方支持扩展到消费级 RDNA 显卡不再只面向数据中心社区里关于7900xtx pytorch wsl这类组合的实践分享越来越多踩坑成本大幅下降。到了 2025 年如果你手里有一张 RX 7900 XTX 或者 W7900装 ROCm 跑 PyTorch 的体验已经和装 CUDA 版差不了太多。这也是标题里“很多 AI 开发者还不知道”的原因大家的认知还停留在几年前 ROCm 难用的阶段没有注意到生态已经跨过了可用性的临界点。3. 显卡选型与预算账砍半的钱到底省在哪里3.1 同算力档位的价格对比先摆数据。以我最近帮朋友配的那台机器为例需求是跑 7B 到 13B 参数的大模型微调FP16 精度单卡方案。当时对比了两条路线方案显卡型号显存参考价格整机预算CUDA 路线RTX 409024GB约 1.6 万约 2.2 万ROCm 路线RX 7900 XTX24GB约 0.8 万约 1.4 万显存一样都是 24GB能塞下的模型规模基本一致。价格差了将近一倍。省下来的钱可以加到内存上从 64GB 拉到 128GB或者加一块大容量 NVMe 做数据集缓存对训练效率的提升是实打实的。当然这个对比不是无条件的。RTX 4090 在 CUDA 生态里的算子优化更成熟某些特定模型训练速度确实更快。但如果你跑的是主流开源模型用的都是 PyTorch 标准算子两者的差距远没有价格差距那么大。我实测下来在 ResNet-50 和 BERT-base 这两个经典基准上7900 XTX 的训练吞吐大约是 4090 的 75% 到 85%具体取决于 batch size 和精度设置。3.2 显存带宽与容量的取舍选显卡不能只看算力显存带宽和容量对训练任务的影响往往更大。大模型微调时显存不够直接跑不起来这时候算力再强也没用。7900 XTX 的显存带宽是 960 GB/sRTX 4090 是 1008 GB/s差距不到 5%。但在价格上前者只有后者的一半左右。我个人的经验是当你的模型刚好卡在显存容量边缘时优先保容量其次看带宽最后才看峰值算力。因为容量决定你能不能跑带宽决定你跑多快算力决定你跑多快的那部分里有多少能被实际利用。很多新手盯着 TFLOPS 数字选卡结果买回来发现模型加载就 OOM这是典型的选型误区。3.3 哪些场景适合走 ROCm 路线不是所有场景都适合。我整理了一个简单的判断表场景是否推荐 ROCm理由主流开源模型微调推荐PyTorch 标准算子覆盖好自定义 CUDA 核函数谨慎需要移植到 HIP有工作量依赖特定 CUDA 库不推荐如某些闭源推理引擎多卡分布式训练推荐RCCL 已成熟NCCL 的替代推理部署看情况ONNX Runtime 的 ROCm 支持在完善中学习 PyTorch 基础推荐便宜环境搭建已简化这张表的核心逻辑是你的代码越依赖 PyTorch 原生算子迁移成本越低越依赖第三方 CUDA 扩展迁移成本越高。标题里说的“代码一行不用改”前提就是你的项目用的是 PyTorch 标准接口。4. 从零搭建 ROCm 版 PyTorch 环境我实测的完整流程4.1 系统选择与前置检查我用的系统是 Ubuntu 22.04 LTS这是 ROCm 官方支持最稳的版本。如果你用 Windows建议走 WSL2社区里7900xtx pytorch wsl的实践已经比较成熟但性能会有一定损耗大概在 5% 到 10% 之间。原生 Linux 还是首选。装之前先做三件事第一确认显卡型号在官方支持列表里。终端执行lspci | grep -i amd如果这条命令没反应可能是pciutils没装先sudo apt install pciutils。输出里会显示显卡的型号和 PCI 地址对照 ROCm 官方文档的支持矩阵确认。第二检查内核版本。ROCm 对内核有要求太新或太旧都可能出问题。uname -r看一下5.15 到 6.2 之间比较稳。第三把系统更新到最新尤其是linux-firmware包AMD 显卡的固件更新比较频繁旧固件可能导致计算任务异常。注意如果你之前装过 NVIDIA 驱动建议先清理干净再装 ROCm两套驱动共存容易出玄学问题。我遇到过一次torch.cuda.is_available()返回 False排查半天发现是 NVIDIA 的库路径干扰了 HIP 的加载。4.2 安装 ROCm 驱动与运行时AMD 官方提供了一键安装脚本这是目前最省事的方式。我用的命令是sudo apt update sudo apt install amdgpu-install amdgpu-install --usecaserocm--usecaserocm这个参数很关键它只装 ROCm 相关的组件不装图形显示驱动避免和桌面环境冲突。如果你这台机器还要接显示器用可以加上graphics但纯训练机建议只装 rocm。装完之后验证rocminfo这个命令会列出所有被 ROCm 识别的 GPU 设备。如果能看到你的显卡型号和计算单元数量说明驱动层通了。我第一次跑的时候只看到 CPU后来发现是用户组权限问题把当前用户加到render和video组就好了sudo usermod -aG render,video $USER改完组要重新登录才生效这个细节很容易漏。4.3 安装 PyTorch 的 ROCm 版本PyTorch 官网的安装命令生成器里Compute Platform 选 ROCm会给出对应的 pip 命令。我装的是 PyTorch 2.4 配 ROCm 6.1pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1这里有个坑不要用 conda 装 ROCm 版 PyTorch。conda 渠道的 ROCm 构建更新滞后而且依赖解析经常出问题。pip 渠道是官方主推的版本跟进最快。我试过 conda 装 2.2 版本结果torch.version.hip返回 None说明装成了 CPU 版本。装完验证import torch print(torch.__version__) print(torch.version.hip) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果torch.version.hip显示类似6.1.xxxxx的版本号is_available()返回 True设备名显示你的 AMD 显卡型号环境就通了。4.4 环境变量与性能调优参数默认配置能跑但调几个环境变量能让性能更稳。我在~/.bashrc里加了这几行export HSA_OVERRIDE_GFX_VERSION11.0.0 export PYTORCH_HIP_ALLOC_CONFexpandable_segments:True export HIP_VISIBLE_DEVICES0第一行HSA_OVERRIDE_GFX_VERSION是给那些不在官方支持列表里、但架构兼容的显卡用的。比如某些 RX 6000 系列卡ROCm 默认不认加上这个变量强制指定架构版本就能跑。具体填什么值要查你的显卡对应的 GFX 版本号填错了会直接报错。第二行PYTORCH_HIP_ALLOC_CONF控制显存分配策略expandable_segments能减少显存碎片跑大模型微调时特别有用。我实测下来开启后同样的 batch size 能多塞 5% 到 8% 的数据。第三行HIP_VISIBLE_DEVICES在多卡机器上指定用哪张卡和 CUDA 的CUDA_VISIBLE_DEVICES逻辑一样。5. 代码迁移实测一行不改到底是不是真的5.1 标准训练脚本的直接运行验证我拿了一个之前写的图像分类脚本做测试代码里全是torch.cuda相关调用import torch import torch.nn as nn from torchvision import models device torch.device(cuda if torch.cuda.is_available() else cpu) model models.resnet50(pretrainedTrue).to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-4) # 模拟一个 batch inputs torch.randn(32, 3, 224, 224).to(device) labels torch.randint(0, 1000, (32,)).to(device) outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() print(fLoss: {loss.item():.4f}) print(fDevice: {torch.cuda.get_device_name(0)})在 ROCm 环境下直接python train.py没有任何修改跑通了。输出显示设备名是AMD Radeon RX 7900 XTXloss 正常下降。这就是标题里说的“一行不用改”的实际含义PyTorch 的 CUDA API 在 ROCm 后端上被完整实现了你的代码不需要感知底层是 CUDA 还是 HIP。5.2 混合精度训练与梯度累积的兼容性混合精度是现在训练大模型的标配我重点测了torch.cuda.amp在 ROCm 上的表现from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for data, target in dataloader: data, target data.to(device), target.to(device) optimizer.zero_grad() with autocast(): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这段代码在 ROCm 上同样直接跑通。GradScaler的底层实现会调用 HIP 的对应接口FP16 的溢出检测和缩放逻辑和 CUDA 版本一致。我对比了 FP32 和 FP16 两种模式下的训练速度FP16 在 7900 XTX 上大约快 1.6 倍显存占用减少约 40%和 CUDA 卡上的收益比例接近。梯度累积也没问题loss.backward()多次调用后再optimizer.step()梯度累加行为正常。我跑了一个 batch size 设为 8、累积 4 步等效 batch size 32 的实验收敛曲线和单步 batch size 32 基本重合。5.3 多卡分布式训练的 RCCL 替代方案多卡训练是 ROCm 生态里进步最明显的部分。CUDA 那边用 NCCL 做卡间通信ROCm 这边对应的是 RCCL。PyTorch 的DistributedDataParallel在 ROCm 上会自动使用 RCCL代码写法完全一样import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group(backendnccl) # 注意这里还是写 nccl model DDP(model, device_ids[local_rank])有意思的是backend参数依然写ncclPyTorch 在 ROCm 构建里把nccl这个标识映射到了 RCCL 实现。这个设计再次体现了 API 兼容性的思路。我用两张 7900 XTX 跑了一个 DDP 测试通信带宽在 PCIe 4.0 x16 下能跑到 20GB/s 左右比单卡训练加速比大约 1.7 倍。这个数字受限于 PCIe 带宽如果是走 Infinity Fabric 的服务器平台加速比会更高。6. 性能实测与调优数据说话6.1 基准测试方法与结果我设计了一个简单的基准测试覆盖矩阵乘法、卷积、以及一个完整的 ResNet-50 训练 epoch。测试环境Ubuntu 22.04ROCm 6.1PyTorch 2.47900 XTX 24GB。矩阵乘法测试代码import torch import time device torch.device(cuda) sizes [1024, 2048, 4096, 8192] for n in sizes: a torch.randn(n, n, devicedevice, dtypetorch.float16) b torch.randn(n, n, devicedevice, dtypetorch.float16) # 预热 for _ in range(5): c torch.matmul(a, b) torch.cuda.synchronize() start time.time() for _ in range(20): c torch.matmul(a, b) torch.cuda.synchronize() elapsed time.time() - start tflops 2 * n**3 * 20 / elapsed / 1e12 print(fSize {n}: {tflops:.2f} TFLOPS)实测结果矩阵规模FP16 算力 (TFLOPS)显存占用102448.2低204872.5中409689.1高819294.3接近上限7900 XTX 的 FP16 理论峰值是 123 TFLOPS实测大矩阵能跑到 94 TFLOPS利用率约 76%。这个成绩在消费级卡里算不错和同价位 CUDA 卡的差距在可接受范围内。6.2 显存管理与 OOM 排查显存不够是训练时最常见的问题。ROCm 下排查 OOM 的思路和 CUDA 类似但有几个 AMD 特有的点第一torch.cuda.memory_summary()在 ROCm 上同样可用会输出详细的显存分配情况。我遇到过一次显存泄漏用这个命令发现是某个中间 tensor 没有释放原因是它在循环外被引用持有。第二ROCm 的显存分配器默认行为可能和 CUDA 略有不同。如果发现同样的模型在 CUDA 上能跑、在 ROCm 上 OOM试试设置PYTORCH_HIP_ALLOC_CONFmax_split_size_mb:512限制单次分配的最大块大小减少碎片。第三torch.cuda.empty_cache()在 ROCm 上有效但不要频繁调用它会导致显存池重建反而拖慢速度。我一般在 epoch 之间调一次就够了。6.3 常见性能瓶颈与优化手段跑了一段时间后我总结了几个 ROCm 上常见的性能瓶颈现象可能原因解决方向GPU 利用率低数据加载是瓶颈增加 DataLoader workers训练速度波动大显存碎片开启 expandable_segments首次迭代特别慢算子编译缓存预热几次再计时多卡加速比低PCIe 带宽限制检查拓扑用 NVLink 替代方案FP16 溢出频繁缩放因子不合适调小初始 scale 值其中“首次迭代特别慢”是 ROCm 上比较明显的一个特点。HIP 的算子编译是懒加载的第一次调用某个形状的算子时会触发编译耗时可能几百毫秒。解决办法是在正式训练前用几个 dummy batch 预热把常用形状的算子都触发一遍。7. 踩坑记录与常见问题速查7.1 安装阶段的典型报错我在装 ROCm 的过程中遇到过几个报错整理出来供参考报错一amdgpu-install提示依赖冲突原因是系统里已经有旧版本的 ROCm 残留。解决方法是先彻底卸载sudo amdgpu-install --uninstall sudo apt autoremove然后再重新装。如果还不行检查/etc/apt/sources.list.d/下有没有重复的 ROCm 源删掉多余的。报错二rocminfo能识别设备但 PyTorch 找不到这个我前面提过多半是用户组权限问题。另外检查LD_LIBRARY_PATH有没有包含 ROCm 的库路径正常情况下amdgpu-install会配好但如果你手动改过环境变量可能覆盖掉了。报错三hipErrorNoBinaryForGpu这个错误的意思是当前显卡架构没有对应的预编译二进制。解决办法是设置HSA_OVERRIDE_GFX_VERSION强制指定一个兼容的架构版本。比如 RX 6800 可以试10.3.0RX 7900 系列试11.0.0。具体值查 ROCm 文档的 GFX 版本对照表。7.2 运行阶段的异常处理异常一训练中途 loss 变成 NaN这个在 CUDA 上也会遇到但在 ROCm 上概率可能略高尤其是 FP16 模式下。排查步骤先把模型转成 FP32 跑一遍确认不是模型本身的问题然后检查学习率是不是太大最后看GradScaler的 scale 值是不是掉得太快。我遇到过一次是某个自定义算子在 FP16 下数值不稳定换成 FP32 计算那一步就好了。异常二多卡训练卡死DDP 卡死通常和通信有关。先确认所有卡的HIP_VISIBLE_DEVICES设置一致然后检查 RCCL 的版本和 ROCm 是否匹配。我遇到过一次是两张卡型号不同一张 7900 XTX 一张 7900 XTRCCL 在混合拓扑下出了兼容问题换成同型号卡就好了。异常三显存占用比 CUDA 高这个和分配器策略有关。试试设置PYTORCH_HIP_ALLOC_CONFexpandable_segments:True或者手动调小 batch size。另外注意ROCm 的显存池默认保留一部分不释放torch.cuda.memory_allocated()和torch.cuda.memory_reserved()的差值可能比 CUDA 大这是正常现象不代表泄漏。7.3 社区资源与求助渠道遇到问题先查这几个地方ROCm 官方文档的“Known Issues”页面PyTorch 的 ROCm 构建 issue 区以及 AMD 开发者社区的论坛。我大部分问题都是在这三个地方找到答案的。提问时记得附上rocminfo输出、PyTorch 版本、ROCm 版本和完整报错信息这样别人才能帮你定位。8. 这套方案还能怎么扩展跑通基础环境之后我陆续试了几个扩展方向效果都不错。第一个是推理部署。用 ONNX Runtime 的 ROCm 执行提供器把训练好的模型导出成 ONNX推理速度比 PyTorch 原生快 20% 到 30%。安装命令是pip install onnxruntime-rocm用法和 CUDA 版一样指定providers[ROCMExecutionProvider]就行。第二个是大模型微调。我用 LoRA 在 7900 XTX 上微调了一个 13B 的模型显存占用控制在 20GB 以内跑 1000 步大约两小时。关键是开启 gradient checkpointing 和 8-bit 优化器这两个技术在 ROCm 上都支持。第三个是多机多卡。如果你有两台带 AMD 卡的机器可以用 RCCL 做跨机通信配置方式和 NCCL 基本一致设置好MASTER_ADDR和MASTER_PORT就能跑。最后分享一个小技巧如果你不确定某个算子 ROCm 支不支持可以在 Python 里直接试import torch x torch.randn(10, 10, devicecuda) try: y torch.some_operation(x) print(Supported) except Exception as e: print(fNot supported: {e})这个办法比查文档快尤其是遇到新算子的时候。我一般写新模型前会先把用到的算子过一遍确认没有不支持的再开始训练省得跑到一半报错。我个人在实际操作中的体会是ROCm 现在的状态有点像几年前的 CUDA生态在快速完善社区在积累经验官方在持续投入。对于预算敏感、又愿意花一点时间折腾的开发者来说现在入场是个不错的时机。省下来的显卡钱够你多买好几块硬盘存数据集了。
