Quanto: PyTorch Optimum 量化后端

量化是一种通过将深度学习模型的权重和激活表示为低精度数据类型（如 8 位整数 int8）而不是常用的 32 位浮点数（float32）来降低计算和内存成本的技术。

减少位数意味着生成的模型需要更少的内存存储，这对于在消费设备上部署大型语言模型至关重要。它还为低位宽数据类型（如 CUDA 设备上的 int8 或 float8 矩阵乘法）实现了特定的优化。

有许多开源库可用于量化 PyTorch 深度学习模型，每个库都提供了非常强大的功能，但通常仅限于特定的模型配置和设备。

此外，尽管它们都基于相同的设计原则，但不幸的是，它们通常彼此不兼容。

今天，我们很高兴地推出 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 激活。

最近的量化方法似乎专注于量化大型语言模型（LLM），而 quanto 旨在为任何模态都可适应的简单量化方案（线性量化、按组量化）提供极其简单的量化原语。

量化工作流程

Quanto 可作为 pip 包使用。

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)

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

5. 重新加载量化模型

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

import json

from safetensors.torch import load_file

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

# Create an empty model from your modeling code and requantize it
with torch.device('meta'):
  new_model = ...
requantize(new_model, state_dict, quantization_map, device=torch.device('cuda'))

有关量化工作流程的实例化，请参阅示例。您还可以查看此 notebook，其中我们展示了如何使用 quanto 量化 BLOOM 模型！

性能

以下是评估 meta-llama/Meta-Llama-3.1-8B 不同量化配置准确性的两个图表。

注意：每组中的第一个条形图始终对应于未量化模型。

这些结果是在不应用任何训练后优化算法（如 hqq 或 AWQ）的情况下获得的。

下图显示了在 NVIDIA A10 GPU 上测量的每 token 延迟。

请继续关注最新结果，因为我们正在不断改进 quanto，加入优化器和优化内核。

有关不同模型架构和配置的详细结果，请参阅 quanto 基准测试。

在 transformers 中的集成

Quanto 已无缝集成到 Hugging Face transformers 库中。您可以通过将 QuantoConfig 传递给 from_pretrained 来量化任何模型！

目前，您需要使用最新版本的 accelerate 以确保完全兼容。

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
)

您可以通过在 QuantoConfig 中简单地传递正确的参数，将权重和/或激活量化为 int8、float8、int4 或 int2。激活可以是 int8 或 float8。对于 float8，您需要拥有与 float8 精度兼容的硬件，否则 quando 在执行矩阵乘法时（仅当权重被量化时）将默默地将权重和激活提升到 torch.float32 或 torch.float16（取决于模型的原始数据类型）。如果您尝试使用 MPS 设备上的 float8，torch 目前会引发错误。

Quanto 与设备无关，这意味着无论您是在 CPU/GPU/MPS（Apple Silicon）上，您都可以量化和运行您的模型。

Quanto 也对 torch.compile 友好。您可以使用 quanto 量化模型，并调用 torch.compile 对模型进行编译以实现更快的生成。如果涉及动态量化（即启用量化感知训练或量化激活），此功能可能无法开箱即用。在使用 transformers 集成时，请确保在创建 QuantoConfig 时将 activations=None。

使用 quanto 还可以量化任何模态的模型！我们演示了如何使用 quanto 将 openai/whisper-large-v3 模型量化为 int8。

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
)

查看此 notebook，了解如何将 quanto 与 transformers 集成正确使用的完整教程！

为 quanto 贡献

非常欢迎对 quanto 的贡献，尤其是在以下领域：

• 针对特定设备的 quanto 操作的优化内核，
• 训练后量化优化器，以恢复量化过程中损失的精度，
• 用于 transformers 或 diffusers 模型的辅助类。