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

摘要

Hugging Face 示範了透過結合 Optimum Intel/OpenVINO、系統層級的記憶體與執行緒最佳化、Intel Extension for PyTorch (IPEX) 搭配 BF16/AMX，以及 DPMSolverMultistepScheduler，將 Intel Sapphire Rapids (Xeon Sapphire Rapids) CPU 上的 Stable Diffusion 推理速度從 32.3 秒 加速至 約 5 秒。

Diffusers 函式庫 – 基線效能

diffusers 函式庫提供簡單的 Python API 以使用 Stable Diffusion。於 Amazon EC2 r7iz.metal-16xl 實例（64 vCPU、512 GB 記憶體、Ubuntu 20.04）上使用預設的 float32 pipeline，平均每張影像（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 秒 降至 11.8 秒，約提升 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、文字編碼器、安全檢查器）皆會轉換為 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 秒 改善至 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 秒）提升 6.5 倍，且大約比 Ice Lake 快 10 倍。

完整環境快照

此基準測試於以下環境執行：

  • Instance：Amazon EC2 r7iz.metal-16xl（64 vCPU、512 GB RAM）
  • OS：Ubuntu 20.04，Linux kernel 5.15.0-1031-aws
  • Libraries：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 推理縮短至數秒的商品化 CPU，開啟了新應用情境：

  • 在面向客戶的應用中，無需 GPU 基礎設施即可即時產生影像。
  • 以具成本效益的方式產生合成資料，用於訓練其他模型。
  • 在資料隱私或延遲限制不允許使用雲端 GPU 的情況下，於本地部署。

開發者可依循提供的 GitLab 示範倉庫開始，該倉庫包含部落格文章中使用的完整腳本。欲進一步擴展，未來的 Hugging Face 文章將說明在 Sapphire Rapids 叢集上分散式微調 diffusion 模型的方式。

資源

Sources