加速Intel CPU上的稳定扩散推理

'Accelerate stable inference on Intel CPUs'

最近,我们介绍了最新一代的英特尔 Xeon 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秒。正如这个英特尔 Space 所示,同样的代码在之前一代的英特尔 Xeon(代号 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。在加载模型时,如果将 export=True,可以动态地将 PyTorch 模型转换为 OpenVINO 格式。

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,您的应用程序通常会限制为一个(或少数几个)不同的输出分辨率,例如 512×512 或 256×256。因此,通过将流水线重新调整为固定分辨率,解锁显著加速是非常有意义的。如果您需要多个输出分辨率,可以简单地维护几个流水线实例,每个分辨率一个。

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推理的一种简单高效的方式。与Ice Lake Xeons上的原始推理相比,与Sapphire Rapids CPU结合使用可提供近10倍的加速效果。

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

系统级优化

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

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* Runtime的一部分。

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核Xeon上表现非常出色。

我们还远未完成。让我们将Intel Extension for PyTorch添加到混合中。

IPEX和BF16

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

让我们安装它。

pip install intel_extension_for_pytorch==1.13.100

然后,我们更新我们的代码,使用IPEX优化每个流水线元素(您可以通过打印pipe对象进行列举)。这需要将它们转换为通道最后的格式。

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处理器上的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论坛上阅读。

感谢阅读!