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 有开销) |
| 适用场景 | 大规模推理/训练 | 个人工作站/原型开发 |
注意事项
- Diffusers 分支:需要切换到
add-ernie-image分支,官方 Diffusers 尚未合并 - ROCm 版本:ROCm 7.2 是最低要求,建议升级到最新补丁版本
- 内存交换:R9700 使用 CPU Offload 会引入额外延迟,推理时间比全精度加载长约 30-50%
- 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 用户提供了一个高质量、免费商用的文本到图像生成选择。