在 Intel Sapphire Rapids CPU 上加速 Stable Diffusion 推理

TL;DR

Hugging Face 演示了通过结合 Optimum Intel/OpenVINO、系统级内存和线程优化、Intel Extension for PyTorch (IPEX) 的 BF16/AMX 支持以及 DPMSolverMultistepScheduler，能够将 Stable Diffusion 推理在 Intel Sapphire Rapids（Xeon Sapphire Rapids）CPU 上的耗时从 32.3 s 加速至 约 5 s。

Diffusers 库 – 基准性能

diffusers 库提供了一个简洁的 Python API 用于 Stable Diffusion。使用默认的 float32 pipeline 在 Amazon EC2 r7iz.metal-16xl 实例（64 vCPU，512 GB RAM，Ubuntu 20.04）上，每张图片（20 步推理）的平均延迟为 32.3 秒。同一代的 Sapphire Rapids 硬件已经比上一代 Ice Lake Xeon 更快，后者完成相同工作负载大约需要 45 秒。

Optimum Intel 和 OpenVINO – 2× 加速

Optimum Intel 提供了 StableDiffusionPipeline 的直接替代实现 OVStableDiffusionPipeline。使用 pip install optimum[openvino] 安装后，可自动将 PyTorch 模型转换为 OpenVINO 格式并以 bfloat16 进行推理。

from optimum.intel.openvino import OVStableDiffusionPipeline
ov_pipe = OVStableDiffusionPipeline.from_pretrained(model_id, export=True)
latency = elapsed_time(ov_pipe, prompt)
print(latency)  # → ~16.7 s

OpenVINO pipeline 将延迟降低至 16.7 秒，相较于原始 Diffusers 基准提升了 2 倍。

固定形状优化 – 再提升 3.5×

OpenVINO 还支持静态输入形状。将 pipeline 重塑为固定分辨率（例如 512 × 512）后，延迟进一步下降至 4.7 秒。

ov_pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
latency = elapsed_time(ov_pipe, prompt)

结合 Sapphire Rapids，这相当于比 Ice Lake 提升了近 10× 的速度。

系统级优化 – 代码无需改动即可提升 3 倍

内存分配和线程调度会显著影响 Stable Diffusion 使用的数十 GB 大模型。以下系统调优在同一 EC2 实例上完成：

  • 安装并预加载 jemalloc 以获得高性能内存管理。
  • 安装 intel‑mkl 并预加载 libiomp5.so 以启用 Intel OpenMP 运行时。
  • 将 OMP_NUM_THREADS=32 设置为与核心数匹配。
  • 使用 numactl -C 0-31 将 Python 进程绑定到部分核心。
sudo apt-get install -y libjemalloc-dev intel-mkl
export LD_PRELOAD=$LD_PRELOAD:/usr/lib/x86_64-linux-gnu/libjemalloc.so
export LD_PRELOAD=$LD_PRELOAD:/usr/lib/x86_64-linux-gnu/libiomp5.so
export OMP_NUM_THREADS=32
numactl -C 0-31 python sd_blog_1.py

这些改动将原始 Diffusers 的延迟从 32.3 s 降至 11.8 s，约提升 3 倍。

IPEX 与 BF16 – 再提升 2 倍

Intel Extension for PyTorch (IPEX) 利用 Sapphire Rapids 上的 AVX‑512 VNNI 与 AMX（Advanced Matrix Extensions）。在安装 intel_extension_for_pytorch==1.13.100 后，pipeline 的每个组件（UNet、VAE、文本编码器、safety checker）都会被转换为 channels‑last 布局，并在 bfloat16 模式下通过 IPEX 进行优化。

import torch, intel_extension_for_pytorch as ipex
# Convert modules to channels‑last
pipe.unet = pipe.unet.to(memory_format=torch.channels_last)
# ... repeat for vae, text_encoder, safety_checker
# Optimize with IPEX
pipe.unet = ipex.optimize(pipe.unet.eval(), dtype=torch.bfloat16, inplace=True, sample_input=input_example)
# Run under autocast
with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    latency = elapsed_time(pipe, prompt)
    print(latency)  # → ~5.4 s

延迟从 11.8 s 降至 5.4 秒，得益于 BF16 与 AMX 的双重加速，提升超过 2 倍。

更快的调度器 – 最终整体提升 6.5×

Diffusers 支持可互换的调度器来控制去噪步骤的安排。DPMSolverMultistepScheduler 在仅 20 步的情况下提供了最佳的速度/质量平衡。

from diffusers import DPMSolverMultistepScheduler
scheduler = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler)
latency = elapsed_time(pipe, prompt)
print(latency)  # → ~5.05 s

使用该调度器后，最终延迟为 5.05 秒，相较于最初的 Sapphire Rapids 基准（32.3 s）提升 6.5×，并且大约 10× 快于 Ice Lake。

完整环境快照

基准在以下环境中完成：

  • 实例：Amazon EC2 r7iz.metal-16xl（64 vCPU，512 GB RAM）
  • 操作系统：Ubuntu 20.04，Linux kernel 5.15.0-1031-aws
  • 库：libjemalloc‑dev 5.2.1‑1，intel‑mkl 2020.0.166‑1，PyTorch 1.13.1，Intel Extension for PyTorch 1.13.1，Transformers 4.27.2，Diffusers 0.14，Accelerate 0.17.1，OpenVINO 2023.0.0.dev20230217，Optimum 1.7.1，Optimum‑Intel 1.7

含义与后续工作

将 Stable Diffusion 推理时间压缩至几秒钟的 commodity CPU 上，为以下场景打开了可能性：

  • 在无需 GPU 基础设施的面向客户的应用中实现实时图像生成。
  • 以成本效益高的方式生成合成数据，用于训练其他模型。
  • 在对数据隐私或延迟有严格要求、无法使用云 GPU 的本地部署环境中使用。

开发者可以先参考提供的 GitLab 示例仓库，其中包含本文博客中使用的完整脚本。后续将有 Hugging Face 文章进一步介绍在 Sapphire Rapids 集群上进行分布式微调扩散模型的方案。

资源

Sources