Quanto: PyTorch 量化工具包

HuggingFace發表於2024-04-09
PyTorch

量化技術透過用低精度資料型別 (如 8 位整型 (int8)) 來表示深度學習模型的權重和啟用，以減少傳統深度學習模型使用 32 位浮點 (float32) 表示權重和啟用所帶來的計算和記憶體開銷。

減少位寬意味著模型的記憶體佔用更低，這對在消費裝置上部署大語言模型至關重要。量化技術也使得我們可以針對較低位寬資料型別進行特殊的計算最佳化，例如 CUDA 裝置有針對 int8 或 float8 矩陣乘法的硬體最佳化。

市面上有許多可用於量化 PyTorch 深度學習模型的開源庫，它們各有特色及侷限。通常來講，每個庫都僅實現了針對特定模型或裝置的特性，因而普適性不強。此外，儘管各個庫的設計原理大致相同，但不幸的是，它們彼此之間卻互不相容。

因此，quanto 庫應運而出，其旨在提供一個多功能的 PyTorch 量化工具包。目前 quanto 包含如下特性:

  • 在 eager 模式下可用 (適用於無法成圖的模型)，
  • 生成的量化模型可以執行於任何裝置 (包括 CUDA 裝置和 MPS 裝置) 上，
  • 自動插入量化和反量化結點，
  • 自動插入量化後的 torch.nn.functional 運算元，
  • 自動插入量化後的 torch.nn 模組 (具體支援列表見下文)，
  • 提供無縫的模型量化工作流，支援包含靜態量化、動態量化在內的多種模型量化方案，
  • 支援將量化模型序列化為 state_dict ，
  • 不僅支援 int8 權重，還支援 int2 以及 int4 ，
  • 不僅支援 int8 啟用，還支援 float8 。

最近，出現了很多僅專注於大語言模型 (LLM) 的量化演算法，而 quanto 的目標為那些適用於任何模態的、易用的量化方案 (如線性量化，分組量化等) 提供簡單易用的量化原語。

我們無意取代其他量化庫，而是想透過新演算法的實現門檻來促進創新，使得大家能夠輕鬆地實現新模組，抑或是輕鬆組合現有模組來實現新演算法。

毫無疑問，量化很困難。當前，如要實現模型的無縫量化，需要大家對 PyTorch 的內部結構有深入瞭解。但不用擔心，quanto 的目標就是為你完成大部分繁重的工作，以便你可以集中精力在最重要的事情上，即: 探索低位元 AI 從而找出惠及 GPU 窮人的解決方案。

量化工作流

大家可以 pip 安裝 quanto 包。

pip install quanto

quanto 沒有對動態和靜態量化進行明確區分。因為靜態量化可以首先對模型進行動態量化，隨後再將權重 凍結 為靜態值的方式來完成。

典型的量化工作流包括以下步驟:

1. 量化

將標準浮點模型轉換為動態量化模型。

quantize(model, weights=quanto.qint8, activations=quanto.qint8)

此時，我們會對模型的浮點權重進行動態量化以用於後續推理。

2. 校準 (如果上一步未量化啟用，則可選)

quanto 支援校準模式。在校準過程中，我們會給量化模型傳一些代表性樣本，並在此過程中記錄各運算元啟用的統計資訊 (如取值範圍)。

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

上述程式碼會自動使能量化模組的啟用量化功能。

3. 微調，即量化感知訓練 (可選)

如果模型的效能下降太多，可以嘗試將其微調幾輪以恢復原浮點模型的效能。

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. 凍結整型權重

模型凍結後，其浮點權重將替換為量化後的整型權重。

freeze(model)

請參閱 該例 以深入瞭解量化工作流程。你還可以檢視此 notebook，其提供了一個完整的用 quanto 量化 BLOOM 模型的例子。

效果

下面我們列出了一些初步結果，我們還在緊鑼密鼓地更新以進一步提高量化模型的準確性和速度。但從這些初步結果中，我們仍能看出 quanto 的巨大潛力。

下面兩幅圖評估了 mistralai/Mistral-7B-v0.1 在不同的量化引數下的準確度。注意: 每組的第一根柱子均表示非量化模型。

mistralai/Mistral-7B-v0.1 在 Lambada 資料集上的預測準確度

mistralai/Mistral-7B-v0.1 在 Lambada 資料集上的預測準確度

上述結果均未使用任何高階訓後量化演算法 (如 hqq 或 AWQ)。

下圖給出了在英偉達 A100 GPU 上測到的詞元延遲。

mistralai/Mistral-7B-v0.1 平均詞元延遲

這些測試結果都尚未利用任何最佳化的矩陣乘法運算元。可以看到，量化位寬越低，開銷越大。我們正在持續改進 quanto，以增加更多的最佳化器和最佳化運算元，請持續關注我們的效能演進。

請參閱 quanto 基準測試 以瞭解在不同模型架構及配置下的詳細結果。

整合進 transformers

我們已將 quanto 無縫整合至 Hugging Face transformers 庫中。你可以透過給 from_pretrained API 傳 QuantoConfig 引數來對任何模型進行量化！

目前，你需要使用最新版本的 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 精度的硬體，否則當執行 matmul (僅當量化權重時) 時，我們會預設將權重和啟用都轉成 torch.float32 或 torch.float16 (具體視模型的原始精度而定) 再計算。目前 MPS 裝置不支援 float8 ， torch 會直接丟擲錯誤。

quanto 與裝置無關，這意味著無論用的是 CPU/GPU 還是 MPS (Apple 的晶片)，你都可以對模型進行量化並執行它。

quanto 也可與 torch.compile 結合使用。你可以先用 quanto 量化模型，然後用 torch.compile 來編譯它以加快其推理速度。如果涉及動態量化 (即使用量化感知訓練或對啟用進行動態量化)，該功能可能無法開箱即用。因此，請確保在使用 transformers API 建立 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，以詳細瞭解如何在 transformers 中正確使用 quanto ！

實現細節

量化張量

quanto 的核心是一些 Tensor 子類，其主要做下面兩件事:

  • 將源張量按最優比例 投影至給定量化資料型別的取值範圍內。
  • 將投影后的值對映至目標資料型別。

當目標型別是浮點型時，對映由 PyTorch 原生轉換介面 (即 Tensor.to() ) 完成。而當目標型別是整型時，對映可以用一個簡單的舍入操作 (即 torch.round() ) 來完成。

投影的目標是提高資料型別轉換的精確度，具體可以透過最小化以下兩個值來達成:

  • 飽和值的個數 (即有多少個數最終對映為目標資料型別的最小值/最大值)，
  • 歸零值的個數 (即有多少個數因為小於目標資料型別可以表示的最小數字，所以被對映成了 0)。

為了提高效率起見， 8 位元 量化時，我們使用對稱投影，即以零點為中心進行投影。一般而言，對稱量化張量與許多標準運算元相容。

在使用較低位寬的量化 (如 int2 或 int4 ) 時，一般使用的是仿射投影。此時，會多一個 zeropoint 引數以對齊投影值和原值的零點。這種方法對量化範圍的覆蓋度會好些。仿射量化張量通常更難與標準運算元相容，因此一般需要自定義很多運算元。

量化 torch.nn 模組

quanto 實現了一種通用機制，以用能夠處理 quanto 張量的 quanto 模組替換相應的 torch 模組 ( torch.nn.Module )。

quanto 模組會動態對 weights 進行資料型別轉換，直至模型被凍結，這在一定程度上會減慢推理速度，但如果需要微調模型 (即量化感知訓練)，則這麼做是需要的。

此外，我們並未量化 bias 引數，因為它們比 weights 小得多，並且對加法進行量化很難獲得太多加速。

我們動態地將啟用量化至固定取值範圍 (預設範圍為 [-1, 1] )，並透過校準過程決定最佳的比例 (使用二階動量更新法)。

我們支援以下模組的量化版:

  • Linear (QLinear)。僅量化權重，不量化偏置。輸入和輸出可量化。
  • Conv2d (QConv2D)。僅量化權重，不量化偏置。輸入和輸出可量化。
  • LayerNorm。權重和偏至均 不 量化。輸出可量化。

定製運算元

得益於 PyTorch 出色的排程機制，quanto 支援在 transformers 或 diffusers 的模型中最常用的函式，無需過多修改模型程式碼即可啟用量化張量。

大多數“排程”功能可透過標準的 PyTorch API 的組合來完成。但一些複雜的函式仍需要使用 torch.ops.quanto 名稱空間下的自定義操作。其中一個例子是低位寬的融合矩陣乘法。

訓後量化最佳化

quanto 中尚未支援高階的訓後量化演算法，但該庫足夠通用，因此與大多數 PTQ 最佳化演算法相容，如 hqq、AWQ

展望未來，我們計劃無縫整合這些最流行的演算法。

為 Quanto 作出貢獻

我們非常歡迎大家對 quanto 作出貢獻，尤其歡迎以下幾類貢獻:

  • 實現更多針對特定裝置的 quanto 最佳化運算元，
  • 支援更多的 PTQ 最佳化演算法，
  • 擴大量化張量可排程操作的覆蓋面。

英文原文: https://hf.co/blog/quanto-introduction
原文作者: David Corvoysier，Younes Belkada，Marc Sun
譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。

相關文章