Diffusers 与 PEFT 优化 Flux LoRA 推理，H100 上提速 2.23x

LoRA 适配器为各种规模和形态的模型提供了大量的定制能力。在图像生成方面，它们可以赋予模型不同的风格、不同的角色以及更多。有时，它们还可以被用来降低推理延迟。因此，它们的重要性至关重要，尤其是在定制和微调模型时。

在这篇文章中，我们选用 Flux.1-Dev 模型进行文本到图像生成，因为它广受欢迎并被广泛采用，并探讨如何在使用 LoRA 时优化其推理速度（约 2.3 倍）。在 Hugging Face Hub 平台上，据报道已有超过 3 万个适配器基于它训练。因此，它对社区的重要性非常显著。

请注意，尽管我们展示了 Flux 上的加速效果，但我们相信我们的方案足够通用，也可以应用于其他模型。

如果您迫不及待想开始使用代码，请查看配套代码仓库。

目录

• 优化 LoRA 推理的障碍
• 优化方案
• 在消费级 GPU 上优化 LoRA 推理
• 结论

优化 LoRA 推理的障碍

在服务 LoRA 时，热插拔（换入和换出不同的 LoRA）是很常见的做法。LoRA 会改变基础模型的架构。此外，LoRA 之间也可能各不相同——每个 LoRA 可能具有不同的秩以及不同的目标适配层。为了应对 LoRA 的这些动态特性，我们必须采取必要措施，确保所应用的优化具有鲁棒性。

例如，我们可以在加载了特定 LoRA 的模型上应用 torch.compile，以在推理延迟上获得加速。然而，一旦我们将该 LoRA 换成另一个（可能具有不同配置的）LoRA，就会遇到重新编译问题，导致推理变慢。

也可以将 LoRA 参数融合到基础模型参数中，运行编译，然后在加载新 LoRA 时取消融合。然而，由于可能存在架构层面的变化，这种方法在每次运行推理时都会再次遇到重新编译的问题。

我们的优化方案考虑了上述情况，以尽可能贴近实际。以下是我们的优化方案的关键组成部分：

• Flash Attention 3（FA3）
• torch.compile
• 来自 TorchAO 的 FP8 量化
• 支持热插拔

请注意，在上述内容中，FP8 量化是有损的，但通常能提供最强大的速度-内存权衡。尽管我们主要在 NVIDIA GPU 上测试了该方案，但它也应该能在 AMD GPU 上运行。

优化方案

在我们之前的博客文章（文章 1 和 文章 2）中，我们已经讨论了使用优化方案前三个组成部分的好处。逐一应用它们只需几行代码：

from diffusers import DiffusionPipeline, TorchAoConfig
from diffusers.quantizers import PipelineQuantizationConfig
from utils.fa3_processor import FlashFluxAttnProcessor3_0
import torch

# quantize the Flux transformer with FP8
pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16,
    quantization_config=PipelineQuantizationConfig(
        quant_mapping={"transformer": TorchAoConfig("float8dq_e4m3_row")}
    )
).to("cuda")

# use Flash-attention 3
pipe.transformer.set_attn_processor(FlashFluxAttnProcessor3_0())

# use torch.compile()
pipe.transformer.compile(fullgraph=True, mode="max-autotune")

# perform inference
pipe_kwargs = {
    "prompt": "A cat holding a sign that says hello world",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 28,
    "max_sequence_length": 512,
}

# first time will be slower, subsequent runs will be faster
image = pipe(**pipe_kwargs).images[0]

FA3 处理器来自这里。

当我们尝试在不触发重新编译的情况下，将 LoRA 换入和换出已编译的扩散 Transformer（pipe.transformer）时，问题就开始浮现了。

通常，加载和卸载 LoRA 需要重新编译，这抵消了编译带来的任何速度优势。幸运的是，有一种方法可以避免重新编译。通过传递 hotswap=True，diffusers 将保持模型架构不变，仅交换 LoRA 适配器本身的权重，这不需要重新编译。

pipe.enable_lora_hotswap(target_rank=max_rank)
pipe.load_lora_weights(<lora-adapter-name1>)
# compile *after* loading the first LoRA
pipe.transformer.compile(mode="max-autotune", fullgraph=True)
image = pipe(**pipe_kwargs).images[0]
# from this point on, load new LoRAs with `hotswap=True`
pipe.load_lora_weights(<lora-adapter-name2>, hotswap=True)
image = pipe(**pipe_kwargs).images[0]

（提醒一下，首次调用 pipe 会很慢，因为 torch.compile 是即时编译器。但是，后续调用应该会快得多。）

这通常允许在不重新编译的情况下交换 LoRA，但有一些限制：

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

有关 Diffusers 中热交换及其限制的更多信息，请访问文档的热交换部分。

当我们查看不使用编译但使用热交换的推理延迟时，这种工作流程的优势就变得明显了。

关键要点：

• “常规 + 编译”选项比常规选项提供了不错的加速，但它会引发重新编译问题，从而增加整体执行时间。在我们的基准测试中，我们没有展示编译时间。
• 当通过热交换（也称为“优化”选项）消除重新编译问题时，我们实现了最高的加速。
• 在“优化”选项中，启用了 FP8 量化，这可能导致质量损失。即使不使用 FP8，我们也能获得相当不错的加速（“no_fp8”选项）。
• 为了演示目的，我们使用两个 LoRA 的池进行热交换与编译。完整代码请参考附带的代码仓库。

到目前为止我们讨论的优化方案假设可以访问像 H100 这样的强大 GPU。然而，当我们只能使用像 RTX 4090 这样的消费级 GPU 时，我们能做什么呢？让我们来看看。

在消费级 GPU 上优化 LoRA 推理

Flux.1-Dev（不带任何 LoRA），使用 Bfloat16 数据类型，运行需要约 33GB 内存。根据 LoRA 模块的大小，并且不使用任何优化，这个内存占用可能会进一步增加。许多消费级 GPU 如 RTX 4090 只有 24GB。在本节的剩余部分，我们将考虑使用 RTX 4090 机器作为我们的测试平台。

首先，为了实现 Flux.1-Dev 的端到端执行，我们可以应用 CPU 卸载，其中执行当前计算不需要的组件被卸载到 CPU，以释放更多加速器内存。这样做允许我们在 RTX 4090 上以约 22GB 内存运行整个管道，耗时 35.403 秒。启用编译可以将延迟降低到 31.205 秒（1.12 倍加速）。在代码方面，只需几行：

pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()
# Instead of full compilation, we apply regional compilation
# here to take advantage of `fullgraph=True` and also to reduce
# compilation time. More details can be found here:
# https://hf.co/docs/diffusers/main/en/optimization/fp16#regional-compilation
pipe.transformer.compile_repeated_blocks(fullgraph=True)
image = pipe(**pipe_kwargs).images[0]

请注意，我们在这里没有应用 FP8 量化，因为 CPU 卸载和编译不支持它（相关 issue 讨论）。因此，仅对 Flux Transformer 应用 FP8 量化也不足以缓解内存耗尽问题。在这种情况下，我们决定将其移除。

因此，要利用 FP8 量化方案，我们需要找到一种无需 CPU 卸载的方法。对于 Flux.1-Dev，如果我们额外对 T5 文本编码器应用量化，我们应该能够在 24GB 内加载并运行完整流水线。下面是对 T5 文本编码器进行量化与不进行量化的结果对比（NF4 量化来自 bitsandbytes）。

如上图所示，量化 T5 文本编码器不会带来太大的质量损失。将量化后的 T5 文本编码器与 FP8 量化的 Flux Transformer 结合使用 torch.compile，可以得到还算合理的结果——从 32.27 秒降至 9.668 秒（大幅加速约 3.3 倍），且没有明显的质量下降。

即使不对 T5 文本编码器进行量化，也可以在 24 GB 显存下生成图像，但这会让我们的生成流水线稍微复杂一些。

现在我们有了在 RTX 4090 上以 FP8 量化运行整个 Flux.1-Dev 流水线的方法。我们可以应用之前建立的优化方案，在同一硬件上优化 LoRA 推理。由于 RTX 4090 不支持 FA3，我们将坚持使用以下优化方案，并新加入 T5 量化：

• FP8 量化
• torch.compile
• 支持热插拔
• T5 量化（使用 NF4）

在下表中，我们展示了应用上述不同组件组合时的推理延迟数据。

快速说明：

• 编译相比基线提供了 2 倍的巨大加速。
• 其他选项即使在启用卸载的情况下也产生了 OOM 错误。

热插拔的技术细节

要在不触发重新编译的情况下启用热插拔，必须克服两个障碍。首先，必须将 LoRA 缩放因子从浮点数转换为 torch 张量，这相当容易实现。其次，LoRA 权重的形状需要填充到所需的最大形状。这样，权重中的数据就可以被替换，而无需重新分配整个属性。这就是上面讨论的 max_rank 参数至关重要的原因。由于我们用零填充值，结果保持不变，尽管根据填充的大小，计算会稍微变慢。

由于没有添加新的 LoRA 属性，这也要求第一个之后的每个 LoRA 只能针对第一个所针对的相同层或其子集。因此，请明智地选择加载顺序。如果 LoRA 针对不相交的层，则有可能创建一个针对所有目标层并集的虚拟 LoRA。

要了解此实现的细节，请访问 PEFT 中的 hotswap.py 文件。

结论

本文概述了一种使用 Flux 进行快速 LoRA 推理的优化方案，展示了显著的加速效果。我们的方法结合了 Flash Attention 3、torch.compile 和 FP8 量化，同时确保热插拔能力且不会出现重新编译问题。在 H100 等高端 GPU 上，此优化设置相比基线提供了 2.23 倍的加速。

对于消费级 GPU，特别是 RTX 4090，我们通过引入 T5 文本编码器量化（NF4）并利用区域编译解决了内存限制。这套综合方案实现了 2.04 倍的显著加速，使 Flux 上的 LoRA 推理即使在显存有限的情况下也变得可行且高效。关键洞察在于，通过精心管理编译和量化，LoRA 的优势可以在不同硬件配置上得到充分发挥。

希望本文中的方案能启发你优化基于 LoRA 的使用场景，享受快速推理带来的好处。

资源

以下是我们在本文中引用的重要资源列表：

• Presenting Flux Fast: Making Flux go brrr on H100s
• torch.compile and Diffusers: A Hands-On Guide to Peak Performance
• LoRA guide in Diffusers
• 对于有兴趣尝试此方法的读者，Parag Ekbote 部署的演示可在 Replicate 上获取。该演示已在 NVIDIA L40 和 A100 GPU 上设计和测试，展现出相当的性能。