Quanto模型量化：PyTorch内存优化技术详解

原创

用户11764306

发布于 2026-02-16 14:11:44

990

引言

基于Transformer的扩散模型正日益完善，并已被证明能够彻底改变文本到图像的生成模式。Transformer的能力提升了任何模型的可扩展性和性能，同时也增加了模型的复杂性。

"能力越强，责任越大"

在这种情况下，模型复杂性越高，功耗和内存消耗也越大。

例如，使用像Stable Diffusion 3这样的模型运行推理需要巨大的GPU内存，因为它涉及多个组件——文本编码器、扩散主干网络和图像解码器。这种高内存需求给使用消费级GPU的用户带来了阻碍，影响了可访问性和实验性。

模型量化技术应运而生。想象一下，将一个资源密集型模型缩小到更易于管理的大小，同时又不牺牲其有效性。量化就像将高分辨率图像压缩成更紧凑的格式，将模型的参数转换为更低精度的表示形式。这不仅减少了内存使用，还加速了计算，使复杂模型更易于访问和使用。

本文将探讨Quanto的量化工具如何显著提升基于Transformer的扩散 pipeline 的内存效率。

先决条件

AI模型基础理解：熟悉深度学习框架，特别是PyTorch。
量化知识：了解模型量化的基本概念及其在优化AI模型中的目的。
环境设置：可以访问安装了Python的机器，以及PyTorch和Optimum等必要依赖项。
计算资源：用于性能基准测试的GPU环境（例如，某机构的H100）。
示例模型/数据集：用于Quanto测试和实验的预训练模型或数据集。

Quanto简介：一个多功能的PyTorch量化后端

量化是降低深度学习模型计算和内存需求的关键技术，使其更适合在消费类设备上部署。通过使用低精度数据类型（如8位整数）代替32位浮点数，量化不仅减少了内存存储需求，还能针对特定硬件（如CUDA设备上的int8或float8矩阵乘法）进行优化。

Quanto是Optimum的一个新的量化后端，旨在提供一个通用且直接的量化过程。Quanto因其对各种功能的全面支持而脱颖而出，确保了与多样化模型配置和设备的兼容性：

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

虽然许多工具专注于使大型AI模型变小，但Quanto旨在对所有类型的模型都简单实用。

Quanto工作流程

使用pip安装Quanto，请使用以下代码：

!pip install optimum-quanto

量化模型

以下代码将帮助将标准模型转换为量化模型：

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

校准

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

from optimum.quanto import Calibration

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

量化感知训练

如果模型性能受到影响，可以对模型进行几个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()

冻结整数权重

冻结模型时，浮点权重将转换为量化权重。

from optimum.quanto import freeze
freeze(model)

H100基准测试研究

某机构的H100 GPU是一款专为处理要求苛刻的AI任务而设计的高性能显卡，包括用于大型模型（如Transformer和扩散模型）的训练和推理。选择它进行基准测试的原因如下：

顶级性能：H100提供卓越的速度和算力，非常适合处理大型模型所需的复杂操作。
支持FP16：该GPU能有效处理FP16（半精度浮点）计算，从而减少内存使用并加快计算速度，同时不会显著降低精度。
先进的硬件特性：H100支持针对混合精度训练和推理的优化操作，使其成为旨在减小模型大小同时保持性能的量化技术的绝佳选择。

在基准测试研究中，主要关注点是应用Quanto这种新的量化工具于扩散模型。虽然量化在大型语言模型从业者中众所周知，但在扩散模型中并不常用。本研究使用Quanto来探索它是否能在质量几乎没有损失的情况下，为这些模型节省内存。

研究内容如下：

环境设置

CUDA 12.2：使用的CUDA版本，对于在H100 GPU上运行计算至关重要。
PyTorch 2.4.0：用于模型训练和推理的深度学习框架。
Diffusers：用于构建和运行扩散模型的库，从特定提交安装以确保一致性。
Quanto：量化工具，也从特定提交安装。
FP16计算：默认情况下，所有计算都以FP16执行，以减少内存使用并提高速度。但是，为避免潜在的数字不稳定性，不对VAE进行量化。
关键pipeline：本研究从Diffusers中选择了三个基于Transformer的扩散pipeline。选择这些pipeline是因为它们代表了文本到图像生成的最先进模型，使其成为测试Quanto在减少内存使用同时保持模型质量方面有效性的理想选择。
- PixArt-Sigma
- Stable Diffusion 3
- Aura Flow

通过使用H100 GPU并专注于这些关键模型，本研究旨在展示Quanto在弥合LLM和扩散模型之间量化差距方面的潜力，提供显著的内存节省，同时对性能的影响最小。

应用Quanto量化DiffusionPipeline

使用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)

完成后使用pipeline。

image = pipeline("ghibli风格，一个带有城堡的幻想景观").images[0]

（使用FP8扩散Transformer生成的图像）

以下代码可用于量化文本编码器：

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

文本编码器本身也是一个Transformer模型，因此也可以进行量化。通过量化文本编码器和扩散主干网络，可以获得更大的内存节省。

LLM Pipelines

optimum-quanto提供了辅助类来量化、保存和重新加载某机构的量化模型。

下面的代码将使用Transformers库加载预训练语言模型（Meta-Llama-3-8B）。然后，它使用Optimum Quanto中的QuantizedModelForCausalLM类对模型应用量化，特别将模型的权重转换为qint4数据类型。lm_head（输出层）被排除在量化之外以保持其精度，确保最终输出质量保持高水平。

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

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

# 可以使用 save_pretrained 保存量化模型
qmodel.save_pretrained('./Llama-3-8B-quantized')

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

与Transformers的集成

Quanto与某机构的transformers库无缝集成。只需通过QuantoConfig传递模型，即可对任何模型进行量化。

from transformers import AutoModelForCausalLM, AutoTokenizer, QuantoConfig

model_id = "某机构/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上，都可以量化和运行模型。

from transformers import AutoModelForSpeechSeq2Seq

model_id = "某机构/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：在模型中规范化数据。在Quanto中，输出可以被量化，但权重和偏置保持不变。

这种设置使模型变得更小、更快，同时仍保持良好的准确性水平。

研究的最终观察结果

在扩散模型中量化文本编码器（例如Stable Diffusion中使用的编码器）会显著影响性能和内存效率。对于使用三个不同文本编码器的Stable Diffusion 3，关于量化这些编码器的观察结果如下：

量化第二个文本编码器：这种方法效果不佳，可能是由于第二个编码器的特定特性。
量化第一个文本编码器：此选项推荐使用且通常有效，因为第一个编码器的量化在内存节省和性能之间提供了良好的平衡。
量化第三个文本编码器：此方法也推荐使用。量化第三个编码器可以在不显著影响模型性能的情况下带来显著的内存节省。
同时量化第一个和第三个文本编码器：结合量化第一个和第三个文本编码器可以成为最大化内存效率同时保持可接受性能水平的良好策略。

在内存节省方面，下表详细说明了应用每种量化策略可以节省多少内存，特别关注量化不同文本编码器组合的相对优势和权衡。

在所有情况下量化扩散Transformer可确保观察到的内存节省主要来自文本编码器的量化。
当考虑H100或4090等强大GPU时，使用bfloat16可能更快。
qint8由于高效的整数运算和硬件优化，通常推理速度更快。
融合QKV投影会加厚int8内核，通过减少操作次数和利用高效数据处理进一步优化计算。
在H100 GPU上使用qint4与bfloat16结合时，内存使用有所改善，因为qint4每个值仅使用4位，减少了存储权重所需的内存量。然而，这是以增加推理延迟为代价的。这是因为H100 GPU仍然不支持4位整数的计算。尽管权重以压缩的4位格式存储，但实际计算仍以bfloat16执行，这意味着硬件必须处理更复杂的操作，导致处理时间变慢。