Hugging Face 实测 SDXL 推理优化：fp16、SDPA、torch.compile 与 CPU offload 的显存和延迟对比

Stable Diffusion XL (SDXL) 是 Stability AI 推出的最新潜在扩散模型，用于生成高质量的超级逼真图像。它克服了之前 Stable Diffusion 模型的一些挑战，比如正确生成手部和文本，以及空间上正确的构图。此外，SDXL 还具有更强的上下文感知能力，并且在其提示中需要更少的词语就能生成更好看的图像。

然而，所有这些改进都是以模型显著增大为代价的。大了多少呢？基础 SDXL 模型有 35 亿个参数（尤其是 UNet），大约是之前 Stable Diffusion 模型的 3 倍。

为了探索如何针对推理速度和内存使用优化 SDXL，我们在 A100 GPU（40 GB）上运行了一些测试。每次推理运行，我们生成 4 张图像并重复 3 次。在计算推理延迟时，我们只考虑 3 次迭代中的最后一次。

所以，如果你直接以全精度、使用默认注意力机制开箱即用地运行 SDXL，它会消耗 28GB 内存并耗时 72.2 秒！

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0").to("cuda")
pipe.unet.set_default_attn_processor()

这不太实用，而且会拖慢你的速度，因为你通常生成的图像不止 4 张。而且如果你没有更强大的 GPU，就会遇到令人沮丧的内存不足错误消息。那么，我们如何优化 SDXL 以提高推理速度并减少其内存使用呢？

在 🤗 Diffusers 中，我们有许多优化技巧和技术，可以帮助你运行像 SDXL 这样内存密集型的模型，我们会向你展示如何做到！我们将重点关注的两件事是推理速度和内存。

🧠 本文讨论的技术适用于所有 pipeline。

推理速度

扩散是一个随机过程，所以无法保证你会得到一张喜欢的图像。很多时候，你需要多次运行推理并反复迭代，这就是为什么针对速度进行优化至关重要。本节重点介绍使用较低精度的权重，并结合内存高效的注意力机制以及来自 PyTorch 2.0 的 torch.compile 来提升速度并减少推理时间。

较低精度

模型权重以某种精度存储，精度用浮点数据类型表示。标准的浮点数据类型是 float32（fp32），它可以准确表示范围很广的浮点数。对于推理，你通常不需要那么高的精度，所以应该使用 float16（fp16），它表示的浮点数范围更窄。这意味着与 fp32 相比，fp16 存储时只占用一半的内存，而且由于计算更容易，速度也快一倍。此外，现代 GPU 显卡拥有经过优化的硬件来运行 fp16 计算，使其速度更快。

使用 🤗 Diffusers，你可以通过指定 torch.dtype 参数在模型加载时转换权重，从而使用 fp16 进行推理：

from diffusers import StableDiffusionXLPipeline

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

与完全未优化的 SDXL pipeline 相比，使用 fp16 只占用 21.7GB 内存，并且仅需 14.8 秒。你几乎把推理速度提升了整整一分钟！

内存高效的注意力

transformer 模块中使用的注意力块可能是一个巨大的瓶颈，因为内存会随着输入序列变长而呈二次方增长。这会很快占用大量内存，并让你遇到内存不足的错误消息。😬

内存高效的注意力算法旨在减少计算注意力时的内存负担，无论是通过利用稀疏性还是分块技术。这些优化算法过去大多作为需要单独安装的第三方库提供。但从 PyTorch 2.0 开始，情况不再如此。PyTorch 2 引入了 缩放点积注意力 (SDPA)，它提供了 Flash Attention、内存高效注意力 (xFormers) 以及 C++ 实现的 PyTorch 版本的融合实现。SDPA 可能是加速推理最简单的方法：如果你使用的是 PyTorch ≥ 2.0 和 🤗 Diffusers，它会自动默认启用！

from diffusers import StableDiffusionXLPipeline

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

与完全未优化的 SDXL 管道相比，使用 fp16 和 SDPA 占用相同的内存量，推理时间提升至 11.4 秒。让我们以此作为新的基线，与其他优化进行比较。

torch.compile

PyTorch 2.0 还引入了 torch.compile API，用于将你的 PyTorch 代码即时 (JIT) 编译为更优化的推理内核。与其他编译器解决方案不同，torch.compile 对现有代码的修改要求极少，只需用该函数包装你的模型即可。

通过 mode 参数，你可以在编译期间针对内存开销或推理速度进行优化，这为你提供了更大的灵活性。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

与之前的基线（fp16 + SDPA）相比，用 torch.compile 包装 UNet 将推理时间提升至 10.2 秒。

⚠️ 首次编译模型时速度较慢，但一旦模型编译完成，后续所有调用都会快得多！

模型内存占用

如今的模型越来越大，将其装入内存成为一项挑战。本节重点介绍如何减少这些庞大模型的内存占用，以便你可以在消费级 GPU 上运行它们。这些技术包括 CPU 卸载、分多步而非一次性将潜在表示解码为图像，以及使用自动编码器的蒸馏版本。

模型 CPU 卸载

模型卸载通过将 UNet 加载到 GPU 内存中，同时将扩散模型的其他组件（文本编码器、VAE）加载到 CPU 上来节省内存。这样，UNet 可以在 GPU 上运行多次迭代，直到不再需要为止。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
)
pipe.enable_model_cpu_offload()

与基线相比，现在需要 20.2GB 内存，为你节省了 1.5GB 内存。

顺序 CPU 卸载

另一种类型的卸载是顺序 CPU 卸载，它可以节省更多内存，但代价是推理速度变慢。它不是卸载整个模型（如 UNet），而是将存储在不同 UNet 子模块中的模型权重卸载到 CPU，仅在前向传递之前才加载到 GPU 上。本质上，你每次只加载模型的一部分，这允许你节省更多内存。唯一的缺点是速度明显变慢，因为你需要多次加载和卸载子模块。

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
)
pipe.enable_sequential_cpu_offload()

与基线相比，这需要 19.9GB 内存，但推理时间增加到 67 秒。

切片

在 SDXL 中，变分编码器 (VAE) 将（由 UNet 预测的）精炼潜在表示解码为逼真的图像。此步骤的内存需求随预测图像的数量（批量大小）而变化。根据图像分辨率和可用的 GPU 显存，它可能会非常占用内存。

这就是“切片”技术的用武之地。待解码的输入张量被分割成多个切片，解码计算分多步完成。这样可以节省内存并支持更大的批处理规模。

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

通过切片计算，我们将内存降至 15.4GB。如果再加上顺序 CPU 卸载，内存进一步降至 11.45GB，让你每个提示词可以生成 4 张图像（1024x1024）。不过，使用顺序卸载也会增加推理延迟。

缓存计算

任何文本条件图像生成模型通常都会使用一个文本编码器来从输入提示词计算嵌入。SDXL 使用了两个文本编码器！这对推理延迟贡献不小。不过，由于这些嵌入在反向扩散过程中保持不变，我们可以预先计算它们并在过程中重复使用。这样，在计算完文本嵌入后，我们就可以将文本编码器从内存中移除。

首先，加载文本编码器及其对应的分词器，并从输入提示词计算嵌入：

tokenizers = [tokenizer, tokenizer_2]
text_encoders = [text_encoder, text_encoder_2]

(
    prompt_embeds,
    negative_prompt_embeds,
    pooled_prompt_embeds,
    negative_pooled_prompt_embeds
) = encode_prompt(tokenizers, text_encoders, prompt)

接下来，清空 GPU 内存以移除文本编码器：

del text_encoder, text_encoder_2, tokenizer, tokenizer_2
flush()

现在嵌入可以直接送入 SDXL 流水线了：

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    torch_dtype=torch.float16,
).to("cuda")

call_args = dict(
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        num_images_per_prompt=num_images_per_prompt,
        num_inference_steps=num_inference_steps,
)
image = pipe(**call_args).images[0]

结合 SDPA 和 fp16，我们可以将内存降至 21.9GB。上面讨论的其他内存优化技术也可以与缓存计算配合使用。

微型自编码器

如前所述，VAE 将潜在表示解码为图像。自然，这一步直接受限于 VAE 的大小。所以，我们干脆用一个更小的自编码器！由 madebyollin 开发的微型自编码器，可在Hub上获取，仅有 10MB，它是从 SDXL 使用的原始 VAE 中蒸馏而来的。

from diffusers import AutoencoderTiny

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
)
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16)
pipe.to("cuda")

通过这种设置，我们将内存需求降至 15.6GB，同时降低了推理延迟。

⚠️ 微型自编码器可能会遗漏图像中一些更精细的细节，因此微型自编码器更适合用于图像预览。

结论

最后，总结一下我们各项优化所节省的资源：

⚠️ 在对 GPU 进行性能分析以衡量推理延迟与内存需求之间的权衡时，注意所使用的硬件非常重要。上述发现可能不会在不同硬件之间等同适用。例如，`torch.compile` 似乎只对现代 GPU 有益，至少对 SDXL 而言是如此。

我们希望这些优化能让你轻松运行自己喜爱的流水线。试试这些技术，并把你的图像分享给我们！🤗

致谢：感谢 Pedro Cuenca 对草稿提出的宝贵审阅意见。