基于 Quanto 和 Diffusers 的内存高效 transformer 扩散模型

过去的几个月，我们目睹了使用基于 transformer 模型作为扩散模型的主干网络来进行高分辨率文生图 (text-to-image，T2I) 的趋势。和一开始的许多扩散模型普遍使用 UNet 架构不同，这些模型使用 transformer 架构作为扩散过程的主模型。由于 transformer 的性质，这些主干网络表现出了良好的可扩展性，模型参数量可从 0.6B 扩展至 8B。

随着模型越变越大，内存需求也随之增加。对扩散模型而言，这个问题愈加严重，因为扩散流水线通常由多个模型串成: 文本编码器、扩散主干模型和图像解码器。此外，最新的扩散流水线通常使用多个文本编码器 - 如: Stable Diffusion 3 有 3 个文本编码器。使用 FP16 精度对 SD3 进行推理需要 18.765GB 的 GPU 显存。

这么高的内存要求使得很难将这些模型运行在消费级 GPU 上，因而减缓了技术采纳速度并使针对这些模型的实验变得更加困难。本文，我们展示了如何使用 Diffusers 库中的 Quanto 量化工具脚本来提高基于 transformer 的扩散流水线的内存效率。

目录

• 基于 Quanto 和 Diffusers 的内存高效 transformer 扩散模型

  • 目录

  • 基础知识

  • 用 Quanto 量化 DiffusionPipeline

  • 上述攻略通用吗？

  • 其他发现

    • 在 H100 上 bfloat16 通常表现更好
    • qint8 的前途
    • INT4 咋样？

  • 加个鸡腿 - 在 Quanto 中保存和加载 Diffusers 模型

  • 小诀窍

  • 总结

基础知识

你可参考 这篇文章 以获取 Quanto 的详细介绍。简单来说，Quanto 是一个基于 PyTorch 的量化工具包。它是 Hugging Face Optimum 的一部分，Optimum 提供了一套硬件感知的优化工具。

模型量化是 LLM 从业者必备的工具，但在扩散模型中并不算常用。Quanto 可以帮助弥补这一差距，其可以在几乎不伤害生成质量的情况下节省内存。

我们基于 H100 GPU 配置进行基准测试，软件环境如下:

• CUDA 12.2
• PyTorch 2.4.0
• Diffusers (从源代码安装，参考 此提交)
• Quanto (从源代码安装，参考 此提交)

除非另有说明，我们默认使用 FP16 进行计算。我们不对 VAE 进行量化以防止数值不稳定问题。你可于 此处 找到我们的基准测试代码。

截至本文撰写时，以下基于 transformer 的扩散模型流水线可用于 Diffusers 中的文生图任务:

• PixArt-Alpha 及 PixArt-Sigma
• Stable Diffusion 3
• Hunyuan DiT
• Lumina
• Aura Flow

另外还有一个基于 transformer 的文生视频流水线: Latte。

为简化起见，我们的研究仅限于以下三个流水线: PixArt-Sigma、Stable Diffusion 3 以及 Aura Flow。下表显示了它们各自的扩散主干网络的参数量:

用 Quanto 量化 DiffusionPipeline

使用 Quanto 量化模型非常简单。

from optimum.quanto import freeze, qfloat8, quantize
from diffusers import PixArtSigmaPipeline
import torch

pipeline = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
).to("cuda")

quantize(pipeline.transformer, weights=qfloat8)
freeze(pipeline.transformer)

我们对需量化的模块调用 quantize() ，以指定我们要量化的部分。上例中，我们仅量化参数，保持激活不变，量化数据类型为 FP8。最后，调用 freeze() 以用量化参数替换原始参数。

然后，我们就可以如常调用这个 pipeline 了:

image = pipeline("ghibli style, a fantasy landscape with castles").images[0]

FP16	将 transformer 扩散主干网络量化为 FP8

我们注意到使用 FP8 可以节省显存，且几乎不影响生成质量; 我们也看到量化模型的延迟稍有变长:

我们可以用相同的方式量化文本编码器:

quantize(pipeline.text_encoder, weights=qfloat8)
freeze(pipeline.text_encoder)

文本编码器也是一个 transformer 模型，我们也可以对其进行量化。同时量化文本编码器和扩散主干网络可以带来更大的显存节省:

量化文本编码器后生成质量与之前的情况非常相似:

上述攻略通用吗？

将文本编码器与扩散主干网络一起量化普遍适用于我们尝试的很多模型。但 Stable Diffusion 3 是个特例，因为它使用了三个不同的文本编码器。我们发现 _ 第二个 _ 文本编码器量化效果不佳，因此我们推荐以下替代方案:

• 仅量化第一个文本编码器 (CLIPTextModelWithProjection) 或
• 仅量化第三个文本编码器 (T5EncoderModel) 或
• 同时量化第一个和第三个文本编码器

下表给出了各文本编码器量化方案的预期内存节省情况 (扩散 transformer 在所有情况下均被量化):

量化文本编码器: 1	量化文本编码器: 3	量化文本编码器: 1 和 3

其他发现

在 H100 上 bfloat16 通常表现更好

对于支持 bfloat16 的 GPU 架构 (如 H100 或 4090)，使用 bfloat16 速度更快。下表列出了在我们的 H100 参考硬件上测得的 PixArt 的一些数字: Batch Size 精度 量化 显存 (GB) 延迟 (秒) 是否量化文本编码器

qint8 的前途

我们发现使用 qint8 (而非 qfloat8 ) 进行量化，推理延迟通常更好。当我们对注意力 QKV 投影进行水平融合 (在 Diffusers 中调用 fuse_qkv_projections() ) 时，效果会更加明显，因为水平融合会增大 int8 算子的计算维度从而实现更大的加速。我们基于 PixArt 测得了以下数据以证明我们的发现:

INT4 咋样？

在使用 bfloat16 时，我们还尝试了 qint4 。目前我们仅支持 H100 上的 bfloat16 的 qint4 量化，其他情况尚未支持。通过 qint4 ，我们期望看到内存消耗进一步降低，但代价是推理延迟变长。延迟增加的原因是硬件尚不支持 int4 计算 - 因此权重使用 4 位，但计算仍然以 bfloat16 完成。下表展示了 PixArt-Sigma 的结果:

但请注意，由于 INT4 量化比较激进，最终结果可能会受到影响。所以，一般对于基于 transformer 的模型，我们通常不量化最后一个投影层。在 Quanto 中，我们做法如下:

quantize(pipeline.transformer, weights=qint4, exclude="proj_out")
freeze(pipeline.transformer)

"proj_out" 对应于 pipeline.transformer 的最后一层。下表列出了各种设置的结果:

量化文本编码器: 否 , 不量化的层: 无	量化文本编码器: 否 , 不量化的层: "proj_out"	量化文本编码器: 是 , 不量化的层: 无	量化文本编码器: 是 , 不量化的层: "proj_out"

为了恢复损失的图像质量，常见的做法是进行量化感知训练，Quanto 也支持这种训练。这项技术超出了本文的范围，如果你有兴趣，请随时与我们联系！

本文的所有实验结果都可以在 这里 找到。

加个鸡腿 - 在 Quanto 中保存和加载 Diffusers 模型

以下代码可用于对 Diffusers 模型进行量化并保存量化后的模型:

from diffusers import PixArtTransformer2DModel
from optimum.quanto import QuantizedPixArtTransformer2DModel, qfloat8

model = PixArtTransformer2DModel.from_pretrained("PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", subfolder="transformer")
qmodel = QuantizedPixArtTransformer2DModel.quantize(model, weights=qfloat8)
qmodel.save_pretrained("pixart-sigma-fp8")

此代码生成的 checkpoint 大小为 587MB ，而不是原本的 2.44GB。然后我们可以加载它:

from optimum.quanto import QuantizedPixArtTransformer2DModel
import torch

transformer = QuantizedPixArtTransformer2DModel.from_pretrained("pixart-sigma-fp8")
transformer.to(device="cuda", dtype=torch.float16)

最后，在 DiffusionPipeline 中使用它:

from diffusers import DiffusionPipeline
import torch

pipe = DiffusionPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
    transformer=None,
    torch_dtype=torch.float16,
).to("cuda")
pipe.transformer = transformer

prompt = "A small cactus with a happy face in the Sahara desert."
image = pipe(prompt).images[0]

将来，我们计划支持在初始化流水线时直接传入 transformer 就可以工作:

pipe = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
- transformer=None,
+ transformer=transformer,
    torch_dtype=torch.float16,
).to("cuda")

QuantizedPixArtTransformer2DModel 实现可参考 此处。如果你希望 Quanto 支持对更多的 Diffusers 模型进行保存和加载，请在 此处 提出需求并 @sayakpaul 。

小诀窍

• 根据应用场景的不同，你可能希望对流水线中不同的模块使用不同类型的量化。例如，你可以对文本编码器进行 FP8 量化，而对 transformer 扩散模型进行 INT8 量化。由于 Diffusers 和 Quanto 的灵活性，你可以轻松实现这类方案。
• 为了优化你的用例，你甚至可以将量化与 Diffuser 中的其他 内存优化技术 结合起来，如 enable_model_cpu_offload() 。

总结

本文，我们展示了如何量化 Diffusers 中的 transformer 模型并优化其内存消耗。当我们同时对文本编码器进行量化时，效果变得更加明显。我们希望大家能将这些工作流应用到你的项目中并从中受益🤗。

感谢 Pedro Cuenca 对本文的细致审阅。