Optimum Quanto

🤗 Optimum Quanto 是 optimum 的 PyTorch 量化后端。

它的设计注重多功能性和简单性：

所有功能在即时模式下可用（适用于不可追踪的模型），
量化后的模型可以放置在任何设备上（包括 CUDA 和 MPS），
自动插入量化和反量化存根，
自动插入量化的函数操作，
自动插入量化的模块（支持的模块列表见下文），
提供从浮点模型到动态再到静态量化模型的无缝工作流程，
序列化兼容 PyTorch 的 weight_only 和 🤗 的 safetensors，
在 CUDA 设备上加速矩阵乘法（int8-int8、fp16-int4、bf16-int8、bf16-int4），
支持 int2、int4、int8 和 float8 权重，
支持 int8 和 float8 激活。

尚未实现的功能：

动态激活平滑，
所有设备上所有混合矩阵乘法的内核，
与 torch compiler（又称 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 pipeline 内的任何子模型，并稍后无缝地将它们包含在另一个 pipeline 中。

这里我们量化 Pixart pipeline 的 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")

之后，我们可以重新加载量化模型并重新创建 pipeline：

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

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)

# 从你的建模代码创建一个空模型并重新量化它
with torch.device('meta'):
  new_model = ...
requantize(new_model, state_dict, quantization_map, device=torch.device('cuda'))