在英特爾 CPU 上加速 Stable Diffusion 推理

HuggingFace發表於2023-04-13

前一段時間，我們向大家介紹了最新一代的 英特爾至強 CPU (代號 Sapphire Rapids)，包括其用於加速深度學習的新硬體特性，以及如何使用它們來加速自然語言 transformer 模型的 分散式微調 和 推理。

本文將向你展示在 Sapphire Rapids CPU 上加速 Stable Diffusion 模型推理的各種技術。後續我們還計劃釋出對 Stable Diffusion 進行分散式微調的文章。

在撰寫本文時，獲得 Sapphire Rapids 伺服器的最簡單方法是使用 Amazon EC2 R7iz 系列例項。由於它仍處於預覽階段，你需要 註冊 才能獲得訪問許可權。與之前的文章一樣，我使用的是 r7iz.metal-16xl 例項 (64 個 vCPU，512GB RAM)，作業系統映象為 Ubuntu 20.04 AMI (ami-07cd3e6c4915b2d18)。

本文的程式碼可從 Gitlab 上獲取。我們開始吧！

Diffusers 庫

Diffusers 庫使得用 Stable Diffusion 模型生成影像變得極其簡單。如果你不熟悉 Stable Diffusion 模型，這裡有一個很棒的 圖文介紹。

首先，我們建立一個包含以下庫的虛擬環境: Transformers、Diffusers、Accelerate 以及 PyTorch。

virtualenv sd_inference
source sd_inference/bin/activate
pip install pip --upgrade
pip install transformers diffusers accelerate torch==1.13.1

然後，我們寫一個簡單的基準測試函式，重複推理多次，最後返回單張影像生成的平均延遲。

import time

def elapsed_time(pipeline, prompt, nb_pass=10, num_inference_steps=20):
    # warmup
    images = pipeline(prompt, num_inference_steps=10).images
    start = time.time()
    for _ in range(nb_pass):
        _ = pipeline(prompt, num_inference_steps=num_inference_steps, output_type="np")
    end = time.time()
    return (end - start) / nb_pass

現在，我們用預設的 float32 資料型別構建一個 StableDiffusionPipeline，並測量其推理延遲。

from diffusers import StableDiffusionPipeline

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id)
prompt = "sailing ship in storm by Rembrandt"
latency = elapsed_time(pipe, prompt)
print(latency)

平均延遲為 32.3 秒。正如這個英特爾開發的 Hugging Face Space 所展示的，相同的程式碼在上一代英特爾至強 (代號 Ice Lake) 上執行需要大約 45 秒。

開箱即用，我們可以看到 Sapphire Rapids CPU 在沒有任何程式碼更改的情況下速度相當快！

現在，讓我們繼續加速它吧！

Optimum Intel 與 OpenVINO

Optimum Intel 用於在英特爾平臺上加速 Hugging Face 的端到端流水線。它的 API 和 Diffusers 原始 API 極其相似，因此所需程式碼改動很小。

Optimum Intel 支援 OpenVINO，這是一個用於高效能推理的英特爾開源工具包。

Optimum Intel 和 OpenVINO 安裝如下:

pip install optimum[openvino]

相比於上文的程式碼，我們只需要將 StableDiffusionPipeline 替換為 OVStableDiffusionPipeline 即可。如需載入 PyTorch 模型並將其實時轉換為 OpenVINO 格式，你只需在載入模型時設定 export=True。

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

# Don't forget to save the exported model
ov_pipe.save_pretrained("./openvino")

OpenVINO 會自動最佳化 bfloat16 模型，最佳化後的平均延遲下降到了 16.7 秒，相當不錯的 2 倍加速。

上述 pipeline 支援動態輸入尺寸，對輸入影像 batch size 或解析度沒有任何限制。但在使用 Stable Diffusion 時，通常你的應用程式僅限於輸出一種 (或幾種) 不同解析度的影像，例如 512x512 或 256x256。因此，透過固定 pipeline 的輸出解析度來解鎖更高的效能增益有其實際意義。如果你需要不止一種輸出解析度，您可以簡單地維護幾個 pipeline 例項，每個解析度一個。

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

固定輸出解析度後，平均延遲進一步降至 4.7 秒，又獲得了額外的 3.5 倍加速。

如你所見，OpenVINO 是加速 Stable Diffusion 推理的一種簡單有效的方法。與 Sapphire Rapids CPU 結合使用時，和至強 Ice Lake 的最初效能的相比，推理效能加速近 10 倍。

如果你不能或不想使用 OpenVINO，本文下半部分會展示一系列其他最佳化技術。繫好安全帶！

系統級最佳化

擴散模型是數 GB 的大模型，影像生成是一種記憶體密集型操作。透過安裝高效能記憶體分配庫，我們能夠加速記憶體操作並使之能在 CPU 核之間並行處理。請注意，這將更改系統的預設記憶體分配庫。你可以透過解除安裝新庫來返回預設庫。

jemalloc 和 tcmalloc 是兩個很有意思的記憶體最佳化庫。這裡，我們使用 jemalloc，因為我們測試下來，它的效能比 tcmalloc 略好。 jemalloc 還可以用於針對特定工作負載進行調優，如最大化 CPU 利用率。詳情可參考 jemalloc調優指南。

sudo apt-get install -y libjemalloc-dev
export LD_PRELOAD=$LD_PRELOAD:/usr/lib/x86_64-linux-gnu/libjemalloc.so
export MALLOC_CONF="oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms: 60000,muzzy_decay_ms:60000"

接下來，我們安裝 libiomp 庫來最佳化多核並行，這個庫是 英特爾 OpenMP 執行時庫 的一部分。

sudo apt-get install intel-mkl
export LD_PRELOAD=$LD_PRELOAD:/usr/lib/x86_64-linux-gnu/libiomp5.so
export OMP_NUM_THREADS=32

最後，我們安裝 numactl 命令列工具。它讓我們可以把我們的 Python 程式繫結到指定的核，並避免一些上下文切換開銷。

numactl -C 0-31 python sd_blog_1.py

使用這些最佳化後，原始的 Diffusers 程式碼只需 11.8 秒 就可以完成推理，快了幾乎 3 倍，而且無需任何程式碼更改。這些工具在我們的 32 核至強 CPU 上執行得相當不錯。

我們還有招。現在我們把 英特爾 PyTorch 擴充套件 (Intel Extension for PyTorch， IPEX) 引入進來。

IPEX 與 BF16

IPEX 擴充套件了 PyTorch 使之可以進一步充分利用英特爾 CPU 上的硬體加速功能，包括 AVX-512 、向量神經網路指令 (Vector Neural Network Instructions，AVX512 VNNI) 以及 先進矩陣擴充套件 (AMX)。

我們先安裝 IPEX。

pip install intel_extension_for_pytorch==1.13.100

裝好後，我們需要修改部分程式碼以將 IPEX 最佳化應用到 pipeline 的每個模組 (你可以透過列印 pipe 物件羅列出它有哪些模組)，其中之一的最佳化就是把資料格式轉換為 channels-last 格式。

import torch
import intel_extension_for_pytorch as ipex
...
pipe = StableDiffusionPipeline.from_pretrained(model_id)

# to channels last
pipe.unet = pipe.unet.to(memory_format=torch.channels_last)
pipe.vae = pipe.vae.to(memory_format=torch.channels_last)
pipe.text_encoder = pipe.text_encoder.to(memory_format=torch.channels_last)
pipe.safety_checker = pipe.safety_checker.to(memory_format=torch.channels_last)

# Create random input to enable JIT compilation
sample = torch.randn(2,4,64,64)
timestep = torch.rand(1)*999
encoder_hidden_status = torch.randn(2,77,768)
input_example = (sample, timestep, encoder_hidden_status)

# optimize with IPEX
pipe.unet = ipex.optimize(pipe.unet.eval(), dtype=torch.bfloat16, inplace=True, sample_input=input_example)
pipe.vae = ipex.optimize(pipe.vae.eval(), dtype=torch.bfloat16, inplace=True)
pipe.text_encoder = ipex.optimize(pipe.text_encoder.eval(), dtype=torch.bfloat16, inplace=True)
pipe.safety_checker = ipex.optimize(pipe.safety_checker.eval(), dtype=torch.bfloat16, inplace=True)

我們使用了 bloat16 資料型別，以利用 Sapphire Rapids CPU 上的 AMX 加速器。

with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    latency = elapsed_time(pipe, prompt)
    print(latency)

經過此番改動，推理延遲從 11.9 秒進一步減少到 5.4 秒。感謝 IPEX 和 AMX，推理速度提高了 2 倍以上。

還能榨點效能出來嗎？能，我們將目光轉向排程器 (scheduler)！

排程器

Diffusers 庫支援為每個 Stable Diffusion pipiline 配置 排程器 (scheduler)，用於在去噪速度和去噪質量之間找到最佳折衷。

根據文件所述: “ 截至本文件撰寫時，DPMSolverMultistepScheduler 能實現最佳的速度/質量權衡，只需 20 步即可執行。” 我們可以試一下 DPMSolverMultistepScheduler。

from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
...
dpm = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=dpm)

最終，推理延遲降至 5.05 秒。與我們最初的 Sapphire Rapids 基線 (32.3 秒) 相比，幾乎快了 6.5 倍！

執行環境: Amazon EC2 r7iz.metal-16xl, 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

總結

在幾秒鐘內生成高質量影像的能力可用於許多場景，如 2C 的應用程式、營銷和媒體領域的內容生成，或生成合成資料以擴充資料集。

如你想要在這方面起步，以下是一些有用的資源:

如果你有任何問題或反饋，請透過 Hugging Face 論壇 告訴我們。

感謝垂閱！

英文原文: https://hf.co/blog/stable-diffusion-inference-intel

作者: Julien Simon, Ella Charlaix

譯者: MatrixYao

審校/排版: zhongdongy (阿東)

相關文章