AMD显卡加速Qwen模型推理：从ROCm环境配置到性能调优实战

1次阅读

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

AMD 显卡在深度学习领域长期面临两大核心挑战：

CUDA 生态缺失 ：NVIDIA 的 CUDA 经过 15 年积累形成完整工具链（cuDNN/cuBLAS/TensorRT），而 ROCm 的 HIP 接口兼容层需要额外转换开销
显存管理粗放 ：RDNA 架构的 Infinity Cache 对连续大矩阵运算不友好，显存带宽利用率常低于 60%（对比 NVIDIA 的 90%+）

在 Qwen 这类 Transformer 模型上，关键差异体现在：

矩阵乘法加速
CUDA 使用 Tensor Core 的 WMMA API 进行 4 ×4 矩阵块计算
ROCm 通过 HIPBLAS 的 hgemm 实现，依赖 GPU 的矩阵引擎（CDNA 架构才有专用单元）
Attention 层优化
CUDA 有 FlashAttention- 2 等高度优化的内核
ROCm 需手动实现融合算子（参考代码片段）

# AMD 专用 attention 优化（需开启__HIP_PLATFORM_AMD__）def attention_forward(query, key, value):
    with torch.hip.amp.autocast(dtype=torch.float16):  # 必须启用 AMP
        scores = torch.matmul(query, key.transpose(-2, -1)) 
        attn = torch.softmax(scores, dim=-1)
        return torch.matmul(attn, value)  # HipBLAS 自动选择最优内核

FROM rocm/pytorch:rocm5.6_ubuntu20.04_py3.8_pytorch_2.0.1

# 必须安装的依赖
RUN apt-get update && apt-get install -y \
    hipblas \
    rocrand \
    hipsparse  # 稀疏计算支持

# 解决常见 lib 冲突
RUN ln -sf /opt/rocm/lib/libhsa-runtime64.so.1 /opt/rocm/lib/libhsa-runtime64.so

关键配置检查：

# 验证 ROCm 识别
rocminfo | grep 'gfx90a'  # RX7900XTX 应返回 gfx1100

# 检查 PyTorch 链接
ldd /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch.so | grep libamdhip64

核心优化点：

from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen-7B",
    torch_dtype=torch.float16,
    device_map="auto",
    attn_implementation="eager"  # 必须关闭 flash_attention
).eval()

# 关键：启用 HIP 优化
torch._C._jit_set_profiling_executor(False)
model = torch.compile(model, mode="max-autotune", fullgraph=True)

# 显存优化配置（AMD 专用）torch.hip.set_per_process_memory_fraction(0.8)  # 保留 20% 显存防 OOM

# 编译配置（batch_size=1/8/32 自适应）opt_config = {
    "max_autotune": True,
    "dynamic": True,
    "backend": "inductor",
    "options": {
        "triton.cudagraphs": False,  # AMD 必须关闭
        "shape_padding": True       # 动态 padding 优化
    }
}
model = torch.compile(model, **opt_config)

内存 coalescing

torch.backends.hip.enabled = True
torch.backends.hip.allow_tf32 = True  # 开启矩阵加速

Kernel 融合

export PYTORCH_HIP_KERNEL_SPLIT=1  # 强制内核合并

错误：hipErrorNoBinaryForGpu
原因：ROCm 未针对当前 GPU 编译内核

修复：

export HSA_OVERRIDE_GFX_VERSION=11.0.0  # RX7900XTX 专用

错误：ROCblas_status_invalid_size
原因：矩阵尺寸非 64 倍数

修复：

pad_len = (64 - seq_len % 64) % 64
inputs = torch.nn.functional.pad(inputs, (0, pad_len))

# 必须同时设置这三项
torch.backends.hip.matmul.allow_fp16_reduced_precision_reduction = True
torch.backends.hip.allow_tf32 = True
torch.set_float32_matmul_precision('high')

测试设备：RX7900XTX(24G) vs A100(40G)