Diffusers 文档

加速推理

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

加速推理

扩散模型在推理时速度较慢,因为生成是一个迭代过程,其中噪声在一定数量的“步”中逐渐细化为图像或视频。为了加快此过程,您可以尝试使用不同的调度器,降低模型权重的精度以加快计算,使用更节省内存的注意力机制等。

结合使用这些技术,推理速度将比单独使用任何一种技术都要快。

本指南将介绍如何加速推理。

模型数据类型

模型权重的精度和数据类型会影响推理速度,因为更高的精度需要更多的内存来加载和更多的时间来执行计算。PyTorch 默认以 float32 或全精度加载模型权重,因此更改数据类型是快速获得更快推理的简单方法。

bfloat16
float16
TensorFloat-32

bfloat16 类似于 float16,但对数值误差更鲁棒。bfloat16 的硬件支持各不相同,但大多数现代 GPU 都能够支持 bfloat16。

import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

缩放点积注意力

内存高效注意力机制优化了推理速度*和*内存使用

缩放点积注意力(SDPA)实现了多种注意力后端,包括FlashAttentionxFormers和一个原生的 C++ 实现。它会自动为您的硬件选择最优的后端。

如果您使用 PyTorch >= 2.0,SDPA 默认启用,无需对代码进行额外更改。不过,如果您想选择自己的注意力后端,也可以尝试其他后端。下面的示例使用 torch.nn.attention.sdpa_kernel 上下文管理器来启用高效注意力。

from torch.nn.attention import SDPBackend, sdpa_kernel
import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
  image = pipeline(prompt, num_inference_steps=30).images[0]

torch.compile

torch.compile 通过将 PyTorch 代码和操作编译成优化过的内核来加速推理。Diffusers 通常会编译计算密集型模型,如 UNet、Transformer 或 VAE。

启用以下编译器设置以获得最大速度(更多选项请参阅完整列表)。

import torch
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

加载并编译 UNet 和 VAE。您可以选择多种不同的模式,但 `"max-autotune"` 通过编译到 CUDA 图来优化最快速度。CUDA 图通过单个 CPU 操作启动多个 GPU 操作,有效降低了开销。

使用 PyTorch 2.3.1,您可以控制 torch.compile 的缓存行为。这对于像 `"max-autotune"` 这样的编译模式特别有益,它会对多个编译标志进行网格搜索以找到最佳配置。在torch.compile 中的编译时间缓存教程中了解更多信息。

将内存布局更改为channels_last也可以优化内存和推理速度。

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
).to("cuda")
pipeline.unet.to(memory_format=torch.channels_last)
pipeline.vae.to(memory_format=torch.channels_last)
pipeline.unet = torch.compile(
    pipeline.unet, mode="max-autotune", fullgraph=True
)
pipeline.vae.decode = torch.compile(
    pipeline.vae.decode,
    mode="max-autotune",
    fullgraph=True
)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

第一次编译速度较慢,但一旦编译完成,速度会显著加快。尝试仅在相同类型的推理操作上使用编译后的流水线。在不同图像尺寸上调用编译后的流水线会重新触发编译,这既慢又低效。

区域编译

区域编译通过仅编译模型的特定重复区域(或块)而不是整个模型来减少冷启动编译时间。编译器会重用其他块的缓存和编译代码。

Accelerate 提供了 compile_regions 方法,用于自动顺序编译 `nn.Module` 的重复块。模型的其余部分则单独编译。

# pip install -U accelerate
import torch
from diffusers import StableDiffusionXLPipeline
from accelerate.utils import compile regions

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
).to("cuda")
pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph=True)

图断裂

重要的是在 torch.compile 中指定 `fullgraph=True`,以确保底层模型中没有图断裂。这允许您利用 torch.compile 的优势,而不会有任何性能下降。对于 UNet 和 VAE,这会改变您访问返回变量的方式。

- latents = unet(
-   latents, timestep=timestep, encoder_hidden_states=prompt_embeds
-).sample

+ latents = unet(
+   latents, timestep=timestep, encoder_hidden_states=prompt_embeds, return_dict=False
+)[0]

GPU 同步

在去噪器做出预测后,调度器会每次调用 `step()` 函数,并对 `sigmas` 变量进行索引。当放置在 GPU 上时,由于 CPU 和 GPU 之间的通信同步,它会引入延迟。当去噪器已被编译时,这会变得更加明显。

一般来说,`sigmas` 应该保留在 CPU 上,以避免通信同步和延迟。

基准测试

请参阅diffusers/benchmarks数据集,查看编译后流水线的推理延迟和内存使用数据。

diffusers-torchao 存储库还包含 Flux 和 CogVideoX 编译版本的基准测试结果。

动态量化

动态量化通过降低精度以实现更快的数学运算来提高推理速度。这种特殊类型的量化会根据运行时的数据而不是使用固定缩放因子来确定如何缩放激活。因此,缩放因子与数据更准确地对齐。

以下示例使用 torchao 库对 UNet 和 VAE 应用 动态 int8 量化

请参阅我们的torchao文档,了解有关如何使用 Diffusers torchao 集成的更多信息。

配置编译器标签以获得最大速度。

import torch
from torchao import apply_dynamic_quant
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True
torch._inductor.config.force_fuse_int_mm_with_mul = True
torch._inductor.config.use_mixed_mm = True

使用 dynamic_quant_filter_fn 过滤掉 UNet 和 VAE 中一些不从动态量化中受益的线性层。

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")

apply_dynamic_quant(pipeline.unet, dynamic_quant_filter_fn)
apply_dynamic_quant(pipeline.vae, dynamic_quant_filter_fn)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

融合投影矩阵

fuse_qkv_projections 方法是实验性的,目前仅限于 Stable Diffusion 流水线。请查看此 PR 以了解如何在其他流水线中启用它。

在注意力块中,输入被投影到三个子空间,由投影矩阵 Q、K 和 V 表示。这些投影通常是分开计算的,但您可以将它们水平合并为一个矩阵,并一步执行投影。这增加了输入投影的矩阵乘法的大小,也提高了量化的影响。

pipeline.fuse_qkv_projections()
< > 在 GitHub 上更新