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¶

文档

教程

资源