量化概述

首先，我們想介紹一下 torchao 棧

Quantization Algorithms/Flows: weight only/dynamic/static quantization, hqq, awq, gptq etc.
---------------------------------------------------------------------------------------------
        Quantized Tensors (derived dtypes): AffineQuantizedTensor, CodebookQuantizedTensor
---------------------------------------------------------------------------------------------
  Quantization Primitive Ops/Efficient Kernels: matmul, quantize, dequantize
---------------------------------------------------------------------------------------------
            Basic dtypes: uint1-uint7, int1-int8, float3-float8

任何量化演算法都將使用上述棧中的一些元件，例如 int4_weight_only 量化使用：(1) 僅權重量化流程 (2) tinygemm bf16 啟用 + int4 權重核 和 量化原語操作 (3) AffineQuantizedTensor 張量子類以及 TensorCoreTiledLayout (4) torch.uint4 dtype（目前透過 quant_min/quant_max 模擬）

注意：我們還將在量化張量 (Quantized Tensors) 部分討論如何將稀疏性與量化結合

基本資料型別 (Basic DTypes)

dtype 有點過載，這裡說的基本資料型別 (basic dtype) 是指那些沒有額外元資料也能理解的資料型別（例如，當人們呼叫 torch.empty(.., dtype) 時能理解），更多詳情請查閱：dev-discuss.pytorch.org/t/supporting-new-dtypes-in-pytorch/1833

無論我們進行何種量化，最終都將使用一些低精度資料型別 (low precision dtypes) 來表示量化資料，torchao 中我們打算支援的資料型別包括

• torch.uint1 到 torch.uint8 在 pytorch 2.3 及更高版本中可用

• torch.int1 到 torch.int8 在 pytorch 2.6 及更高版本中可用

• torch.float3_e2_m0, torch.float4_e2_m1, torch.float4_e3_m0, torch.float5_e2_m2, torch.float5_e3_m1, torch.float6_e2_m3, torch.float6_e3_m2, torch.float8_e4m3fn, torch.float8_e5m2, torch.float8_e4m3fnuz, torch.float8_e5m2fnuz （float8 已新增到 torch，如果 float4 和 float6 流行起來，我們也計劃將它們新增到 torch）

注意，上述部分型別目前僅為原型。當它們變得流行並獲得硬體支援時，我們將考慮將其新增到 PyTorch 核心中。

當前支援 (Current Support)

在實際實現方面，有兩個部分：1). 在 PyTorch 中，我們需要將 dtype 新增到 torch.dtype，例如 torch.uint2，示例：pytorch/pytorch#117208，但這只是佔位符，以便我們可以使用 torch.uint2。2). 在 PyTorch 之外（例如在 torchao 中），我們使用張量子類實現這些 dtype 的張量操作，同時還需要一種標準的打包格式 (packing format)。

在 PyTorch 中新增佔位符 dtype

如 dev-discuss.pytorch.org/t/supporting-new-dtypes-in-pytorch/1833 中所述，在 PyTorch 中新增 dtype 的標準是它已被廣泛採用。對於上面提到的基本資料型別，PyTorch 中支援的包括

• torch.uint1 到 torch.uint8, torch.int1 到 torch.int8, torch.float8_e4m3fn, torch.float8_e5m2, torch.float8_e4m3fnuz, torch.float8_e5m2fnuz

對於其他型別，我們計劃等到有更多證據表明其被廣泛採用並獲得硬體支援時再進行新增。

使用張量子類實現這些 dtype 的張量操作

為此，要求是我們確定一種“標準”打包格式 (packing format)，並且希望它易於高效實現，但對於 uintx 和 floatx，我們尚未整合足夠的核心來決定這一點。因此，當前的 打包實現 尚非最終版本。在將更多 uintx、intx 和 floatx 核心整合到 torchao 後，我們可以重新討論這一點。

將張量子類整合到 PyTorch 原生工廠函式中

之後，我們可以將工廠函式與張量子類連線起來，例如： torch.empty(..., dtype=torch.int4, ...) 可以建立一個 Int4Tensor 張量子類，其打包格式在之前的步驟中確定。

量化原語操作 (Quantization Primitive Ops)

量化原語操作 (Quantization primitive ops) 是指用於在低精度量化張量和高精度張量之間進行轉換的運算子。我們將主要有以下量化原語運算子：choose_qparams ops：根據原始張量選擇量化引數的運算子，通常用於動態量化，例如 affine 量化的 scale 和 zero_point；quantize op：根據量化引數將原始高精度張量量化為前一節提到的 dtype 的低精度張量；dequantize op：根據量化引數將低精度張量反量化回高精度張量。

上述操作可能存在變體以適應特定用例，例如對於靜態量化，我們可能擁有 choose_qparams_affine_with_min_max，它將根據觀察過程得出的 min/max 值來選擇量化引數。

高效核心 (Efficient kernels)

我們還將擁有適用於低精度張量的高效核心，例如

_weight_int4pack_mm tinygemm int4 核 (bf16 啟用 + int4 權重) int_matmul 接受兩個 int8 張量並輸出一個 int32 張量 int_scaled_matmul 執行 matmul 並對結果應用 scale。

注意：我們也可以依靠 torch.compile 來生成核心（透過 triton），例如當前的 int8 僅權重量化 核心 僅依賴 torch.compile 來獲得加速。在這種情況下，沒有與量化型別對應的特定“高效核心”。

量化張量（派生資料型別）(Quantized Tensors)

在基本資料型別、量化原語運算子和高效核心的基礎上，我們可以將所有這些組合起來，透過繼承 torch.Tensor 構建一個量化（低精度）張量 (Quantized Tensor)，該張量可以由高精度張量和一些用於配置使用者所需特定量化的引數構造。我們也可以將此類稱為派生資料型別 (derived dtypes)，因為它可以用基本資料型別的張量和一些額外的元資料（如 scale）表示。

torchao 中現有的示例是 AffineQuantizedTensor，這意味著低精度張量透過仿射對映從高精度張量量化而來，即：low_precision_val = high_precision_val / scale + zero_point，其中 scale/zero_point 是可以透過量化原語操作或透過某種最佳化過程計算出的量化引數。仿射量化 (Affine quantization) 是一種非常常見的量化型別，因為當我們嘗試從高精度值對映到低精度值時，進行仿射變換 (high_preicsion_val / scale + zero_point) 非常直接。另一種常見的量化型別，尤其對於較低位寬（例如低於 4 位）的量化，是基於碼本 / 查詢表的量化。

佈局和 TensorImpl (Layout and TensorImpl)

原生張量有一份硬編碼的 佈局 選擇列表，最常見的是 strided 佈局 (strided layout)，它提供了儲存的分步式、多維檢視；我們還有一些 sparse 和 mkldnn 佈局。

以 sparse COO 張量 為例，它擁有 torch.sparse_coo 佈局 (layout)，以及 SparseTensorImpl，後者改變了張量的儲存方式。

將張量打包成不同格式的想法與佈局 (layout) 概念非常吻合，這就是我們希望將其用於打包的原因。我們可以使用 Layout 來表示不同型別的打包格式，並使用 TensorImpl 來實現不同的儲存格式。在 Python 級別的張量子類中，無需修改 C++ PyTorch 核心程式碼即可新增以打包格式儲存張量的新 TensorImpl。

例如，對於 _weight_int4pack_mm，我們需要將權重打包成對 Tensor Core 友好的格式，我們稱之為 TensorCoreTiledLayout。我們為量化張量新增一個 tensor_impl 來儲存打包（或未打包）的權重，並使用 layout 來儲存與打包相關的不同引數

class AffineQuantizedTensor(...):
  # tensor_impl is also implemented with tensor subclass
  tensor_impl: torch.Tensor

  # to not conflict with existing layout property, we use `_layout`
  @property
  def _layout(self) -> Layout:
      return self.tensor_impl._layout

注意，佈局 (layout) 不僅是用於自定義資料表示的抽象，它也用於描述 TensorImpl 如何與不同的運算子互動，例如 transpose、quantized_linear，但運算子的語義應保持不變。

量化 + 稀疏張量也可以透過佈局 (Layout) 抽象來支援，例如，int4 僅權重量化 + 稀疏。我們還提供了一些通用工具，幫助人們為量化張量新增不同的佈局，請查閱下面的開發者指南以獲取程式碼示例。

量化演算法/流程 (Quantization Algorithms/Flows)

在棧的最頂層是最終的量化演算法和量化流程。傳統上我們有僅權重量化 (weight only quantization)、動態量化 (dynamic quantization) 和靜態量化 (static quantization)，但現在我們也看到更多型別的量化正在出現。

為了演示目的，假設在之前的步驟之後我們定義了 AffineQuantizedTensor 和 to_affine_quantized 工廠函式。為簡單起見，假設 to_affine_quantized 接受一個高精度浮點張量和一個 target_dtype（例如 torch.int8），並將其轉換為具有相應 dtype 的 AffineQuantizedTensor。

注意：以下內容均為概念性解釋，有關我們提供的工具和示例的更詳細介紹，請參見 Tensor Subclass Developer Guide 部分。

僅權重量化 (Weight Only Quantization)

這是最簡單的量化形式，很容易將僅權重量化應用於模型，特別是當我們擁有量化張量 (Quantized Tensor) 時。我們只需執行以下操作：

linear_module.weight = torch.nn.Parameter(to_affine_quantized_intx(linear_module.weight, …), requires_grad=False))

將上述操作應用於模型中的所有線性模組，即可獲得僅權重量化模型。

動態啟用和權重量化 (Dynamic Activation and Weight Quantization)

這之前被稱為“動態量化”(dynamic quantization)，但它意味著我們在執行時動態量化啟用，同時也量化權重。與僅權重量化相比，主要問題是如何將量化應用於啟用。在 torchao 中，我們常用的模式是在量化權重的頂部應用 to_linear_activation_quantized：

quantized_weight = to_affine_quantized(linear_module.weight) activation_and_weight_quantized = to_linear_activation_quantized(quantized_weight) linear_module.weight = torch.nn.Parameter(activation_and_weight_quantized, requires_grad=False))

to_linear_activation_quantized 用於將量化應用於啟用，它接受一個 input_quant_func，該函式將量化啟用和原始權重，並且在執行時遇到 F.linear 操作時，它將把儲存的 input_qunat_func 應用於啟用，並重新分派給使用量化啟用和權重的 F.linear。

如果上述方法不起作用，使用者也可以進行模組替換 (module swaps)，或使用 torch.fx.symbolic_trace() 獲取一個可以 修改 的跟蹤模組。

但優先推薦使用張量子類，因為它更容易進行序列化/反序列化。如果使用張量子類支援動態量化，我們可以直接載入量化權重，無需對模型進行進一步準備。否則，在載入量化權重之前，需要先進行模組替換或其他模型修改。

靜態啟用和權重量化 (Static Activation Quantization and Weight Quantization)

靜態量化 (Static quantization) 意味著啟用是靜態量化的，而不是在執行時動態量化。在流程方面，靜態量化需要使用樣本資料進行校準 (calibration)，以便我們確定合適的量化引數。

從高層次上看，靜態量化有三個步驟：(1) 插入觀察者 (insert observers) (2) 校準 (calibration) (3) 量化模型 (quantize the model)

插入觀察者 (Insert Observers)

在插入觀察者這一步，我們需要向運算子的輸入（和輸出）啟用以及權重新增觀察者模組 (observer modules)，以收集張量的統計資訊。因此，我們需要解決兩個問題：如何定義觀察者模組？如何將觀察者模組新增到模型中？

如何定義觀察者模組

觀察者是特定的，取決於：(1) 量化型別（例如仿射量化 (affine quantization)，基於查詢表的量化）(2) 我們希望跟蹤的統計資訊型別，例如 min max 觀察者 (min max observer)，移動平均觀察者 (moving average observer)。

通常，觀察者模組應定義 forward 和 calculate_qparams

對於仿射量化，我們定義了 AffineQuantizedMinMaxObserver，它根據仿射量化的粒度記錄 min_val/max_val，並定義瞭如何基於記錄的統計資訊計算 calculate_qparams。

如何將觀察者模組新增到模型中

1. 使用張量子類 如果您只對量化線性運算子感興趣，可以使用 linear activation weight observer，我們還有一個相應的 insert_observer_ API，用於處理修改線性層的權重。

2. 模組替換 (Module Swap) 另外，您也可以定義並使用 ObservedLinear 模組（或其他模組型別），並將未觀察的模組替換為已觀察的模組

校準 (Calibration)

校準步驟通常很簡單，通常我們只需讓模型執行透過校準資料集 (calibration dataset)。對於更復雜的校準（例如，我們記錄所有輸入並基於所有輸入進行最佳化），我們將在下一節中介紹其中一些。

量化 (Quantize)

我們可以重用 quantize_ API，但提供一個不同的 apply_tensor_subclass 函式，將已觀察的線性模組轉換為帶有量化權重和靜態量化輸入啟用的線性模組，這可以透過與動態量化相同的方式完成（使用 to_linear_activation_quantized），請參閱 示例。

另外，使用者也可以進行 模組替換。

其他量化流程 (Other Quantization Flows)

對於不屬於上述任何一種的其他量化流程/演算法，我們也打算提供常見模式的示例。例如，類似 GPTQ 的量化流程 由 Autoround 採用，它使用 MultiTensor 和模組 hook 來最佳化模組。

如果您正在研究新的量化演算法/流程，並且不確定如何在 PyTorch 原生方式下實現，請隨時提出 issue，描述您的演算法的工作方式，我們可以幫助提供實現細節方面的建議。

訓練 (Training)

上述流程主要側重於推理 (inference)，但低位寬資料型別張量 (low bit dtype Tensors) 也可以用於訓練 (training)。

量化感知訓練 (Quantization Aware Training)

待辦 (TODO)

低位寬最佳化器 (Low Bit Optimizers)

目前我們有一些低位寬最佳化器 (low bit optimizers) 的原型：main/torchao/prototype/low_bit_optim，實現了特定型別的 4 位、8 位和 float8 最佳化器，並且可以與 FSDP 組合使用（帶有查詢表量化）。

量化訓練 (Quantized Training)

與低位寬最佳化器類似，我們在 main/torchao/prototype/quantized_training 中提供了量化訓練 (quantized training) 原型，我們可以擴充套件 AffineQuantizedTensor 以支援訓練。初步支援正在進行中，但還需要大量後續工作，包括使其適用於不同的核心等。

您還可以檢視 量化訓練 (Quantized Training) 的教程，該教程介紹瞭如何使 dtype 張量子類可訓練。

案例研究：torchao 中的 int4 僅權重量化是如何工作的？

為了將所有內容串聯起來，這裡更詳細地講解一下 int4 僅權重量化在 torchao 中的實現方式。

量化流程： quantize_(model, int4_weight_only())

• 發生的情況是： linear.weight = torch.nn.Parameter(to_affine_quantized_intx(linear.weight, …), requires_grad=False))

• 量化原語操作：呼叫 choose_qparams 和 quantize_affine 來量化張量

• 量化後的張量將是 AffineQuantizedTensor，這是一個帶有派生資料型別（例如帶有 scale 和 zero_point 的 int4）的量化張量

• 打包運算子 _convert_weight_to_int4pack，用於將量化權重打包以實現高效執行

模型執行期間： model(input)

• 對輸入和打包的權重呼叫 torch.ops.aten._weight_int4pack_mm

量化期間 (During Quantization)

首先我們從 API 呼叫開始： quantize_(model, int4_weight_only())，它的作用是將模型中 nn.Linear 模組的權重轉換為 int4 量化張量（一個 AffineQuantizedTensor，它是 int4 dtype，非對稱，按組量化），使用 tinygemm 核的佈局： tensor_core_tiled 佈局。

• quantize_：模型級別的 API，透過應用使用者提供的轉換函式（第二個引數）來量化線性層的權重

• int4_weight_only：返回一個函式的函式，該函式將線性層的權重轉換為 int4 僅權重量化權重 * 呼叫量化原語運算子（如 choose_qparams_affine 和 quantize_affine）來量化張量

• TensorCoreTiledLayout：Tensor Core 平鋪佈局型別，儲存打包格式的引數

• TensorCoreTiledAQTTensorImpl：Tensor Core 平鋪的 TensorImpl，儲存打包的權重以用於高效的 int4 僅權重核 (tinygemm 核)

模型執行期間 (During Model Execution)

當我們執行量化模型 model(inputs) 時，我們將透過 nn.Linear 中的函式式線性運算子。

return F.linear(input, weight, bias)

其中輸入是一個 bfloat16 張量，權重是一個 int4 AffineQuantizedTensor，它會呼叫 AffineQuantizedTensor 子類的 __torch_function__，當輸入之一是 AffineQuantizedTensor 時，這將最終進入 F.linear 的實現，因此它會呼叫：

return weight_tensor._quantized_linear_op(input_tensor, weight_tensor, bias)

_quantized_linear_op 會遍歷 _AQT_QLINEAR_DISPATCH_TABLE 並檢查每個排程條件，如果排程條件滿足，它將呼叫使用 input/weight/bias 的實現。請檢視此文件以瞭解 dispatch_condition 和 impl 的解釋。

int4 僅權重的排程條件檢查輸入是否為 bfloat16 張量且權重是否為 uint4 AffineQuantizedTensor。int4 僅權重量化的核心實現接受一個 bfloat16 輸入張量和一個 int4 仿射量化張量，並呼叫 torch.ops.aten._weight_int4pack_mm，傳入輸入張量以及儲存在 weight_tensor.tensor_impl 中的打包權重。

儲存/載入期間

由於 AffineQuantizedTensor 權重仍然是一個 torch.Tensor，因此儲存/載入的工作方式與原始高精度浮點模型相同。有關更多詳細資訊，請參閱序列化文件。