Optimum Quanto

该项目目前处于维护模式。我们仅接受用于修复小错误、改进文档及其他维护任务的拉取请求。重大的新功能或破坏性更改不太可能被合并。对于生产就绪的量化功能或活跃的开发工作，请考虑使用其他项目，例如 bitsandbytes 或 torchAO。

🤗 Optimum Quanto 是一个用于 optimum 的 PyTorch 量化后端。

它在设计时充分考虑了通用性和简洁性：

• 所有功能均可在 eager 模式下使用（适用于不可跟踪的模型）；
• 量化后的模型可以部署到任何设备上（包括 CUDA 和 MPS）；
• 自动插入量化和反量化的桩代码；
• 自动插入量化后的函数操作；
• 自动插入量化模块（支持的模块列表见下文）；
• 提供从浮点模型到动态量化再到静态量化的无缝工作流；
• 序列化与 PyTorch 的 weight_only 格式及 🤗 safetensors 兼容；
• 在 CUDA 设备上加速矩阵乘法运算（int8-int8、fp16-int4、bf16-int8、bf16-int4）；
• 支持 int2、int4、int8 和 float8 权重；
• 支持 int8 和 float8 激活值。

尚未实现的功能：

• 动态激活平滑处理；
• 所有混合精度矩阵乘法的内核在所有设备上的实现；
• 与 PyTorch 编译器（即 dynamo）的兼容性。

性能

简而言之：

• 准确性：使用 int8/float8 权重和 float8 激活值编译的模型，其性能非常接近全精度模型；
• 延迟：在存在优化内核的情况下，仅对模型权重进行量化时，量化模型的推理延迟可与全精度模型相媲美；
• 设备内存：大约为浮点位数除以整数位数。

以下段落仅为示例。有关按模型用例划分的详细结果，请参阅 bench 文件夹。

meta-llama/Meta-Llama-3.1-8B

安装

Optimum Quanto 可通过 pip 包管理器安装。

pip install optimum-quanto

Hugging Face 模型的量化流程

optimum-quanto 提供了辅助类，用于量化、保存和重新加载 Hugging Face 量化模型。

LLM 模型

第一步是量化模型：

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')

注意：量化后的模型权重将被冻结。如果希望保持权重未冻结以便继续训练，则需要直接使用 optimum.quanto.quantize 方法。

量化后的模型可以使用 save_pretrained 方法保存：

qmodel.save_pretrained('./Llama-3-8B-quantized')

之后可以通过 from_pretrained 方法重新加载：

from optimum.quanto import QuantizedModelForCausalLM

qmodel = QuantizedModelForCausalLM.from_pretrained('Llama-3-8B-quantized')

Diffusers 模型

您可以对 Diffusers 流水线中的任意子模型进行量化，并随后无缝地将其集成到另一个流水线中。这里我们对 Pixart 流水线中的 transformer 进行量化：

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")

稍后，我们可以重新加载量化后的模型并重建流水线：

from diffusers import PixArtTransformer2DModel
from optimum.quanto import QuantizedPixArtTransformer2DModel

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

适用于原生 PyTorch 模型的量化工作流（低层级 API）

使用低层级 Quanto API 时，需要注意的一点是：默认情况下，模型权重会进行动态量化；必须显式调用“冻结”操作来固定已量化的权重。

典型的量化工作流包括以下步骤：

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)

5. 序列化量化模型

量化后的模型权重可以序列化为 state_dict，并保存到文件中。支持 pickle 和 safetensors（推荐）两种格式。

from safetensors.torch import save_file

save_file(model.state_dict(), 'model.safetensors')

为了能够重新加载这些权重，还需要存储量化模型的量化映射。

import json

from optimum.quanto import quantization_map

with open('quantization_map.json', 'w') as f:
  json.dump(quantization_map(model), f)

6. 重新加载量化模型

可以通过 state_dict 和 quantization_map 使用 requantize 辅助函数重新加载序列化的量化模型。请注意，首先需要实例化一个空模型。

import json

from safetensors.torch import load_file
from optimum.quanto import requantize

state_dict = load_file('model.safetensors')
with open('quantization_map.json', 'r') as f:
  quantization_map = json.load(f)

# 根据你的建模代码创建一个空模型，并对其进行重新量化
with torch.device('meta'):
  new_model = ...
requantize(new_model, state_dict, quantization_map, device=torch.device('cuda'))

有关该工作流的具体实现，请参阅 示例。

设计概览

张量

Quanto 的核心是一个张量子类，它对应于：

• 将源张量投影到给定目标类型的最优范围内；
• 将投影后的值映射到目标类型。

对于浮点目标类型，映射由原生 PyTorch 的类型转换完成（即 Tensor.to()）。

对于整数目标类型，映射则是一个简单的四舍五入操作（即 torch.round()）。

投影的目标是通过减少以下情况的数量来提高转换精度：

• 饱和值（即被映射到目标类型的最小值或最大值）；
• 被清零的值（因为它们低于目标类型所能表示的最小数值）。

投影方式根据数据类型而有所不同：对于 int8 和 float8，采用逐张量或逐通道对称投影；而对于更低的位宽，则采用分组仿射投影（带有偏移量或“零点”）。

使用较低位宽表示的一个好处是，可以利用针对目标类型加速的运算，通常比高精度版本更快。

Quanto 不支持使用混合目标类型对张量进行转换。

模块

Quanto 提供了一种通用机制，用于将 torch 模块替换为能够处理 Quanto 张量的 optimum-quanto 模块。

optimum-quanto 模块会在模型冻结之前动态转换其权重，这会略微降低推理速度，但如果需要对模型进行微调，则这是必要的。

权重通常按通道沿第一维度（输出特征）进行量化。

偏置不进行转换，以保持典型 addmm 运算的精度。

解释：为了与未量化算术运算保持一致，偏置需要以等于输入和权重尺度乘积的尺度进行量化，这会导致尺度小得离谱，从而需要非常高的位宽才能避免截断。例如，当输入和权重为 int8 时，偏置至少需要以 12 位进行量化，即 int16。由于大多数偏置目前都是 float16，这样做实际上并无意义。

激活值则按张量动态量化，使用静态尺度（默认范围为 [-1, 1]）。

为保持精度，需要对模型进行校准，以评估最佳的激活尺度（使用动量参数）。

以下模块可以被量化：

• Linear（QLinear）。权重始终量化，偏置不量化。输入和输出均可量化。
• Conv2d（QConv2D）。权重始终量化，偏置不量化。输入和输出均可量化。
• LayerNorm，权重和偏置 不 量化。输出可以量化。

量化激活值时需避免的陷阱

激活值始终按张量整体进行量化，这是因为模型图中的大多数线性代数运算都不支持按轴输入：你无法对不同基数表示的数值进行加法运算（“不能把苹果和橘子相加”）。

与此相反，参与矩阵乘法的权重则始终沿其第一个轴进行量化，这是因为所有输出特征都是独立计算的。

无论如何，即使激活值已被量化，量化后的矩阵乘法结果仍会先被反量化，原因如下：

• 结果的累加值采用的位宽远高于激活值的位宽（通常为 int32 或 float32，而激活值通常是 int8 或 float8)；
• 这些值还可能与一个 float 类型的偏置项相加。

如果对应的张量中包含较大的异常值，将激活值按张量整体量化为 int8 格式，可能会导致严重的量化误差。典型情况下，这会使量化后的张量中大部分值都被置为零，仅保留那些异常值。

解决这一问题的一种方法是像 SmoothQuant 所示，对激活值进行静态平滑处理。你可以在 external/smoothquant 目录下找到用于平滑某些模型架构的脚本。

更好的选择则是使用 float8 来表示激活值。

Optimum Quanto 快速上手指南

Optimum Quanto 是 Hugging Face optimum 库的 PyTorch 量化后端，旨在提供灵活且简单的模型量化方案。它支持动态和静态量化，兼容多种设备（CUDA、MPS），并支持 int2/4/8 及 float8 等多种精度。

注意：该项目目前处于维护模式，仅接受次要 bug 修复和文档改进。如需生产级量化功能或活跃开发的新特性，建议考虑 bitsandbytes 或 torchAO。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS, Windows
• Python 版本：3.8 及以上
• PyTorch：已安装与您的硬件（CPU/CUDA/MPS）匹配的 PyTorch 版本
• 依赖库：
  • transformers (用于处理 Hugging Face 模型)
  • diffusers (可选，用于处理图像生成模型)
  • safetensors (推荐，用于高效保存量化权重)

安装步骤

您可以直接通过 pip 安装 Optimum Quanto。国内开发者建议使用清华或阿里镜像源以加速下载。

使用官方源安装：

pip install optimum-quanto

使用国内镜像源安装（推荐）：

pip install optimum-quanto -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

Optimum Quanto 提供了高层 API 以便快速量化 Hugging Face 模型，也提供了底层 API 用于自定义 PyTorch 模型的量化流程。

场景一：快速量化大语言模型 (LLM)

这是最简单的使用方式，适用于 transformers 加载的因果语言模型。以下示例将 Llama-3-8B 模型量化为 int4 精度。

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

# 1. 加载原始浮点模型
model = AutoModelForCausalLM.from_pretrained('meta-llama/Meta-Llama-3-8B')

# 2. 执行量化 (weights=qint4 表示权重量化为 4-bit 整数)
# exclude='lm_head' 通常用于保持输出层精度以提升生成质量
qmodel = QuantizedModelForCausalLM.quantize(model, weights=qint4, exclude='lm_head')

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

# 4. 后续加载量化模型
# qmodel = QuantizedModelForCausalLM.from_pretrained('./Llama-3-8B-quantized')

场景二：量化 Diffusers 模型

您也可以量化图像生成管道中的子模块（如 Transformer）。以下示例将 PixArt 模型的 transformer 部分量化为 float8 精度。

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

# 1. 加载子模型
model = PixArtTransformer2DModel.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", 
    subfolder="transformer"
)

# 2. 执行量化 (weights=qfloat8)
qmodel = QuantizedPixArtTransformer2DModel.quantize(model, weights=qfloat8)
qmodel.save_pretrained("./pixart-sigma-fp8")

# 3. 重新加载并集成到管道中
transformer = QuantizedPixArtTransformer2DModel.from_pretrained("./pixart-sigma-fp8")
transformer.to(device="cuda")

pipe = PixArtSigmaPipeline.from_pretrained(
  "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
  transformer=None, # 不加载原始 transformer
  torch_dtype=torch.float16,
).to("cuda")

pipe.transformer = transformer

场景三：原生 PyTorch 模型量化 (底层 API)

对于非 Hugging Face 架构的原生 PyTorch 模型，需要遵循“量化 -> 校准 (可选) -> 冻结 -> 保存”的标准工作流。

import torch
from optimum.quanto import quantize, freeze, Calibration, qint8

# 假设 model 是一个标准的 torch.nn.Module
# 1. 动态量化：将权重和激活值设置为动态量化模式
quantize(model, weights=qint8, activations=qint8)

# 2. 校准 (可选但推荐)：如果量化了激活值，需通过代表性数据确定缩放比例
# 这能显著提升量化后的精度
with Calibration(momentum=0.9):
    # 传入一些代表性样本进行推理
    # model(samples) 
    pass

# 3. 冻结模型：将浮点权重永久替换为整数量化权重，提升推理速度
freeze(model)

# 4. 保存状态
# 推荐使用 safetensors 保存权重
from safetensors.torch import save_file
save_file(model.state_dict(), 'model.safetensors')

# 同时必须保存量化映射表以便重新加载
import json
from optimum.quanto import quantization_map

with open('quantization_map.json', 'w') as f:
  json.dump(quantization_map(model), f)