ERNIE-Image AMD GPU 部署全指南:ROCm 零修改运行

May 22, 2026

ERNIE-Image AMD GPU 部署全指南:ROCm 零修改运行

摘要:AMD 官方宣布 ERNIE-Image 在其 GPU 上 Day-0 支持——无需任何代码修改,通过 Diffusers + ROCm 即可运行。本文从环境搭建到内存优化,手把手教你在 AMD Instinct MI355X 和 Radeon AI PRO R9700 上部署 ERNIE-Image。


一、为什么 AMD GPU 部署很重要?

长期以来,AI 图像生成领域被 NVIDIA CUDA 生态垄断。ERNIE-Image 作为百度开源的 8B 参数 DiT 模型,最初仅验证了 NVIDIA GPU 的兼容性。2026 年 4 月,AMD 官方博客宣布 Day-0 支持——ERNIE-Image 在 AMD GPU 上零代码修改运行。

这意味着:

  • 打破 CUDA 垄断:AMD GPU 用户可以直接使用 ERNIE-Image
  • 降低硬件门槛:AMD Radeon AI PRO R9700 仅需 ~$1,000,远低于 NVIDIA RTX 6000 Ada 的 $6,800
  • 数据中心级选择:Instinct MI355X 提供 288GB HBM3e,适合大规模部署

二、硬件平台

已验证的 AMD GPU

GPU 型号 架构 VRAM 目标场景
AMD Instinct MI355X CDNA 4 (gfx950) 288GB HBM3e 数据中心 AI 训练/推理
AMD Radeon AI PRO R9700 RDNA 4 (gfx1201) 32GB GDDR6 专业工作站 AI 推理

R9700 规格亮点:64 个计算单元、4096 流处理器、128 个 AI 加速器、64MB Infinity Cache、PCIe 5.0 x16、300W TBP、Linux ECC 内存支持。

ERNIE-Image 模型大小

组件 类型 大小 作用
Transformer ErnieImageTransformer2DModel 15 GB 扩散骨干网络
Text Encoder Mistral3Model 7.2 GB 文本编码
Prompt Enhancer Ministral3ForCausalLM 7.2 GB 自动提示词增强
VAE AutoencoderKLFlux2 161 MB 变分自编码器
总计 — ~29.5 GB —

三、软件环境搭建

Docker 容器设置

# MI355X 使用 rocm7.2.1
docker run -d --name ernie-image-test \
  --device=/dev/kfd --device=/dev/dri \
  --group-add video --group-add render \
  --shm-size=64G \
  -v /path/to/model:/workspace \
  rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.12_pytorch_release_2.9.1 \
  sleep infinity

R9700 使用 rocm7.2

docker run -d --name ernie-image-test
--device=/dev/kfd --device=/dev/dri
--group-add video --group-add render
--shm-size=64G
-v /path/to/model:/workspace
rocm/pytorch:rocm7.2
sleep infinity

验证 GPU 可用性

import torch
print(f'PyTorch: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
print(f'Device: {torch.cuda.get_device_name(0)}')
print(f'VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.0f} GB')

注意:ROCm 的 HIP 兼容层确保 torch.cuda.* API 在 AMD 硬件上原生工作,无需修改代码。

安装依赖

# 克隆包含 ERNIE-Image 支持的 Diffusers 分支
git clone https://github.com/HsiaWinter/diffusers /workspace/diffusers-ernie
cd /workspace/diffusers-ernie && git checkout add-ernie-image

安装依赖

pip install -e . accelerate Pillow transformers

解压模型

tar xf ERNIE-Image.tar

软件栈版本

组件 版本
Docker Image rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.12_pytorch_release_2.9.1
PyTorch 2.9.1+rocm7.2.x
ROCm (HIP) 7.2.53211
Diffusers 0.38.0.dev0 (add-ernie-image 分支)
Transformers 5.5.3
Accelerate 1.13.0
Python 3.12

四、推理脚本

MI355X(288GB,全精度加载)

from diffusers import ErnieImagePipeline
import torch

加载模型

pipe = ErnieImagePipeline.from_pretrained(
"/workspace/ERNIE-Image",
torch_dtype=torch.bfloat16
)
pipe = pipe.to("cuda")

设置为评估模式

pipe.transformer.eval()
pipe.vae.eval()
pipe.text_encoder.eval()
pipe.pe.eval()

生成图像

seed = 42
generator = torch.Generator(device="cuda").manual_seed(seed)
output = pipe(
prompt="a beautiful sunset over the ocean, golden light, photorealistic",
height=1024,
width=1024,
num_inference_steps=50,
guidance_scale=5.0,
generator=generator
)
output.images[0].save("ernie_amd_output.png")
print("Image generated successfully!")

R9700(32GB,CPU Offload 优化)

R9700 的 32GB VRAM 对 ~29.5GB 的 BF16 模型非常紧张。加载所有组件后仅剩 ~1.17GB,不足以容纳中间张量。

解决方案:使用 Diffusers 的 enable_model_cpu_offload() 技术在 GPU 上时间共享组件。

from diffusers import ErnieImagePipeline
import torch

加载模型

pipe = ErnieImagePipeline.from_pretrained(
"/workspace/ERNIE-Image",
torch_dtype=torch.bfloat16
)

启用 CPU Offload(关键!)

pipe.enable_model_cpu_offload()

设置为评估模式

pipe.transformer.eval()
pipe.vae.eval()
pipe.text_encoder.eval()
pipe.pe.eval()

生成图像

seed = 42
generator = torch.Generator(device="cuda").manual_seed(seed)
output = pipe(
prompt="a beautiful sunset over the ocean, golden light, photorealistic",
height=1024,
width=1024,
num_inference_steps=50,
guidance_scale=5.0,
generator=generator
)
output.images[0].save("ernie_r9700_output.png")
print("Image generated successfully!")

R9700 VRAM 详细分析

组件 VRAM 占用 累计
Text Encoder 7.18 GiB 7.18 GiB
Prompt Enhancer 6.38 GiB 13.56 GiB
VAE 0.17 GiB 13.73 GiB
Transformer 14.96 GiB 28.69 GiB
剩余 — ~1.17 GiB

使用 CPU Offload 后,峰值 VRAM 只需容纳最大组件(Transformer ~15GB)+ 中间张量,完全在 32GB 范围内。


五、CUDA → ROCm 迁移指南

从 NVIDIA CUDA 迁移到 AMD ROCm 时,需要移除以下 NVIDIA 特定配置:

项目 NVIDIA (CUDA) AMD (ROCm) 操作
CUBLAS_WORKSPACE_CONFIG 需要设置 不适用 移除
torch.backends.cudnn.* cuDNN 配置 使用 MIOpen 移除
torch.use_deterministic_algorithms 支持 部分支持 如有需要则移除
torch.cuda.* API 原生 HIP 兼容层 无需修改
Attention 后端 Flash Attention / cuDNN AOTriton 自动选择

AOTriton 是 AMD 基于 Triton 的注意力内核,针对 ROCm 优化。PyTorch 会自动将其选为 Scaled Dot-Product Attention 后端。


六、性能对比与注意事项

MI355X vs R9700

指标 MI355X R9700
架构 CDNA 4 RDNA 4
VRAM 288GB HBM3e 32GB GDDR6
精度 BF16 全精度 BF16 + CPU Offload
推理速度 最快 中等(CPU Offload 有开销)
适用场景 大规模推理/训练 个人工作站/原型开发

注意事项

  1. Diffusers 分支:需要切换到 add-ernie-image 分支,官方 Diffusers 尚未合并
  2. ROCm 版本:ROCm 7.2 是最低要求,建议升级到最新补丁版本
  3. 内存交换:R9700 使用 CPU Offload 会引入额外延迟,推理时间比全精度加载长约 30-50%
  4. Linux 驱动:确保安装最新的 AMD GPU 驱动(amdgpu 内核模块)

七、总结

AMD 对 ERNIE-Image 的 Day-0 支持是一个重要里程碑,标志着 AI 图像生成领域 CUDA 垄断的松动。对于拥有 AMD GPU 的用户来说,ERNIE-Image 提供了零代码修改的即用体验。

  • 数据中心用户:MI355X 提供 288GB HBM3e,可全精度加载,适合大规模部署
  • 工作站用户:R9700 通过 CPU Offload 在 32GB VRAM 上运行,性价比突出
  • 开发者:torch.cuda.* API 通过 HIP 兼容层无缝工作,迁移成本极低

随着 ROCm 生态的持续改进,AMD GPU 在 AI 推理领域的位置将越来越重要。ERNIE-Image 作为首批 Day-0 支持的模型,为 AMD 用户提供了一个高质量、免费商用的文本到图像生成选择。


参考资料

  1. AMD: Day-0 Support for Baidu ERNIE-Image on AMD GPUs
  2. AMD ROCm on Consumer GPUs 2026 Guide
  3. GitHub: baidu/ERNIE-Image
  4. HuggingFace: baidu/ERNIE-Image

ERNIE-Image Team