per tensor symmetric activation quantization

[jivercx]

Is there a w8a8 QAT config that support the following describe?
int8 per tensor symmetric activation quantization and int8 per channel weight symmetric quantization

[supriyar]

cc @andrewor14

[andrewor14]

Hi @jivercx, I don't think there's a pre-set config for this but you can express this using something like the following:

from torchao.quantization.qat import (
    IntxFakeQuantizeConfig,
    QATConfig,
)
from torchao.quantization import (
    PerAxis,
    PerTensor,
    quantize_,
)

act_config = IntxFakeQuantizeConfig(
    dtype=torch.int8,
    granularity=PerTensor(),
    is_symmetric=True,
)

weight_config = IntxFakeQuantizeConfig(
    dtype=torch.int8,
    granularity=PerAxis(0),
    is_symmetric=True,
)

qat_config = QATConfig(
    activation_config=act_config,
    weight_config=weight_config,
    step="prepare",
)

# This will perform a simple module swap nn.Linear -> FakeQuantizedLinear
quantize_(model, qat_config)

However, we don't have a corresponding convert path for this since there's no kernel that's for these specific configs. If you have one such kernel, you can manually swap FakeQuantizedLinear in the model to your specific MyQuantizedLinear class that does call this kernel in the forward pass after training.