# 用 AMD 显卡跑 PyTorch：ROCm 环境搭建与实战踩坑指南

## 为什么选 AMD + ROCm？

很多 AI 初学者以为只有 NVIDIA 能跑深度学习。其实 AMD 的 ROCm（Radeon Open Compute）生态在过去两年进步巨大——PyTorch 2.x 已经原生支持 ROCm，`torch.cuda.*` 的代码不用改就能在 AMD 显卡上跑。

2026 年最新的 ROCm 7.2 支持 RDNA 4 消费级显卡（如 RX 9060 XT），这对预算有限的学生开发者来说是个好消息。

## 环境搭建：Docker 方案（推荐）

如果你不想折腾驱动依赖，直接用 AMD 官方预编译的 Docker 镜像：

```bash
docker pull rocm/pytorch:rocm7.2_ubuntu24.04_py3.13_pytorch_release_2.10.0

docker run -it \
    --cap-add=SYS_PTRACE \
    --security-opt seccomp=unconfined \
    --device=/dev/kfd \
    --device=/dev/dri \
    --group-add video \
    --ipc=host \
    --shm-size 8G \
    rocm/pytorch:rocm7.2_ubuntu24.04_py3.13_pytorch_release_2.10.0
```

容器启动后验证 GPU 是否被识别：

```python
import torch
print(torch.cuda.is_available())   # True
print(torch.cuda.device_count())   # 输出你的 AMD GPU 数量
```

## 直接安装：pip wheels（适合 Ubuntu 22.04/24.04）

如果你已经在用 Linux 桌面环境，可以直接装：

```bash
pip3 install --pre torch torchvision torchaudio \
    --index-url https://download.pytorch.org/whl/nightly/rocm7.0
```

装完后跑个简单测试确认 ROCm 栈正常工作：

```python
# 矩阵乘法性能对比：CPU vs AMD GPU
import torch
import time

size = 4096
a = torch.randn(size, size)
b = torch.randn(size, size)

# CPU
start = time.time()
c = a @ b
print(f"CPU: {time.time() - start:.2f}s")

# AMD GPU
a_gpu = a.to("cuda")
b_gpu = b.to("cuda")
torch.cuda.synchronize()
start = time.time()
c_gpu = a_gpu @ b_gpu
torch.cuda.synchronize()
print(f"GPU: {time.time() - start:.2f}s")
```

我的 RX 9060 XT 上矩阵乘法快了约 40 倍。

## 2026 年关键配置

几个 ROCm 7.2 环境变量值得注意：

```bash
export HSA_OVERRIDE_GFX_VERSION=12.0.0    # 启用 RDNA 4 指令集
export PYTORCH_ALLOC_CONF=max_split_size_mb:32  # 修复 VRAM 预分配 Bug
export TORCH_ROCM_AOTRITON=1              # 开启实验性 Flash Attention
```

第一个变量尤其关键——如果不设，RDNA 4 新卡可能被当成旧架构，性能打折扣。

## 踩过的坑

**1. ROCm 和 NVIDIA 驱动不能共存。** 如果你之前装过 CUDA，建议用 Docker 隔离，或者双系统切纯净 Ubuntu。

**2. `hipErrorOutOfMemory` 错误**——设置 `PYTORCH_ALLOC_CONF=max_split_size_mb:32` 解决。这是 ROCm 7.2 早期版本的一个 Bug，已在 7.2.4 中修复，升级可解决。

**3. 不要混装 pip 和 conda 的 PyTorch。** 要么全用 pip，要么全用 conda。我踩过混装导致 `libtorch_cpu.so` 符号冲突的坑，排查了两个小时。

## 总结

AMD ROCm + PyTorch 现在已经是一个可用的 AI 开发栈，尤其是对学生来说——一张 RX 显卡的价格远低于同性能的 NVIDIA 卡。加上 Docker 降低了环境门槛，不用担心驱动问题。如果你手头有 AMD 显卡，不妨试试，`torch.cuda.*` 的代码一行都不用改。

> 本文为 AMD AI 开发者征文大赛参赛作品，使用 ROCm 7.2 + PyTorch 2.10 环境。