Quanto 量化削减 Transformer 扩散管线的内存占用

TL;DR

Quantization with Hugging Face Quanto reduces the GPU memory needed for Transformer‑based diffusion pipelines (e.g., PixArt‑Sigma, Stable Diffusion 3, Aura Flow) from ~12 GB to as low as ~5 GB with only minor latency impact and negligible quality loss.

引言 – 为什么内存对扩散 Transformer 很重要

Transformer backbones have become the dominant architecture for high‑resolution text‑to‑image diffusion models, scaling from 0.6 B to 8 B parameters. Larger models increase GPU memory consumption dramatically; a full Stable Diffusion 3 inference in FP16 occupies 18.8 GB. This memory barrier limits adoption on consumer GPUs and hampers rapid experimentation. The post demonstrates that Quanto’s quantization utilities, integrated in the Diffusers library, can dramatically shrink the memory footprint while preserving visual quality.

Quanto 量化基础

Quanto is a PyTorch‑based quantization toolkit that lives inside Hugging Face Optimum. It supports weight‑only quantization to several low‑precision formats (FP8, INT8, INT4) and can be applied to any Diffusers module.

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 the diffusion transformer to FP8
quantize(pipeline.transformer, weights=qfloat8)
freeze(pipeline.transformer)

The same quantize/freeze calls work for text encoders or any other sub‑module.

内存与延迟结果 – FP8 权重量化

仅扩散 Transformer

批量大小 内存 (GB) 延迟 (s)
1 11.55 (FP8) vs 12.09 (FP16) 1.54 vs 1.20
4 11.55 (FP8) vs 12.09 (FP16) 5.11 vs 4.48

FP8 权重将内存削减约 0.5 GB，延迟略有增加。

添加文本编码器量化

批量大小 是否量化 TE? 内存 (GB) 延迟 (s)
1 No 11.55 1.54
1 Yes 5.36 1.60
4 No 11.55 5.11
4 Yes 5.36 5.14

同时对扩散 Transformer 和文本编码器进行量化可将内存使用减半，而延迟基本保持不变。

跨模型的通用性

作者评估了三个管线：

  • PixArt‑Sigma（0.61 B 参数）
  • Stable Diffusion 3 (medium)（2.03 B 参数，三个文本编码器）
  • Aura Flow（6.84 B 参数）

对于 PixArt‑Sigma 和 Aura Flow，量化文本编码器始终能带来大量内存节省。Stable Diffusion 3 需要有选择地进行量化，因为第二个文本编码器（中间的 CLIP 变体）在量化后会降低质量。推荐的策略有：

  • 仅量化第一个 CLIP 编码器，或
  • 仅量化第三个 T5 编码器，或
  • 同时量化第一和第三个编码器。

一张针对 SD‑3（批量 1，扩散 Transformer 始终使用 FP8）的示例表格显示，内存范围从 8.20 GB（全部三个编码器均量化）到 16.40 GB（未量化）。

其他发现

bfloat16 与 fp16 在 H100 上的比较

在 NVIDIA H100 GPU 上，使用 bfloat16 搭配 INT8 或 FP8 权重可提升延迟表现：

精度 量化方式 内存 (GB) 延迟 (s)
FP16 INT8 5.363 1.538
BF16 INT8 5.364 1.454
FP16 FP8 5.363 1.601
BF16 FP8 5.363 1.495

INT8（qint8）与融合的 QKV 投影

INT8 权重比 FP8 更快，尤其在注意力 QKV 投影被融合时（fuse_qkv_projections()）。对于 PixArt‑Sigma（批量 1），延迟从 1.538 s（INT8，未融合）下降至 1.504 s（INT8 融合 QKV）。

INT4（qint4）用于激进压缩

在 H100 上与 bfloat16 结合使用时，INT4 可显著降低内存（例如 PixArt‑Sigma 从 9.38 GB 降至 3.06 GB），但延迟上升（≈7.6 s），因为计算仍在 bfloat16 中进行。质量损失明显；作者建议在量化时排除最终投影层（proj_out），以减轻退化。

保存与加载量化的 Diffusers 模型

Quanto 提供了可持久化并重新加载的模型类：

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")  # 587 MB checkpoint

加载遵循相同的模式，量化后的 Transformer 可注入到 DiffusionPipeline 中。

开发者实用技巧

  • 根据模块混合使用不同的量化类型（例如，对文本编码器使用 FP8，对扩散 Transformer 使用 INT8），以在内存和速度之间取得平衡。
  • 将 Quanto 量化与 Diffusers 已有的内存节省工具（如 enable_model_cpu_offload()）结合使用，以进一步降低内存占用。
  • 对于 INT4 部署，始终排除最终投影层（exclude="proj_out"），以保持图像保真度。

结论

使用 Hugging Face Quanto 对基于 Transformer 的扩散管线进行量化，可将 GPU 内存需求降低至多 70 %（从约 12 GB 降至约 5 GB），同时保持低延迟和基本不受影响的视觉质量。该技术适用于多种最先进模型，且生成的检查点体积大幅缩小（例如 587 MB 对比 2.44 GB）。开发者可以轻松将这些工作流集成到现有的 Diffusers 管线中，并进一步结合其他内存优化策略。

致谢

感谢 Pedro Cuenca 对本文的深入审阅。

Sources