跳转到主要内容
博客

torch.compile 和 Diffusers：达到巅峰性能的实践指南

Diffusers 是首选的库，它为前沿和开放的图像、视频和音频扩散模型提供了一个统一的接口。在过去的几个月里，我们加深了它与torch.compile的集成。通过根据扩散模型的架构定制编译工作流，torch.compile在对用户体验影响最小的情况下，显著提升了速度。在本文中，我们将展示如何释放这些性能增益。本文的目标受众是

  • 扩散模型作者 – 学习一些小的代码更改，使您的模型对编译器友好，以便最终用户可以从性能提升中受益。
  • 扩散模型用户 – 了解编译时间与运行时间的权衡，如何避免不必要的重新编译，以及其他可以帮助您选择正确设置的方面。我们将通过Diffusers中两个备受社区喜爱的管道来展示其效果。

虽然示例位于Diffusers仓库中，但大多数原则也适用于其他深度学习工作负载。

目录

  • 背景
  • 为扩散模型有效使用 torch.compile
    • 普通编译
    • 模型作者必读：使用 fullgraph=True
    • 模型用户必读：使用区域编译
    • 模型用户必读：减少重新编译
  • 将torch.compile扩展到热门的 Diffusers 功能
    • 内存受限的 GPU
    • LoRA 适配器
  • 运行加固
  • 结论
  • 重要资源链接

背景

torch.compile在您充分了解模型并编译正确的子模块时，能发挥最大效用。在本节中，我们将首先概述影响torch.compile用户体验的因素，然后剖析扩散模型架构，以确定哪些组件从编译中受益最多。

torch.compile的性能和可用性核心因素

torch.compile将 Python 程序转换为优化的计算图，然后生成机器码，但其加速效果和易用性取决于以下因素：

编译延迟 – 作为 JIT（即时）编译器，torch.compile在首次运行时才会启动，因此用户会首先体验到编译成本。这个启动成本可能很高，尤其是在处理大型计算图时。

缓解方法：尝试区域编译，针对小而重复的区域进行编译。虽然这可能会限制与编译整个模型相比的最大可能加速，但它通常能在性能和编译时间之间取得更好的平衡，因此在决定前请评估权衡。

图中断 — 动态 Python 或不支持的操作会将 Python 程序分割成许多小图，从而大幅降低潜在的加速效果。模型开发者应力求使模型的计算密集型部分免于任何图中断。

缓解方法：开启fullgraph=True，识别中断点，并在准备模型时消除它们。

重新编译 – torch.compile会根据确切的输入形状专门生成代码，因此将分辨率从512 × 512更改为1024 × 1024会触发重新编译及其伴随的延迟。

缓解方法：启用dynamic=True来放宽形状约束。请注意，dynamic=True对扩散模型效果很好，但推荐的方法是使用mark_dynamic来有选择地对模型应用动态性。

设备到主机 (DtoH) 同步也可能妨碍最佳性能，但这些问题不容小觑，必须逐案处理。Diffusers 中最广泛使用的扩散管道没有这些同步。感兴趣的读者可以查看这份文档了解更多。由于这些同步与其他因素相比，对延迟的增加影响较小，因此本文的其余部分将不再关注它们。

扩散模型架构

我们将使用Flux‑1‑Dev，这是一个来自 Black Forest Labs 的开源文本到图像模型，作为我们的运行示例。一个扩散管道不是单一网络；它是一组模型的集合：

  • 文本编码器 – CLIP-Text 和 T5 将用户提示转换为嵌入。
  • 去噪器 – 扩散变换器 (DiT) 在这些嵌入的条件下，逐步优化一个带噪声的潜在表示。
  • 解码器 (VAE) – 将最终的潜在表示转换为 RGB 像素。

在这些组件中，DiT占据了大部分计算预算。你可以编译管道中的每个组件，但这只会增加编译延迟、重新编译次数和潜在的图中断，这些开销几乎无关紧要，因为这些部分已经只占总运行时间的一小部分。因此，我们将torch.compile的使用限制在 DiT 组件上，而不是整个管道。

为扩散模型有效使用 torch.compile

普通编译

让我们建立一个基准，以便在此基础上逐步改进，同时保持流畅的torch.compile用户体验。加载 Flux‑1‑Dev 检查点并按常规方式生成图像：

import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16
).to("cuda")

prompt = "A cat holding a sign that says hello world"
out = pipe(
    prompt=prompt,
    guidance_scale=3.5,
    num_inference_steps=28,
    max_sequence_length=512,

).images[0]

out.save("image.png")

现在编译计算密集型的扩散变换器子模块：

pipe.transformer.compile(fullgraph=True)

仅这一行代码就将 H100 上的延迟从 6.7 秒减少到 4.5 秒，实现了大约1.5 倍的加速，而且没有牺牲图像质量。在底层，torch.compile融合了内核并消除了 Python 开销，从而提高了内存和计算效率。

模型作者必读：使用 fullgraph=True

DiT 的前向传播在结构上很简单，所以我们期望它能形成一个连续的图。这个标志指示torch.compile在发生任何图中断时抛出错误，让您能够及早发现不支持的操作，而不是默默地放弃潜在的性能提升。我们建议扩散模型作者在模型准备阶段早期设置 fullgraph=True并修复图中断。请参阅torch.compile故障排除文档和手册来修复图中断。

模型用户必读：使用区域编译

如果您跟着操作，您会注意到第一次推理调用非常慢，在 H100 机器上需要 67.4 秒。这是编译开销。编译延迟随交给编译器的图的大小而增加。减少这种成本的一个实用方法是编译更小、重复的块，我们称之为区域编译策略。

DiT 本质上是一堆相同的 Transformer 层。如果我们编译一次层，并将其内核重用于每个后续层，我们就可以大幅减少编译时间，同时几乎保留了全图编译所见的所有运行时增益。

Diffusers 通过一个单行辅助函数来暴露此功能

pipe.transformer.compile_repeated_blocks(fullgraph=True)

在 H100 上，这将编译延迟从67.4 秒减少到 9.6 秒，将冷启动时间减少了7 倍，同时仍然提供了全模型编译所实现的 1.5 倍运行时加速。如果您想深入了解或为您的新模型启用区域编译，实现讨论位于 PR中。

请注意，上面的编译时间是冷启动测量值：我们使用torch._inductor.utils.fresh_inductor_cache API 清除了编译缓存，因此torch.compile必须从头开始。或者，在热启动中，缓存的编译器构件（存储在本地磁盘或远程缓存中）让编译器跳过部分编译过程，从而减少编译延迟。对于我们的模型，区域编译在冷启动时需要 9.6 秒，但一旦缓存变热，只需 2.4 秒。有关有效使用编译缓存的详细信息，请参阅链接的指南。

模型用户必读：减少重新编译

因为 torch.compile是一个即时编译器，它会根据所见输入的属性——包括形状、数据类型和设备（更多细节请参考这篇博客）来特化编译器构件。改变其中任何一个都会导致重新编译。虽然这在幕后自动发生，但这种重新编译会导致更高的编译时间成本，从而带来糟糕的用户体验。

如果您的应用程序需要处理多种图像尺寸或批处理形状，请在编译时传递 dynamic=True。对于通用模型，PyTorch 推荐使用 mark_dynamic，但是dynamic=True在扩散模型中效果很好。

pipe.transformer.compile_repeated_blocks(
    fullgraph=True, dynamic=True
)

我们对Flux DiT 在形状变化时的前向传播进行了全编译基准测试，并获得了令人信服的结果。

将 torch.compile扩展到热门的 Diffusers 功能

到现在，您应该对如何使用torch.compile加速扩散模型有了一个清晰的认识，而不会影响用户体验。接下来，我们将讨论两个社区最喜欢的 Diffusers 功能，并使它们与torch.compile完全兼容。我们将默认使用区域编译，因为它能提供与完全编译相同的加速效果，但编译延迟要小 8 倍。

  1. 内存受限的 GPU – 许多 Diffusers 用户在显存无法容纳整个模型的显卡上工作。我们将研究 CPU 卸载和量化，以确保在这些设备上生成图像的可行性。
  2. 使用 LoRA 适配器进行快速个性化 – 通过低秩适配器进行微调是使扩散模型适应新风格或任务的首选方法。我们将演示如何在不触发重新编译的情况下更换 LoRA。

内存受限的 GPU

CPU 卸载：一个完整的 bfloat16 格式的 Flux 管道大约消耗33GB的显存，这超出了大多数消费级 GPU 的承受能力。幸运的是，并非每个子模块都必须在整个前向传播过程中占用 GPU 内存。一旦文本编码器完成提示嵌入的生成，它们就可以被移至系统内存。同样，在 DiT 优化完潜在表示后，它可以将 GPU 内存让给 VAE 解码器。

Diffusers 用一行代码就实现了这种卸载

pipe.enable_model_cpu_offload()

峰值 GPU 使用量降至大约22.7 GB，使得在较小显卡上进行高分辨率生成成为可能，代价是额外的PCIe流量。卸载用内存换取时间，端到端运行现在需要大约21.5 秒，而不是 6.7 秒。

您可以通过同时启用torch.compile和卸载来弥补部分时间损失。编译器的内核融合抵消了一部分 PCIe 开销，将延迟缩短至大约18.7 秒，同时保持了 22.6 GB 的较小内存占用。

pipe.enable_model_cpu_offload()

pipe.transformer.compile_repeated_blocks(fullgraph=True)

Diffusers 提供了多种卸载模式，每种模式都有其独特的“速度与内存”平衡点。请查看卸载指南以获取完整的选项列表。

量化：CPU 卸载可以释放 GPU 内存，但它仍然假设最大的组件，即 DiT，能够装入 GPU 内存。缓解内存压力的另一种方法是利用权重 量化，前提是可以容忍一定程度的输出损失。

Diffusers 支持多种量化后端；这里我们使用来自bitsandbytes的 4 位NF4量化。它将 DiT 的权重占用减少了大约一半，将峰值 GPU 内存从大约 33 GB 降低到15 GB，同时保持了图像质量。

与 CPU 卸载相比，权重 量化将权重保留在 GPU 内存中，导致运行时损失较小——从基准的 6.7 秒增加到 7.3 秒。在此基础上添加torch.compile会融合 4 位操作，将推理时间从7.279 秒减少到 5.048 秒，实现了大约 1.5 倍的加速。

您可以在此处找到不同的后端和代码指针。

（我们对 DiT 和 T5 启用了量化，因为它们都比 CLIP 和 VAE 消耗更多的内存。）

量化 + 卸载：正如您可能预期的那样，您可以将 NF4 量化与 CPU 卸载相结合，以获得最大的内存效益。这种组合技术将内存占用减少到 12.2 GB，推理时间为 12.2 秒。无缝应用 torch.compile，将运行时间减少到9.8 秒，实现了 1.24 倍的加速。

基准测试是使用这个脚本在单台 H100 机器上进行的。以下是我们到目前为止讨论的数据摘要。灰色框表示基准数据，绿色框表示最佳配置。

仔细观察上表，我们可以立即看出，区域编译提供的加速几乎与完全编译相当，而在编译时间方面则显著更快。

LoRA 适配器

LoRA 适配器让你可以在不微调数百万参数的情况下，对基础扩散模型进行个性化。缺点是，切换适配器通常会交换 DiT 内部的权重张量，迫使 torch.compile重新追踪和重新编译。

Diffusers 现在集成了PEFT的 LoRA 热插拔功能来避免这个问题。你只需声明你将需要的最大LoRA 秩，编译一次，然后就可以动态切换适配器。无需重新编译。

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16
).to("cuda")

pipe.enable_lora_hotswap(target_rank=max_rank)
pipe.load_lora_weights(<lora-adapter-name1>)
pipe.transformer.compile(fullgraph=True)

# from this point on, load new LoRAs with `hotswap=True`
pipe.load_lora_weights(<lora-adapter-name2>, hotswap=True)

因为只有LoRA 权重张量发生变化，而它们的形状保持不变，所以编译好的内核仍然有效，推理延迟也保持平稳。

注意事项

  • 我们需要提前提供所有 LoRA 适配器中的最大秩。因此，如果我们有一个秩为 16 的适配器和另一个秩为 32 的适配器，我们需要传递 `max_rank=32`。
  • 热插拔的 LoRA 适配器只能针对与第一个 LoRA 相同的层，或其子集。
  • 尚不支持热插拔文本编码器。

更多详情，请参阅LoRA 热插拔文档和测试套件。

LoRA 推理与上面讨论的卸载和量化功能无缝集成。如果您的 GPU VRAM 受限，请考虑使用它们。

运行加固

Diffusers 每晚运行一个专门的 CI 作业，以保持 torch.compile支持的稳健性。该套件会测试最流行的管道，并自动检查：

鼓励感兴趣的读者查看此链接，以了解近期改进torch.compile覆盖范围的 PR。

基准测试

正确性只是一半；我们同样关心性能。一个经过改进的基准测试工作流现在与 CI 一同运行，捕获本文中涵盖的每种场景的延迟和峰值内存。结果会导出到一个整合的 CSV 文件中，以便轻松发现性能回归（或提升！）。设计和早期数据位于这个PR中。

结论

torch.compile可以将标准的 Diffusers 管道转变为高性能、内存高效的主力。通过将编译重点放在 DiT 上，利用区域编译和动态形状，并将编译器与卸载、量化和 LoRA 热插拔相结合，您可以在不牺牲图像质量或灵活性的情况下，释放显著的速度提升和 VRAM 节省。

我们希望这些方法能激励您将 torch.compile融入到您自己的扩散工作流中。我们期待看到您接下来的创作。

编译愉快 ⚡️

Diffusers 团队感谢Animesh和Ryan在改进torch.compile支持方面提供的帮助和支持。

重要资源链接