Hugging Face Diffusers 量化后端

Hugging Face Diffusers 现在集成了多个量化后端，使得大型扩散模型（如 FLUX.1-dev）通过降低内存和计算需求变得更易使用。这些后端让用户在不显著牺牲图像质量的情况下压缩模型，8 位量化往往能够产生几乎与高精度 BF16 模型无异的结果。

FLUX.1-dev 模型架构与内存基准

要了解量化的影响，需要先确定 FLUX.1-dev 模型在 BF16 精度下的内存基准，大约需要 31.447 GB 内存。该模型由四个主要组件组成：

  • Transformer (MMDiT)： 核心生成式多模态扩散 Transformer，需 23.8 GB。
  • Text Encoder 2 (T5)： 用于细致的理解和文本生成，需 9.52 GB。
  • Text Encoder 1 (CLIP)： 用于初步的文本理解，需 246 MB。
  • VAE： 在像素空间和潜在空间之间转换图像，需 168 MB。

量化工作主要针对 Transformer 和 T5 文本编码器，以实现最大的内存节省。

支持的量化后端

bitsandbytes (BnB)

bitsandbytes 提供 8 位和 4 位量化。对于 FLUX.1-dev，4 位量化（使用 NF4）将加载后内存降至 12.584 GB，峰值内存降至 17.281 GB，同时在 NVIDIA H100 80GB GPU 上保持 12 秒的推理时间。8 位量化则会增加内存使用（加载 19.273 GB / 峰值 24.432 GB），并将推理时间放慢至 27 秒。

torchao

torchao 是一个基于 PyTorch 的架构优化库。它支持多种仅权重量化格式：

  • int4_weight_only： 最低内存占用（加载 10.635 GB / 峰值 14.654 GB），但推理最慢（109 秒）。
  • int8_weight_only： 中等内存（加载 17.020 GB / 峰值 21.482 GB），推理快速（15 秒）。
  • float8_weight_only： 与 int8 类似的内存和速度（加载 17.016 GB / 峰值 21.488 GB；15 秒）。

Quanto

通过 optimum 库集成，Quanto 支持多种精度：

  • INT4： 加载 12.254 GB / 峰值 16.139 GB；推理 109 秒。
  • INT8： 加载 17.330 GB / 峰值 21.814 GB；推理 15 秒。
  • FP8： 加载 16.395 GB / 峰值 20.898 GB；推理 16 秒。

GGUF

Diffusers 支持 GGUF 文件格式，允许通过 from_single_file 使用来自 llama.cpp 社区的预量化模型。FLUX.1-dev 的基准测试显示：

  • Q2_k： 加载 13.264 GB / 峰值 17.752 GB；推理 26 秒。
  • Q4_1： 加载 16.838 GB / 峰值 21.326 GB；推理 23 秒。
  • Q8_0： 加载 21.502 GB / 峰值 25.973 GB；推理 15 秒。

FP8 分层转型

enable_layerwise_casting 是一种内存优化方式，将权重存储为 FP8（e4m3），但在计算时动态转换为更高的计算精度（例如 BF16）。这将加载后内存降低至 23.682 GB，峰值内存降低至 28.451 GB，推理时间为 13 秒。

将量化与内存优化相结合

量化后端可以与其他 Diffusers 优化技术结合，以进一步降低 VRAM 使用：

  • 模型 CPU 卸载 (enable_model_cpu_offload)： 在 CPU 与 GPU 之间移动整个组件。将 BnB 4 位量化与此方法结合，可将峰值内存降低至 12.383 GB。
  • 分组卸载 (enable_group_offload)： 将内部层组移动到 CPU。将 FP8 分层转型与分组卸载结合，可将加载后内存降低至 9.264 GB，峰值内存降低至 14.232 GB。
  • torch.compile： 通过 PyTorch 2.x 加速执行。虽然它不降低内存，但显著加快推理速度。例如，torchao int4_weight_only 的推理时间在编译后从 109 秒降至 6 秒（尽管编译耗时约 285 秒）。

后端选择指南

根据硬件和性能目标，推荐使用不同的后端：

  • NVIDIA 用户（最简便）： 使用 bitsandbytes 4/8 位。
  • 推理速度： 使用 torchao、GGUF，或结合 torch.compile() 的 bitsandbytes。
  • 硬件灵活性（CPU/MPS）或 FP8： 使用 Quanto。
  • Hopper/Ada 架构： 使用 FP8 分层转型。
  • 已有 GGUF 模型： 通过 from_single_file 加载 GGUF。

Sources