使用 NNCF 和 🤗 Optimum 优化 Intel CPU 上的 稳定扩散

概览

Hugging Face 提出了一种方法，通过将 Neural Network Compression Framework (NNCF) 量化感知训练与 令牌合并 相结合，在 Intel CPU 上优化 稳定扩散，推理速度最高可提升 5.1 倍，模型大小降至原始 PyTorch 检查点的 0.25 倍。

稳定扩散优化

稳定扩散管道的 UNet 组件是计算成本最高的部分，因此对其进行优化可以带来显著的速度提升。传统的训练后 8 位量化在此模型上效果不佳，因为像素级预测任务对参数变化高度敏感，且由于在数亿样本上训练，模型冗余较少。为了保持精度，需要采用更复杂的量化方法，例如 Quantization‑Aware Training (QAT)。

优化工作流程

我们从在 Pokemon 数据集 (svjack/Stable-Diffusion-Pokemon-en) 上微调的 稳定扩散 模型开始。利用 Diffusers 文本到图像微调示例，我们将基于 NNCF 的 QAT 集成到训练脚本中，添加了知识蒸馏损失（其中原始模型充当教师），并对模型参数（不包括量化器）应用了指数移动平均 (EMA) 以提高训练稳定性。梯度检查点和将 EMA 模型保存在 RAM 中使得整个优化过程能够在单个配备 24 GB VRAM 的 GPU 上运行，在不到一天的时间内完成 4096 次迭代。

超越量化感知训练

仅使用量化可以减小模型占用空间、加载时间、内存消耗和延迟。我们将 8 位量化与 令牌合并 (ToME) 方法结合，该方法在自注意力块之前合并冗余令牌以减少计算。上述组合工作流程包括知识蒸馏、EMA 和梯度检查点。再次从在 Pokemon 上微调的模型出发，我们在量化基础上应用了合并比例为 0.4 的 ToME。得到的模型旨在用于客户端或边缘 CPU 的推理。

结果

将 PyTorch 基线转换为 OpenVINO FP32 可获得 1.9 倍的加速。添加 8 位量化后，相对于 PyTorch 的加速提升至 3.9 倍，并且模型占用空间降至原始检查点的 0.25 倍。在量化基础上叠加 令牌合并 可实现 5.1 倍的推理速度提升，同时保持占用空间仍为 0.25 倍。所有测量均在使用 3rd Generation Intel® Xeon® Scalable 处理器并搭载 Intel® Deep Learning Boost 技术的 Hugging Face Spaces CPU 升级实例上，使用 OpenVINO 2022.3 以及默认的 50 步推理步骤完成。博客指出，减少步骤可以提升速度但可能影响图像质量，并建议尝试不同的步数和调度器。

An example inference pipeline is shown below:

from optimum.intel import OVStableDiffusionPipeline

# Load and compile the pipeline for performance.
name = "OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive"
pipe = OVStableDiffusionPipeline.from_pretrained(name, compile=False)
pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
pipe.compile()

# Generate an image.
prompt = "a drawing of a green pokemon with red eyes"
output = pipe(prompt, num_inference_steps=50, output_type="pil\)).images[0]
output.save("image.png

The training and quantization code are available in the Optimum Intel repository, a demonstration notebook is provided, and optimized models can be found on the Hugging Face Hub under the OpenVINO organization. A live demo runs on Hugging Face Spaces.

关于通用目的的稳定扩散模型？

在 Pokemon 模型上演示的工作流程表明，在适度的训练资源下可以实现显著的优化。从头开始训练通用目的的 稳定扩散 模型成本高昂，但只要有足够的预算和硬件，就可以采用相同的方法。需要注意的是，令牌合并会降低模型容量；因此，对于更复杂的数据集，在优化过程中应使用较低的合并比例。

For further reading on complementary approaches for 4th‑generation Intel Xeon CPUs, see the related Hugging Face blog post on Stable Diffusion inference with Intel.

Sources