Dynamically Quantized Models - Dynamically Quantized Models - 57300

Quantizing models

Hugging Face models can be quantized using Int8DynamicActivationInt8WeightConfig from TorchAO. Zentorch supports dynamically quantized models with INT8 dynamically quantized activation with per-token granularity and INT8 quantized weights with per-channel granularity.

To quantize a model, use the TorchAoConfig integration in Hugging Face Transformers as shown here.

import torch

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TorchAoConfig,
)
from torchao.quantization import Int8DynamicActivationInt8WeightConfig
from torchao.quantization.quant_primitives import MappingType

quantization_config = TorchAoConfig(
    Int8DynamicActivationInt8WeightConfig(
        version=2,
        act_mapping_type=MappingType.SYMMETRIC,
    )
)

quantized_model = AutoModelForCausalLM.from_pretrained(
    args.model_name,
    torch_dtype=torch.bfloat16,
    device_map="cpu",
    quantization_config=quantization_config,
    trust_remote_code=True,
)