# 加速推理 扩散模型在推理阶段速度较慢，因为生成是一个迭代过程，需要在若干“步”中逐步将噪声细化为图像或视频。为了加速这一过程，您可以尝试使用不同的 [调度器](../api/schedulers/overview)、降低模型权重的精度以加快计算、采用更节省显存的注意力机制等。 将这些技术组合使用，可以实现比单独使用任何一项技术更快的推理速度。 本指南将介绍如何加速推理。 ## 模型数据类型 模型权重的精度和数据类型会影响推理速度，因为更高的精度需要更多的内存来加载，并且需要更多的时间来进行计算。PyTorch 默认以 float32 或全精度加载模型权重，因此更改数据类型是一种快速提升推理速度的简单方法。 bfloat16 与 float16 类似，但对数值误差的鲁棒性更强。bfloat16 的硬件支持情况因设备而异，但大多数现代 GPU 均支持 bfloat16。 ```py import torch from diffusers import StableDiffusionXLPipeline pipeline = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16 ).to("cuda") # or "mps", "xpu", "cpu" prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k" pipeline(prompt, num_inference_steps=30).images[0] ``` float16 与 bfloat16 类似，但更容易出现数值误差。 ```py import torch from diffusers import StableDiffusionXLPipeline pipeline = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16 ).to("cuda") # or "mps", "xpu", "cpu" prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k" pipeline(prompt, num_inference_steps=30).images[0] ``` [TensorFloat-32 (tf32)](https://blogs.nvidia.com/blog/2020/05/14/tensorfloat-32-precision-format/) 模式在 NVIDIA Ampere GPU 上受支持，它使用 tf32 执行卷积和矩阵乘法运算。存储和其他操作仍保持为 float32。与 bfloat16 或 float16 结合使用时，可显著提升计算速度。 PyTorch 默认仅对卷积启用 tf32 模式，你需要显式地为矩阵乘法启用该模式。 ```py import torch from diffusers import StableDiffusionXLPipeline torch.backends.cuda.matmul.allow_tf32 = True pipeline = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16 ).to("cuda") # or "mps", "xpu", "cpu" prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k" pipeline(prompt, num_inference_steps=30).images[0] ``` 有关更多详细信息，请参阅[混合精度训练](https://huggingface.co/docs/transformers/en/perf_train_gpu_one#mixed-precision)文档。 ## 缩放点积注意力 > [!TIP] > 内存高效注意力机制针对推理速度 *和* [内存占用](./memory#memory-efficient-attention) 进行了优化！ [缩放点积注意力（SDPA）](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)实现了多种注意力后端，包括 [FlashAttention](https://github.com/Dao-AILab/flash-attention)、[xFormers](https://github.com/facebookresearch/xformers) 以及原生 C++ 实现。它会自动为您的硬件选择最优的后端。 如果您使用的是 PyTorch >= 2.0，SDPA 默认已启用，且无需对代码进行任何额外更改。不过，如果您希望自行选择，可以尝试使用其他 attention 后端。下面的示例使用 [torch.nn.attention.sdpa_kernel](https://pytorch.org/docs/stable/generated/torch.nn.attention.sdpa_kernel.html) 上下文管理器来启用高效的 attention。 ```py 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", dtype=torch.bfloat16 ).to("cuda") # or "mps", "xpu", "cpu" 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](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) 通过将 PyTorch 代码和运算编译为优化后的内核来加速推理。Diffusers 通常会对计算量较大的模型（如 UNet、transformer 或 VAE）进行编译。 启用以下编译器设置以获得最大速度（有关更多选项，请参阅[完整列表](https://github.com/pytorch/pytorch/blob/main/torch/_inductor/config.py)）。 ```py 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 graph 来优化至最快速度。CUDA graph 通过单次 CPU 操作启动多个 GPU 操作，从而有效降低开销。 > [!TIP] > 在 PyTorch 2.3.1 中，您可以控制 torch.compile 的缓存行为。这对于 `"max-autotune"` 等编译模式尤其有益，该模式会对多个编译标志执行网格搜索，以找到最优配置。有关详细信息，请参阅 [torch.compile 中的编译时缓存](https://pytorch.org/tutorials/recipes/torch_compile_caching_tutorial.html) 教程。 将内存布局更改为 [channels_last](./memory#torchchannelslast) 还可以优化内存和推理速度。 ```py pipeline = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16 ).to("cuda") # or "mps", "xpu", "cpu" 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] ``` 首次编译速度较慢，但编译完成后速度会显著提升。请尽量仅将编译后的 pipeline 用于同一类型的推理操作。如果对不同的图像尺寸调用编译后的 pipeline，会重新触发编译，这既缓慢又低效。 ### 动态形状编译 > [!TIP] > 请务必始终使用 PyTorch 的 nightly 版本，以获得更好的支持。 `torch.compile` 会跟踪输入的形状和条件，如果这些条件发生变化，它会重新编译模型。例如，如果模型是在 1024x1024 分辨率的图像上编译的，而随后用于不同分辨率的图像，则会触发重新编译。 为避免重新编译，请添加 `dynamic=True`，以尝试生成更具动态性的内核，从而在条件变化时避免重新编译。 ```diff + torch.fx.experimental._config.use_duck_shape = False + pipeline.unet = torch.compile( pipeline.unet, fullgraph=True, dynamic=True ) ``` 指定 `use_duck_shape=False` 可指示编译器是否应使用相同的符号变量来表示大小相同的输入。有关更多详细信息，请参阅此 [评论](https://github.com/huggingface/diffusers/pull/11327#discussion_r2047659790)。 并非所有模型都能开箱即用地从动态编译中获益，可能需要对其进行修改。请参阅此 [PR](https://github.com/huggingface/diffusers/pull/11297/)，该 PR 改进了 [AuraFlowPipeline](/docs/diffusers/v0.43.0/zh/api/pipelines/aura_flow#diffusers.AuraFlowPipeline) 的实现，使其能够受益于动态编译。 如果动态编译在 Diffusers 模型上未按预期工作，请随时提交 issue。 ### 区域编译 [区域编译](https://docs.pytorch.org/tutorials/recipes/regional_compilation.html)通过仅编译模型中*较小且频繁重复的块*（通常是 Transformer 层）来降低冷启动延迟，并允许在后续每次出现时复用编译产物。对于许多扩散架构而言，这能提供与全图编译相同的运行时加速效果，并将编译时间缩短 8 到 10 倍。 使用 [compile_repeated_blocks()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.compile_repeated_blocks) 方法（一个封装了 `torch.compile` 的辅助函数）对诸如 transformer 模型等任意组件进行编译，如下所示。 ```py # pip install -U diffusers import torch from diffusers import StableDiffusionXLPipeline pipeline = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16, ).to("cuda") # or "mps", "xpu", "cpu" # compile only the repeated transformer layers inside the UNet pipeline.unet.compile_repeated_blocks(fullgraph=True) ``` 要为新模型启用区域编译，请在模型类中添加一个 `_repeated_blocks` 属性，其中包含要编译的块的类名（以字符串形式）。 ```py class MyUNet(ModelMixin): _repeated_blocks = ("Transformer2DModel",) # ← compiled by default ``` > [!TIP] > 有关更多区域编译示例，请参阅参考 [PR](https://github.com/huggingface/diffusers/pull/11705)。 [Accelerate](https://huggingface.co/docs/accelerate/index) 中还提供了一个 [compile_regions](https://github.com/huggingface/accelerate/blob/273799c85d849a1954a4f2e65767216eb37fa089/src/accelerate/utils/other.py#L78) 方法，用于自动选择模型中需要编译的候选块。剩余的图会被单独编译。由于无需手动设置要编译的块或调整编译标志，该方法非常适合快速实验。 ```py # 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", dtype=torch.float16 ).to("cuda") # or "mps", "xpu", "cpu" pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph=True) ``` [compile_repeated_blocks()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.compile_repeated_blocks) 的设计是刻意保持显式的。在 `_repeated_blocks` 中列出需要重复的块，该辅助函数只会编译这些块。它通过一行代码即可提供可预测的行为，并便于推理缓存的复用。 ### 图断点 在 torch.compile 中指定 `fullgraph=True` 非常重要，以确保底层模型中不存在图中断。这样，您就可以充分利用 torch.compile，而不会导致任何性能下降。对于 UNet 和 VAE，这会改变您访问返回变量的方式。 ```diff - 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()` 函数会在去噪器每次做出预测后被[调用](https://github.com/huggingface/diffusers/blob/1d686bac8146037e97f3fd8c56e4063230f71751/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py#L1228)，并且会对 `sigmas` 变量进行[索引](https://github.com/huggingface/diffusers/blob/1d686bac8146037e97f3fd8c56e4063230f71751/src/diffusers/schedulers/scheduling_euler_discrete.py#L476)。当将其放置在 GPU 上时，由于 CPU 与 GPU 之间的通信同步，会引入延迟。当去噪器已经被编译后，这一现象会更加明显。 通常，`sigmas` 应[保留在 CPU 上](https://github.com/huggingface/diffusers/blob/35a969d297cba69110d175ee79c59312b9f49e1e/src/diffusers/schedulers/scheduling_euler_discrete.py#L240)，以避免通信同步和延迟。 > [!TIP] > 请参阅 [torch.compile 与 Diffusers：性能巅峰实战指南](https://pytorch.org/blog/torch-compile-and-diffusers-a-hands-on-guide-to-peak-performance/) 博客文章，了解如何使用 `torch.compile` 最大化扩散模型的性能。 ### 基准测试 请参阅 [diffusers/benchmarks](https://huggingface.co/datasets/diffusers/benchmarks) 数据集，查看编译后流水线的推理延迟和内存使用数据。 [diffusers-torchao](https://github.com/sayakpaul/diffusers-torchao#benchmarking-results) 仓库还包含 Flux 和 CogVideoX 编译版本的基准测试结果。 ## 内核 [Kernels](https://huggingface.co/docs/kernels/index) 是一个用于构建、分发和加载 [Hub](https://huggingface.co/kernels-community) 上优化计算内核的库。它支持 [attention](./attention_backends#setattentionbackend) 内核以及用于 RMSNorm、GEGLU、RoPE 和 AdaLN 等操作的自定义 CUDA 内核。 [Diffusers Pipeline Integration](https://github.com/huggingface/kernels/blob/main/kernel-builder/skills/cuda-kernels/references/diffusers-integration.md) 指南展示了如何将内核与 [add cuda-kernels](https://github.com/huggingface/kernels/blob/main/kernel-builder/skills/cuda-kernels/SKILL.md) 技能进行集成。该技能可让 Claude 或 Codex 等智能体针对特定模型和您的硬件编写自定义内核。 > [!TIP] > 安装 [add cuda-kernels](https://github.com/huggingface/kernels/blob/main/kernel-builder/skills/cuda-kernels/SKILL.md) 技能，以教会智能体如何编写内核。[Codex 和 Claude 的通用自定义内核](https://huggingface.co/blog/custom-cuda-kernels-agent-skills) 博客文章对此进行了更详细的介绍。 例如，一个自定义的 RMSNorm 内核（由 `add cuda-kernels` 技能生成）配合 [torch.compile](#torchcompile) 使用，可在 H100 上将 LTX-Video 的生成速度提升 1.43 倍。