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

最近，我们介绍了最新一代的 英特尔至强 CPU（代号 Sapphire Rapids），及其用于深度学习加速的新硬件特性，以及如何使用它们来加速自然语言处理 Transformers 的分布式微调和推理。

在本文中，我们将向您展示在 Sapphire Rapids CPU 上加速 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 模型生成图像变得极其简单。如果您不熟悉这些模型，这里有一篇很棒的图文并茂的介绍。

首先，让我们创建一个包含所需库的虚拟环境：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 秒。正如这个 Intel Space 所演示的，同样的代码在上一代英特尔至强 CPU（代号 Ice Lake）上运行大约需要 45 秒。

开箱即用，我们可以看到，在没有任何代码更改的情况下，Sapphire Rapids CPU 的速度要快得多！

现在，让我们来加速吧！

Optimum Intel 与 OpenVINO

Optimum Intel 可在英特尔架构上加速端到端流水线。它的 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 倍加速。

上述流水线支持动态输入形状，对图像数量或其分辨率没有限制。对于 Stable Diffusion，您的应用程序通常仅限于一种（或几种）不同的输出分辨率，例如 512x512 或 256x256。因此，通过将流水线重塑为固定分辨率来解锁显著的加速是非常有意义的。如果您需要多种输出分辨率，可以简单地维护几个流水线实例，每种分辨率一个。

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 至强 CPU 上的原生推理相比，它提供了近 10 倍的加速。

如果您不能或不想使用 OpenVINO，本文的其余部分将向您展示一系列其他优化技术。请系好安全带！

系统级优化

Diffuser 模型是数 GB 大的模型，图像生成是内存密集型操作。通过安装一个高性能的内存分配库，我们应该能够加速内存操作，并在至强核心之间并行化它们。请注意，这将更改您系统上的默认内存分配库。当然，您可以通过卸载新库来恢复到默认库。

jemalloc 和 tcmalloc 同样值得关注。在这里，我安装了 jemalloc，因为我的测试表明它有轻微的性能优势。它还可以针对特定工作负载进行调整，例如最大化 CPU 利用率。您可以参考调优指南了解详情。

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 库以优化并行处理。它是 Intel 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 上确实表现出色。

我们还远未完成。让我们将 Intel Extension for PyTorch 也加入进来。

IPEX 与 BF16

Intel Extension for Pytorch (IPEX) 扩展了 PyTorch，并利用了英特尔 CPU 上的硬件加速功能，例如 AVX-512 矢量神经网络指令 (AVX512 VNNI) 和高级矩阵扩展 (AMX)。

让我们安装它。

pip install intel_extension_for_pytorch==1.13.100

然后我们更新代码，用 IPEX 优化每个流水线元素（您可以通过打印 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)

我们还启用了 bfloat16 数据格式，以利用 Sapphire Rapids CPU 上的 AMX 瓦片矩阵乘法单元 (TMMU) 加速器。

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

通过这个更新版本，推理延迟从 11.9 秒进一步减少到 5.4 秒。得益于 IPEX 和 AMX，这实现了超过 2 倍的加速。

我们还能榨取更多性能吗？可以，用调度器！

调度器

Diffusers 库允许我们为 Stable Diffusion 流水线附加一个调度器。调度器试图在去噪速度和去噪质量之间找到最佳平衡。

根据文档：“在撰写本文档时，DPMSolverMultistepScheduler 提供了可以说是最佳的速度/质量权衡，并且只需 20 个步骤即可运行。”

让我们试试看。

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*

结论

在几秒钟内生成高质量图像的能力应该适用于许多用例，例如客户应用、营销和媒体内容生成，或用于数据集增强的合成数据。

以下是一些资源可帮助您入门：

• Diffusers 文档
• Optimum Intel 文档
• Intel IPEX 在 GitHub 上
• Intel 和 Hugging Face 提供的开发人员资源。

如果您有任何问题或反馈，我们非常乐意在 Hugging Face 论坛上阅读。

感谢阅读！