IntXQuantizationAwareTrainingConfig

class torchao.quantization.qat.IntXQuantizationAwareTrainingConfig(activation_config: Optional[FakeQuantizeConfig] = None, weight_config: Optional[FakeQuantizeConfig] = None)

用于将伪量化应用于 torch.nn.Module 的配置。与 quantize_() 一起使用。

使用示例

from torchao.quantization import quantize_
from torchao.quantization.qat import FakeQuantizeConfig
activation_config = FakeQuantizeConfig(
    torch.int8, "per_token", is_symmetric=False,
)
weight_config = FakeQuantizeConfig(
    torch.int4, group_size=32, is_symmetric=True,
)
quantize_(
    model,
    IntXQuantizationAwareTrainingConfig(activation_config, weight_config),
)

注意：如果配置应用于非 torch.nn.Linear 或 torch.nn.Embedding 的模块，或者应用于带有激活配置的 torch.nn.Embedding，则会引发 ValueError，因为这些情况不受支持。