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 在第一次运行时启动，因此用户需要预先承担编译成本。这个启动成本可能很高，尤其是对于大型图。
缓解措施： 尝试使用区域编译来针对小的、重复的区域。虽然这可能会限制与编译整个模型相比可能达到的最大速度提升，但它通常能在性能和编译时间之间取得更好的平衡，因此在做决定前请权衡利弊。
图中断 (Graph breaks) — 动态 Python 或不受支持的操作会将 Python 程序切分成许多小的图，从而削弱了潜在的速度提升。模型开发者应努力使模型中计算密集的部分不包含任何图中断。
缓解措施： 开启 fullgraph=True，在准备模型时识别并消除中断。
重新编译 – torch.compile 会针对确切的输入形状专门化其代码，因此将分辨率从 512 × 512 更改为 1024 × 1024 会触发重新编译和随之而来的延迟。
缓解措施： 启用 dynamic=True 来放宽形状约束。请注意，dynamic=True 在扩散模型中效果很好，但推荐的方式是使用 mark_dynamic 选择性地对模型应用动态性。
设备到主机 (DtoH) 同步 也可能阻碍最佳性能， 但这些是非同小可的，需要根据具体情况处理。Diffusers 中使用最广泛的扩散管道都没有这些同步问题。感兴趣的读者可以查看 本文档以了解更多信息。由于与提及的其他因素相比，这些同步导致的延迟增加很小，因此在本文的其余部分我们将不关注它们。
扩散模型架构
我们将以 Black Forest Labs 的开源文生图模型 Flux‑1‑Dev 作为我们的示例。扩散管道不是单个网络；它是一组模型：
- 文本编码器 – CLIP‑Text 和 T5 将用户提示转换为嵌入。
- 去噪器 – 一个扩散 Transformer (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")
现在编译计算量大的扩散 Transformer子模块：
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 是一个即时编译器，它会根据它看到的输入的属性（包括形状、dtype 和设备）专门化编译器工件（有关更多详细信息，请参阅这篇博客）。更改其中任何一项都会导致重新编译。尽管这会在后台自动发生，但这种重新编译会导致更高的编译时间成本，从而导致糟糕的用户体验。
如果您的应用程序需要处理多种图像尺寸或批次形状，请在编译时传递 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 倍更短）提供了相同的加速效果。
- 内存受限的 GPU – 许多 Diffusers 用户在 VRAM 无法容纳整个模型的卡上工作。我们将研究 CPU 卸载和量化，以在这些设备上保持生成的可行性。
- 使用 LoRA 适配器快速个性化 – 通过低秩适配器进行微调是使扩散模型适应新风格或任务的首选方法。我们将演示如何交换 LoRA 而不会触发重新编译。
内存受限的 GPU
CPU 卸载： 一个完整的 Flux 管道（bfloat16）大约消耗 33 GB，这超过了大多数消费级 GPU 所能提供的内存。幸运的是，并非所有子模块都必须在整个前向传递中占用 GPU 内存。一旦文本编码器完成了提示嵌入的生成，它们就可以被移动到系统 RAM 中。同样，在 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 支持的健壮性。该套件 运行最流行的管道，并自动检查：
- 图中断和静默回退
- 跨常见输入形状的意外重新编译
- 编译、CPU 卸载和我们支持的每个量化后端之间的兼容性
鼓励感兴趣的读者查看此链接，以查看近期改进 torch.compile 覆盖范围的 PR。
基准测试
正确性只是故事的一半；我们同样关心性能。一个新的基准测试工作流程现在与 CI 一起运行，捕获了本文中涵盖的每种场景的延迟和峰值内存。结果被导出到一个汇总 CSV中，因此很容易发现回归（或胜利！）。设计和初步结果见此 PR。
结论
torch.compile 可以将标准的 Diffusers 管道转变为高性能、内存高效的工作负载。通过将编译集中在 DiT 上，利用区域编译和动态形状，并将编译器与卸载、量化和 LoRA 热插拔堆叠，您可以在不牺牲图像质量或灵活性的情况下，释放出显著的速度提升和 VRAM 节省。
我们希望这些秘诀能激励您将 torch.compile 融入您自己的扩散工作流程中。我们期待看到您接下来构建什么。
编译愉快 ⚡️
Diffusers 团队要感谢 Animesh 和 Ryan 在改进 torch.compile 支持方面的帮助和支持。