#
加速推理

Diffusion模型在推理时速度较慢，因为生成是一个迭代过程，需要经过一定数量的”步数”逐步将噪声细化为图像或视频。要加速这一过程，您可以尝试使用不同的调度器、降低模型权重的精度以加快计算、使用更高效的内存注意力机制等方法。

将这些技术组合使用，可以比单独使用任何一种技术获得更快的推理速度。

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

##
模型数据类型

模型权重的精度和数据类型会影响推理速度，因为更高的精度需要更多内存来加载，也需要更多时间进行计算。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）实现了多种注意力后端，包括FlashAttention、xFormers和原生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]

第一次编译时速度较慢，但一旦编译完成，速度会显著提升。尽量只在相同类型的推理操作上使用编译后的管道。在不同尺寸的图像上调用编译后的管道会重新触发编译，这会很慢且效率低下。

###
动态形状编译

确保始终使用PyTorch的nightly版本以获得更好的支持。

`torch.compile`会跟踪输入形状和条件，如果这些不同，它会重新编译模型。例如，如果模型是在1024x1024分辨率的图像上编译的，而在不同分辨率的图像上使用，就会触发重新编译。

为避免重新编译，添加`dynamic=True`以尝试生成更动态的内核，避免条件变化时重新编译。

    + torch.fx.experimental._config.use_duck_shape = False
    + pipeline.unet = torch.compile(
        pipeline.unet, fullgraph=True, dynamic=True
    )

指定`use_duck_shape=False`会指示编译器是否应使用相同的符号变量来表示相同大小的输入。更多详情请参阅此评论。

并非所有模型都能开箱即用地从动态编译中受益，可能需要更改。参考此PR，它改进了`AuraFlowPipeline`的实现以受益于动态编译。

如果动态编译对Diffusers模型的效果不如预期，请随时提出问题。

###
区域编译

区域编译通过仅编译模型中小而频繁重复的块（通常是transformer层）来减少冷启动延迟，并为每个后续出现的块重用编译后的工件。对于许多diffusion架构，这提供了与全图编译相同的运行时加速，并将编译时间减少了8-10倍。

使用`compile_repeated_blocks()`方法（一个包装`torch.compile`的辅助函数）在任何组件（如transformer模型）上，如下所示。

    # pip install -U diffusers
    import torch
    from diffusers import StableDiffusionXLPipeline

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

    # 仅编译UNet中重复的transformer层
    pipeline.unet.compile_repeated_blocks(fullgraph=True)

要为新模型启用区域编译，请在模型类中添加一个`_repeated_blocks`属性，包含您想要编译的块的类名（作为字符串）。

    class MyUNet(ModelMixin):
        _repeated_blocks = ("Transformer2DModel",)  # ← 默认编译

更多区域编译示例，请参阅参考PR。

Accelerate中还有一个compile_regions方法，可以自动选择模型中的候选块进行编译。其余图会单独编译。这对于快速实验很有用，因为您不需要设置哪些块要编译或调整编译标志。

    # 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)

`compile_repeated_blocks()`是故意显式的。在`_repeated_blocks`中列出要重复的块，辅助函数仅编译这些块。它提供了可预测的行为，并且只需一行代码即可轻松推理缓存重用。

###
图中断

在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上，以避免通信同步和延迟。

参阅torch.compile和Diffusers：峰值性能实践指南博客文章，了解如何为扩散模型最大化`torch.compile`的性能。

###
基准测试

参阅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]

##
融合投影矩阵

> [!WARNING][fuse_qkv_projections](https://github.com/huggingface/diffusers/blob/58431f102cf39c3c8a569f32d71b2ea8caa461e1/src/diffusers/pipelines/pipeline_utils.py#L2034)方法是实验性的，目前主要支持Stable Diffusion管道。参阅此PR了解如何为其他管道启用它。

在注意力块中，输入被投影到三个子空间，分别由投影矩阵Q、K和V表示。这些投影通常单独计算，但您可以水平组合这些矩阵为一个矩阵，并在单步中执行投影。这会增加输入投影的矩阵乘法大小，并提高量化的效果。

    pipeline.fuse_qkv_projections()

##
资源

- 阅读Presenting Flux Fast: Making Flux go brrr on H100s博客文章，了解如何结合所有这些优化与TorchInductor和AOTInductor，使用flux-fast的配方获得约2.5倍的加速。

  这些配方支持AMD硬件和Flux.1 Kontext Dev。

- 阅读torch.compile和Diffusers：峰值性能实践指南博客文章，了解如何在使用`torch.compile`时最大化性能。

< > Update on GitHub