AMD显卡跑深度学习：ROCm生态实战指南与性能调优

1次阅读

共计 2919 个字符，预计需要花费 8 分钟才能阅读完成。

当前深度学习领域几乎被 NVIDIA 的 CUDA 生态垄断，这带来了几个显著问题：

高昂的成本：NVIDIA 的专业显卡价格居高不下，特别是最新一代产品
供应链风险：全球芯片短缺情况下，NVIDIA 显卡供应不稳定
生态封闭：CUDA 作为私有技术，限制了硬件选择的多样性

AMD 的 ROCm(RAdeon Open Compute)平台正是为解决这些问题而生。作为一个开源计算平台，ROCm 提供了与 CUDA 类似的功能，但能够在 AMD 显卡上运行。

特性	ROCm	CUDA
API 兼容性	提供 HIP(可转换 CUDA 代码)	原生 CUDA API
框架支持	PyTorch、TensorFlow 部分支持	全面支持主流框架
计算库完备性	基础 BLAS、FFT 等库完整	计算库生态极其丰富
跨平台支持	Linux 为主	全平台支持
开源程度	完全开源	闭源

检查系统兼容性（需要 Ubuntu 20.04/22.04 LTS）：

uname -m  # 确认是 x86_64 架构
lsb_release -a  # 确认 Ubuntu 版本

安装依赖项：

sudo apt update && sudo apt install -y libnuma-dev libpci-dev wget

添加 ROCm 仓库并安装：

wget -q -O - https://repo.radeon.com/rocm/rocm.gpg.key | sudo apt-key add -
echo 'deb [arch=amd64] https://repo.radeon.com/rocm/apt/5.7 ubuntu main' | sudo tee /etc/apt/sources.list.d/rocm.list
sudo apt update && sudo apt install -y rocm-opencl-runtime

验证安装：

/opt/rocm/bin/rocminfo  # 应显示设备信息

常见问题解决：
– 如果遇到clang: error: unsupported option '-fno-plt'，需要降级 LLVM 版本
– 权限问题可通过将用户加入 video 和render组解决

获取源码：

git clone --recursive https://github.com/pytorch/pytorch
cd pytorch

安装编译依赖：

pip install -r requirements.txt
sudo apt install -y hipblas hipsparse rocrand rccl

编译配置：

export PYTORCH_ROCM_ARCH="gfx906"  # 根据你的 AMD 显卡调整
python setup.py install --user --no-cache-dir --cmake-only

执行编译：

python setup.py install --user --no-cache-dir

--no-cache-dir参数可避免使用可能冲突的缓存文件，这在交叉编译时特别重要。

CUDA 原代码：

__global__ void vectorAdd(float* A, float* B, float* C, int numElements) {
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    if (i < numElements) {C[i] = A[i] + B[i];
    }
}

HIP 改写版本（注意内存对齐）：

#include <hip/hip_runtime.h>

__global__ void vectorAdd(float* A, float* B, float* C, int numElements) {
    int i = hipBlockDim_x * hipBlockIdx_x + hipThreadIdx_x;
    // 确保内存访问对齐（AMD 架构敏感点）if (i < numElements && i % 4 == 0) {float4 a = *reinterpret_cast<float4*>(&A[i]);
        float4 b = *reinterpret_cast<float4*>(&B[i]);
        float4 c;
        c.x = a.x + b.x;
        c.y = a.y + b.y;
        c.z = a.z + b.z;
        c.w = a.w + b.w;
        *reinterpret_cast<float4*>(&C[i]) = c;
    }
}

利用矩阵计算单元：MI100/200 系列包含 Matrix Core，适合 4 ×4 FP16 矩阵运算
批处理优化：增大 batch size 以利用并行计算能力
内存访问 ：使用__restrict__ 关键字减少指针别名分析开销

TensorFlow custom op：需要重新编译时指定 ROCm 路径
```
TF_NEED_ROCM=1 ./configure
```

PyTorch 分布式训练：需使用 RCCL 代替 NCCL

torch.distributed.init_process_group(backend='gloo')  # 或 'hccl'

安装分析工具：
```
sudo apt install -y rocm-profiler
```
收集性能数据：
```
rocprof --stats ./your_application
```

生成可视化报告：

rocprof --stats --timestamp on -o output.csv ./your_application

测试脚本示例：

import torch
import torchvision.models as models
import time

model = models.resnet50().to('cuda')  # 自动识别 ROCm 设备
input = torch.randn(64, 3, 224, 224).to('cuda')

# 预热
for _ in range(10):
    _ = model(input)

# 正式测试
start = time.time()
for _ in range(100):
    _ = model(input)
torch.cuda.synchronize()
print(f"Time per batch: {(time.time()-start)/100:.4f}s")

使用 AMP(Automatic Mixed Precision)：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
with autocast():
    output = model(input)
    loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()