QATConfig¶

class torchao.quantization.qat.QATConfig(base_config: Optional[AOBaseConfig] = None, activation_config: Optional[FakeQuantizeConfigBase] = None, weight_config: Optional[FakeQuantizeConfigBase] = None, *, step: QATStep = 'prepare')[原始碼]¶

用於將量化感知訓練 (QAT) 應用於 torch.nn.Module 的配置，與 quantize_() 一起使用。

此配置有兩個步驟：“prepare”（準備）和“convert”（轉換）。prepare 步驟將“偽”量化應用到模型，應在訓練前應用，而 convert 步驟將模型轉換為實際的量化模型。這裡的偽量化是指使用高精度算術（例如 bf16）來模擬量化數值（例如 int4），目的是減少量化帶來的最終退化。

有兩種使用此配置的方法。第一種方法是傳遞一個基礎的訓練後量化 (PTQ) 配置，我們將使用它來自動推斷在 prepare 階段使用的相應偽量化方案。在 convert 階段，我們將把基礎 PTQ 配置應用到模型。這將是最常見的用例。

使用示例

from torchao.quantization import (
    quantize_,
    Int8DynamicActivationInt4WeightConfig,
)
from torchao.quantization.qat import QATConfig

base_config = Int8DynamicActivationInt4WeightConfig(group_size=32)
quantize_(model, QATConfig(base_config, step="prepare"))
train_loop(model)
quantize_(model, QATConfig(base_config, step="convert"))

目前僅支援以下作為基礎配置

Int8DynamicActivationInt4WeightConfig

Int4WeightOnlyConfig

第二種使用此配置的方法是透過直接指定偽量化方案。使用者將為權重和/或啟用傳遞 FakeQuantizeConfigBase，而不是基礎 PTQ 配置。此用例主要用於實驗，例如，當相應的 PTQ 配置尚不存在時。

使用示例

from torchao.quantization import quantize_
from torchao.quantization.qat import IntxFakeQuantizeConfig

activation_config = IntxFakeQuantizeConfig(
    torch.int8, "per_token", is_symmetric=False,
)
weight_config = IntxFakeQuantizeConfig(
    torch.int4, group_size=32, is_symmetric=True,
)
qat_config = QATConfig(
    # must specify one of `base_config` or `weight_config`
    activation_config=act_config,
    weight_config=weight_config,
    step="prepare",
)
quantize_(model, qat_config)

引數:

base_config (Optional[AOBaseConfig]) – 基礎 PTQ 配置，用於在 prepare 階段推斷偽量化配置，並在 convert 階段直接應用。
activation_config (Optional[FakeQuantizeConfigBase]) – 輸入啟用的自定義偽量化配置，始終是可選的。如果使用了 base_config，則必須為 None。
weight_config (Optional[FakeQuantizeConfigBase]) – 權重的自定義偽量化配置。如果使用了 base_config，則必須為 None。

關鍵字引數:

step (str) – “prepare”或“convert”之一，決定 QAT 階段

丟擲:

ValueError – 如果同時指定了 base_config 和 activation_config
ValueError – 如果同時指定了 base_config 和 weight_config
ValueError – 如果沒有指定 base_config、activation_config 或 weight_config 中的任何一個
ValueError – 如果指定了 activation_config 或 weight_config 中的任何一個，並且 step 是“convert”
ValueError – 如果 step 不是“prepare”或“convert”之一
ValueError – 如果配置應用於的模組不是 torch.nn.Linear 或 torch.nn.Embedding，或者應用於 torch.nn.Embedding 並帶啟用配置

QATConfig¶

文件

教程

資源