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 模型的輔助類。