Intel Sapphire Rapids CPUでのStable Diffusion推論の高速化

TL;DR

Hugging Faceは、Optimum Intel/OpenVINO、システムレベルのメモリとスレッドの最適化、BF16/AMXを使用したIntel Extension for PyTorch (IPEX)、およびDPMSolverMultistepSchedulerを組み合わせることで、Intel Sapphire Rapids (Xeon Sapphire Rapids) CPU上でStable Diffusion推論を 32.3 s から ~5 s に高速化できることを示しています。

Diffusersライブラリ – ベースライン性能

diffusers ライブラリは、Stable Diffusion 用のシンプルな Python API を提供します。Amazon EC2 r7iz.metal-16xl インスタンス（64 vCPU、512 GB RAM、Ubuntu 20.04）でデフォルトの float32 パイプラインを使用すると、画像1枚あたりの平均レイテンシは 32.3 seconds となります（推論ステップ20）。同じハードウェア世代（Sapphire Rapids）は、前世代の Ice Lake Xeon よりもすでに高速で、同一ワークロードで約45 seconds かかります。

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 パイプラインはレイテンシを 16.7 s に削減し、ベースラインの Diffusers に比べて2倍の改善となります。

固定形状最適化 – さらに3.5倍

OpenVINO は静的入力形状もサポートしています。パイプラインを固定解像度（例：512 × 512）にリシェイプすることで、レイテンシはさらに 4.7 seconds に低下します。

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 が使用する大規模なマルチギガバイトモデルに影響します。以下のシステム調整が同じ 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 をインストールした後、各パイプラインコンポーネント（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 s から 5.4 seconds に改善し、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 seconds となり、元の Sapphire Rapids ベースライン（32.3 s）に対して 6.5× の改善、そして Ice Lake に比べて約 10× の高速化を示します。

完全な環境スナップショット

  • インスタンス: Amazon EC2 r7iz.metal-16xl（64 vCPU、512 GB RAM）
  • OS: Ubuntu 20.04、Linux カーネル 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

意味合いと次のステップ

一般的な CPU で Stable Diffusion の推論を数秒に短縮することで、新たなユースケースが広がります：

  • GPU インフラなしで、顧客向けアプリケーションにおけるリアルタイム画像生成。
  • 他のモデルのトレーニング用にコスト効果の高い合成データ生成。
  • データプライバシーやレイテンシ制約によりクラウド GPU が利用できないオンプレミス展開。

開発者は、ブログ記事で使用されたフルスクリプトが含まれる提供された GitLab デモリポジトリに従って開始できます。さらにスケールするために、今後の Hugging Face 記事では Sapphire Rapids クラスタ上での拡散モデルの分散ファインチューニングについて取り上げる予定です。

リソース

Sources