quantize.py

import argparse
import copy
import json
from dataclasses import dataclass
from typing import Callable

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchao.core.config import AOBaseConfig
from torchao.dtypes import to_affine_quantized_intx_static
from torchao.quantization import quantize_
from torchao.quantization.granularity import PerAxis, PerTensor
from torchao.quantization.observer import (
    AffineQuantizedMinMaxObserver,
    AffineQuantizedObserverBase,
)
from torchao.quantization.quant_api import _replace_with_custom_fn_if_matches_filter
from torchao.quantization.quant_primitives import MappingType
from torchao.quantization.transform_module import register_quantize_module_handler
from torchvision import datasets
from torchvision.transforms import ToTensor

from model import Device, MNISTModel
from utils import compute_fixed_point_multiplier_and_shift


class ObservedLinear(nn.Linear):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        act_obs: AffineQuantizedObserverBase,
        weight_obs: AffineQuantizedObserverBase,
        output_obs: AffineQuantizedObserverBase,
        bias: bool = True,
        device=None,
        dtype=None,
        last_layer: bool = False,
    ):
        super().__init__(in_features, out_features, bias, device, dtype)
        self.act_obs = act_obs
        self.weight_obs = weight_obs
        self.output_obs = output_obs
        self._last_layer = last_layer

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        observed_input = self.act_obs(input)
        observed_weight = self.weight_obs(self.weight)
        output = F.linear(observed_input, observed_weight, self.bias)

        if self._last_layer:
            self.output_obs(output)
        else:
            self.output_obs(F.relu(output))

        return output

    @classmethod
    def from_float(
        cls,
        float_linear: nn.Linear,
        act_obs: AffineQuantizedObserverBase,
        weight_obs: AffineQuantizedObserverBase,
        output_obs: AffineQuantizedObserverBase,
        last_layer: bool = False,
    ) -> "ObservedLinear":
        observed_linear = cls(
            float_linear.in_features,
            float_linear.out_features,
            act_obs,
            weight_obs,
            output_obs,
            False,
            device=float_linear.weight.device,
            dtype=float_linear.weight.dtype,
            last_layer=last_layer,
        )
        observed_linear.weight = float_linear.weight
        observed_linear.bias = float_linear.bias
        return observed_linear


class QuantizedLinear(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        act_obs: AffineQuantizedObserverBase,
        weight_obs: AffineQuantizedObserverBase,
        output_obs: AffineQuantizedObserverBase,
        weight: torch.Tensor,
        bias: torch.Tensor,
        target_dtype: torch.dtype,
    ):
        super().__init__()
        self.act_scale, self.act_zero_point = act_obs.calculate_qparams()
        weight_scale, weight_zero_point = weight_obs.calculate_qparams()
        assert weight.dim() == 2
        block_size = (1, weight.shape[1])
        self.target_dtype = target_dtype
        self.bias = bias
        self.qbias = torch.round(bias / (weight_scale * self.act_scale)).to(torch.int32)
        self.qweight = to_affine_quantized_intx_static(
            weight, weight_scale, weight_zero_point, block_size, self.target_dtype
        )

        self.out_scale, self.out_zero_point = output_obs.calculate_qparams()
        self.multiplier, self.shift = compute_fixed_point_multiplier_and_shift(
            weight_scale * self.act_scale[0] / self.out_scale[0]
        )

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        block_size = input.shape
        qinput = to_affine_quantized_intx_static(
            input,
            self.act_scale,
            self.act_zero_point,
            block_size,
            self.target_dtype,
        )
        return F.linear(qinput, self.qweight, self.bias)

    @classmethod
    def from_observed(
        cls, observed_linear: ObservedLinear, target_dtype: torch.dtype
    ) -> "QuantizedLinear":
        quantized_linear = cls(
            observed_linear.in_features,
            observed_linear.out_features,
            observed_linear.act_obs,
            observed_linear.weight_obs,
            observed_linear.output_obs,
            observed_linear.weight,
            observed_linear.bias,
            target_dtype,
        )
        return quantized_linear


def _get_last_linear_fqn(model: nn.Module) -> str | None:
    """Find the fully qualified name of the last nn.Linear layer in the model."""
    last_fqn = None
    for fqn, module in model.named_modules():
        if isinstance(module, nn.Linear):
            last_fqn = fqn
    return last_fqn


def insert_observers_(
    model: nn.Module,
    act_obs: AffineQuantizedObserverBase,
    weight_obs: AffineQuantizedObserverBase,
    output_obs: AffineQuantizedObserverBase,
    *,
    filter_fn: Callable[[nn.Module, str], bool] | None = None,
):
    last_linear_fqn = _get_last_linear_fqn(model)

    def _is_linear(m: nn.Module, _fqn: str) -> bool:
        return isinstance(m, torch.nn.Linear)

    def convert_to_linear_observer(linear_module: nn.Linear, fqn: str | None = None):
        return ObservedLinear.from_float(
            float_linear=linear_module,
            act_obs=copy.deepcopy(act_obs),
            weight_obs=copy.deepcopy(weight_obs),
            output_obs=copy.deepcopy(output_obs),
            last_layer=(fqn == last_linear_fqn),
        )

    _replace_with_custom_fn_if_matches_filter(
        model,
        convert_to_linear_observer,
        _is_linear if filter_fn is None else filter_fn,
    )


def _calibrate(model: nn.Module, data_loader: DataLoader, device: Device = "cpu"):
    model.eval()
    with torch.no_grad():
        for images, _ in data_loader:
            images = images.to(device)
            model(images)


@dataclass
class StaticQuantConfig(AOBaseConfig):
    target_dtype: torch.dtype


@register_quantize_module_handler(StaticQuantConfig)
def _apply_static_quant(
    module: ObservedLinear,
    config: StaticQuantConfig,
):
    """
    Define a transformation associated with `StaticQuantConfig`.
    This is called by `quantize_`, not by the user directly.
    """
    return QuantizedLinear.from_observed(module, config.target_dtype)


def _save_quantized_model(
    model: nn.Module,
    output: str,
) -> None:
    model_data = {}
    for layer_name, layer in model.named_modules():
        if isinstance(layer, QuantizedLinear):
            model_data[f"{layer_name}_weight"] = {
                "data": layer.qweight.tensor_impl.data.int_data.tolist(),
                "dtype": str(layer.qweight.tensor_impl.data.int_data.dtype),
            }
            model_data[f"{layer_name}_scale"] = {
                "data": layer.qweight.tensor_impl.data.scale.tolist(),
                "dtype": str(layer.qweight.tensor_impl.data.scale.dtype),
            }
            model_data[f"{layer_name}_zero_point"] = {
                "data": layer.qweight.tensor_impl.data.zero_point.to(
                    torch.int8
                ).tolist(),
                "dtype": str(torch.int8),
            }
            model_data[f"{layer_name}_bias"] = {
                "data": layer.qbias.tolist(),
                "dtype": str(layer.qbias.dtype),
            }
            model_data[f"{layer_name}_act_scale"] = {
                "data": layer.act_scale.tolist(),
                "dtype": str(layer.act_zero_point.dtype),
            }
            model_data[f"{layer_name}_act_zero_point"] = {
                "data": layer.act_zero_point.to(torch.int8).tolist(),
                "dtype": str(torch.int8),
            }
            model_data[f"{layer_name}_out_scale"] = {
                "data": layer.out_scale.tolist(),
                "dtype": str(layer.out_scale.dtype),
            }
            model_data[f"{layer_name}_out_zero_point"] = {
                "data": layer.out_zero_point.to(torch.int8).tolist(),
                "dtype": str(torch.int8),
            }
            model_data[f"{layer_name}_multiplier"] = {
                "data": layer.multiplier.to(torch.int32).tolist(),
                "dtype": str(torch.int32),
            }
            model_data[f"{layer_name}_shift"] = {
                "data": layer.shift.to(torch.int32).tolist(),
                "dtype": str(torch.int32),
            }

    with open(output, "w") as output_file:
        json.dump(model_data, output_file)


def quantize(model: nn.Module, input: str, output: str, device: Device = "cpu") -> None:
    # Load model weights
    model = model.to(device)
    model.load_state_dict(torch.load(input))

    # Create observers
    act_obs = AffineQuantizedMinMaxObserver(
        MappingType.SYMMETRIC,
        torch.int8,
        granularity=PerTensor(),
        eps=torch.finfo(torch.float32).eps,
        scale_dtype=torch.float32,
        zero_point_dtype=torch.float32,
    )

    weight_obs = AffineQuantizedMinMaxObserver(
        MappingType.SYMMETRIC,
        torch.int8,
        granularity=PerAxis(axis=0),
        eps=torch.finfo(torch.float32).eps,
        scale_dtype=torch.float32,
        zero_point_dtype=torch.float32,
    )

    output_obs = AffineQuantizedMinMaxObserver(
        MappingType.SYMMETRIC,
        torch.int8,
        granularity=PerAxis(axis=0),
        eps=torch.finfo(torch.float32).eps,
        scale_dtype=torch.float32,
        zero_point_dtype=torch.float32,
    )

    # Insert observers (automatically detects last linear layer)
    insert_observers_(
        model,
        act_obs=act_obs,
        weight_obs=weight_obs,
        output_obs=output_obs,
    )

    calibration_dataset = datasets.MNIST(
        root="data", train=False, download=True, transform=ToTensor()
    )
    calibration_loader = DataLoader(calibration_dataset, batch_size=1, shuffle=False)

    # Calibrate model
    _calibrate(model, calibration_loader, device=device)

    # Perform static quantization in-place
    quantize_(
        model,
        StaticQuantConfig(torch.int8),
        lambda m, _fqn: isinstance(m, ObservedLinear),
    )

    # Save quantized model
    _save_quantized_model(model, output)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="TorchAO Static Quantization",
        description="Post-Training Static Quantization using TorchAO",
    )

    parser.add_argument(
        "-i",
        "--input",
        type=str,
        default="./data/models/best.pt",
        help="Path to the pre-trained PyTorch model checkpoint",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=str,
        default="./data/models/quantized_weights.json",
        help="Path to save the quantized model",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cpu",
        help="Device to run calibration on: cpu or cuda",
    )

    args = parser.parse_args()

    quantize(
        MNISTModel(),
        input=args.input,
        output=args.output,
        device=args.device,
    )