基于 Transformer 的扩散模型如今已经相当成熟，事实证明，它在文本到图像生成领域确实带来了碘伏性的改变。Transformer 的能力不仅让模型的可扩展性和性能上了一个台阶，同时也让模型的复杂度水涨船高。

问题来了——模型越复杂，功耗和内存消耗就越大。比如用 Stable Diffusion 3 这类模型做推理，显存需求高得吓人，因为它包含了文本编码器、扩散主干、图像解码器等多个组件。对于手里只有消费级 GPU 的玩家来说，这无疑是个门槛，直接影响了模型的可访问性和实验自由度。

这时候，模型量化就派上用场了。想象一下：把资源密集型模型压缩到更易管理的大小，同时还不牺牲效果。量化就像把一张超高分辨率图片转成更紧凑的格式——把模型的参数换成低精度表示。好处很明显：内存占用降了，计算速度提了，复杂模型也变得亲民多了。

下面要聊的，就是 Quanto 这个量化工具，看看它是怎么让基于 Transformer 的扩散管道在内存效率上实现显著提升的。

Quanto：多功能的 PyTorch 量化后端

量化是降低深度学习模型计算和内存需求的关键手段，特别适合部署到消费级设备上。把 32 位浮点数 (float32) 换成 8 位整数 (int8) 这种低精度数据类型，不仅能降低存储要求，还能针对特定硬件做优化——比如 CUDA 设备上的 int8 或 float8 矩阵乘法。

Quanto 是 Hugging Face 团队开发的一个 Python 库，作为 Optimum 的新量化后端，主打的就是多功能和简洁的量化流程。它之所以突出，是因为对多种功能支持得相当全面，兼容各种模型配置和设备：

• Eager Mode 兼容性：与不可追踪的模型也能无缝协作。
• 设备灵活性：量化模型可以部署在任何设备上，包括 CUDA 和 MPS。
• 自动集成：自动插入量化/反量化存根、功能操作和量化模块。
• 简化的工作流：从浮点模型到动态、静态量化模型，过渡非常顺畅。
• 序列化支持：兼容 PyTorch weight_only 和 Safetensors 格式。
• 加速矩阵乘法：支持 CUDA 设备上的多种量化格式（int8-int8、fp16-int4、bf16-int8、bf16-int4）。
• 广泛支持：可处理 int2、int4、int8 和 float8 的权重和激活。

市面上不少工具都在专注缩小大模型，但 Quanto 的设计理念更倾向于简单通用，适用面更广。

Quanto 工作流程

先用 pip 安装 Quanto：

!pip install optimum-quanto

1. 量化模型

以下代码能把标准模型转成量化模型：

from optimum.quanto import quantize, qint8
quantize(model, weights=qint8, activations=qint8)

2. 校准

Quanto 的校准模式可以确保量化参数根据模型中真实的数据分布进行调整，从而提升量化模型的准确性和效率。

from optimum.quanto import Calibration

with Calibration(momentum=0.9):
    model(samples)

3. 量化感知训练

如果量化后模型性能受到了一点影响，可以花几个 epoch 微调一下，把性能拉回来。

import torch
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
    data, target = data.to(device), target.to(device)
    optimizer.zero_grad()
    output = model(data).dequantize()
    loss = torch.nn.functional.nll_loss(output, target)
    loss.backward()
    optimizer.step()

4. 冻结整数权重

冻结模型时，浮点权重会被转换成量化权重。

from optimum.quanto import freeze
freeze(model)

H100 基准测试研究

H100 GPU 专为高强度的 AI 任务设计，训练和推理大型模型（比如 Transformer 和扩散模型）都不在话下。它被选为基准测试的原因很直接：

• 顶级性能：H100 的速度和算力足以应付文本到图像、文本到视频这类复杂生成管道的运算需求。
• 支持 FP16：高效处理半精度浮点计算，减少内存使用、加快计算速度，同时几乎不牺牲准确率。
• 高级硬件特性：支持混合精度训练和推理的优化操作，是量化技术（在减小模型尺寸的同时保持性能）的理想搭档。

言归正传，这次基准测试的重点，是把 Quanto 这个量化工具应用到扩散模型上。量化在大型语言模型 (LLM) 的圈子里已经很常见了，但在扩散模型领域用得还不多。所以这次就是想看看，Quanto 能不能在扩散模型上同样省下内存，同时尽量不损失生成质量。

研究的主要内容包括：

环境设置

• CUDA 12.2
• PyTorch 2.4.0
• Diffusers：从特定提交安装，保证一致性
• Quanto：同样从特定提交安装
• 默认使用 FP16 计算以减少内存使用并提高速度。不过 VAE（变分自动编码器）没有量化，以避免潜在的数值不稳定性。
• 重点测试三条基于 Transformer 的扩散管道：PixArt-Sigma、Stable Diffusion 3、Aura Flow——它们代表了当前最先进的文本到图像生成模型。

通过 H100 GPU 和这几个关键模型，研究的目标是展示 Quanto 如何在 LLM 和扩散模型之间架起量化的桥梁，实现显著的内存节省，同时把性能影响降到最低。

应用 Quanto 量化扩散管道

用 Quanto 量化模型的过程很直接：先装好必要的包——这里安装 transformers、quanto、torch、sentencepiece。当然，根据实际需要可能还要装其他包。

!pip install transformers==4.35.0
!pip install quanto==0.0.11
!pip install torch==2.1.1
!pip install sentencepiece==0.2.0
!pip install optimum-quanto

然后在需要量化的模块上调用 quantize() 函数，指定要量化的组件。在这个例子里，只量化参数，激活保持不变，参数量化为 FP8 数据类型。最后调用 freeze() 函数，用新的量化参数替换原始参数。

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)

完成后就可以直接使用管道了：

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

图：使用 FP8 中的扩散变压器生成的图像

文本编码器也可以量化：

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

文本编码器本身也是一个 Transformer 模型，所以量化没问题。把文本编码器和扩散主干都量化了，内存能省下不少。

LLM 管道

optimum-quanto 提供了帮助类，用来量化、保存和重新加载 Hugging Face 量化模型。下面的代码用 Transformers 加载预训练语言模型 (Meta-Llama-3-8B)，然后通过 QuantizedModelForCausalLM 类把模型权重转为 qint4 数据类型。注意 lm_head（输出层）被排除在量化之外，以保证最终输出的质量。

from transformers import AutoModelForCausalLM
from optimum.quanto import QuantizedModelForCausalLM, qint4

model = AutoModelForCausalLM.from_pretrained('meta-llama/Meta-Llama-3-8B')
qmodel = QuantizedModelForCausalLM.quantize(model, weights=qint4, exclude='lm_head')

# 量化模型可以通过 sa ve_pretrained 保存
qmodel.sa ve_pretrained('./Llama-3-8B-quantized')

# 使用 from_pretrained 重新加载
from optimum.quanto import QuantizedModelForCausalLM
qmodel = QuantizedModelForCausalLM.from_pretrained('Llama-3-8B-quantized')

与 Transformers 的集成

Quanto 和 Hugging Face Transformers 库的集成非常顺滑。任何模型都可以通过 QuantoConfig 进行量化。

from transformers import AutoModelForCausalLM, AutoTokenizer, QuantoConfig

model_id = "facebook/opt-125m"
tokenizer = AutoTokenizer.from_pretrained(model_id)
quantization_config = QuantoConfig(weights="int8")
quantized_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config= quantization_config
)

有了 Quanto，不管你在 CPU、GPU 还是 MPS（Apple Silicon）上，都能量化和运行你的模型。

from transformers import AutoModelForSpeechSeq2Seq

model_id = "openai/whisper-large-v3"
quanto_config = QuantoConfig(weights="int8")
model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cuda",
    quantization_config=quanto_config
)

Quanto 中的张量

• Quanto 的作用：Quanto 使用特殊的张量（数据结构），把它们修改成更小的数据类型（如 int8 或 float8），让模型更省内存、跑得更快。
• 转换过程：
  • 对于浮点数，直接使用 PyTorch 内置的转换。
  • 对于整数，会对数字进行四舍五入以适配更小的数据类型。
• 准确性：Quanto 会尽量避免过多的极端值（太大或太小），防止这些值扭曲模型性能。
• 优化操作：使用这些较低位宽的数据类型，能让模型跑得更快，因为这类计算本身速度就更快。

Quanto 中的模块

• 替换 Torch 模块：Quanto 可以把常规 PyTorch 组件替换成适合这些较小张量的特殊 Quanto 版本。
• 权重转换：模型还在微调时，权重会动态转换，可能会稍微拖慢速度。但一旦模型最终确定（冻结），转换就停止了。
• 偏差：偏差（另一组模型参数）不会被转换，因为转换它要么会降低准确性，要么需要极高精度，导致效率低下。
• 激活：模型的激活（中间输出）也会转换成较小的数据类型，但只有在校准模型找到最佳转换设置后才会进行。

可量化的模块

• 线性层：全连接层，权重始终量化，偏差保持不变。
• Conv2D 层：卷积层（常用于图像处理），权重始终量化，偏差不量化。
• LayerNorm：对数据进行规范化，输出可以量化，但权重和偏差保持不变。

这种设计让模型变得更小、更快，同时仍然保持良好的准确性。

研究的观察结果

量化扩散模型中的文本编码器（比如 Stable Diffusion 里的）会明显影响性能和内存效率。Stable Diffusion 3 用了三个不同的文本编码器，具体观察结果如下：

1. 量化第二个文本编码器：效果不太好，可能是第二个编码器的特性导致的。
2. 量化第一个文本编码器 (CLIPTextModelWithProjection)：建议这么做，在节省内存和保持性能之间平衡得不错。
3. 量化第三个文本编码器 (T5EncoderModel)：同样建议，量化它能在不明显影响性能的情况下省下大量内存。
4. 量化第一个和第三个文本编码器：把两者都量化，是最大化内存效率、同时维持可接受性能的好策略。

在内存节省方面，不同量化策略能省多少一目了然。下面的表格详细列出了每种策略的节省情况，重点对比了量化不同文本编码器组合之间的利弊和取舍。

• 在所有情况下都量化扩散 Transformer，可以确保观察到的内存节省主要来自文本编码器的量化。
• 对于 H100 或 4090 这类强力 GPU，使用 bfloat16 速度更快。
• 由于整数运算和硬件优化的效率优势，qint8 通常推理速度更快。
• 融合 QKV 投影会加厚 int8 内核，进一步优化计算（减少运算次数、提高数据处理效率）。
• 在 H100 GPU 上使用 bfloat16 搭配 qint4 时，内存使用率有所提高——qint4 每个值只用 4 位，减少了存储权重所需的内存。但这是以增加推理延迟为代价的。原因是 H100 目前还不支持直接用 4 位整数 (int4) 做计算：权重虽然压缩到了 4 位，但实际运算还是要以 bfloat16 执行，硬件需要处理更复杂的操作，导致速度变慢。

结论

Quanto 为 PyTorch 提供了一个功能强大的量化后端，通过把权重转换成更低精度的格式来优化模型性能。借助 qint8、qint4 这些技术，Quanto 能减少内存消耗、加快推理速度。而且 Quanto 在不同设备（CPU、GPU、MPS）上都能工作，兼容各种环境。不过要注意，在 MPS 设备上使用 float8 会报错。

总的来说，Quanto 让深度学习模型的部署更高效了，在内存节省和性能权衡之间找到了不错的平衡点。