Add static quantization as an example for calibration flow

Summary:
So far quantization flow API that we provided (`quantize_`) does not require calibration (calibrate a model with sample data), this PR added a static quantization
example that serves as an example for calibration flow

* 1. first prepare the model for calibration
* 2. calibrate the prepared model with sample data
* 3. convert the calibrated model to quantized model

Test Plan:
python torchao/prototype/calibration_flow/static_quant.py

Reviewers:

Subscribers:

Tasks:

Tags:

Author: jerryzh168

torchao/dtypes/__init__.py

@@ -1,12 +1,13 @@
 from .nf4tensor import NF4Tensor, to_nf4
 # from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor
 from .uint4 import UInt4Tensor
-from .affine_quantized_tensor import AffineQuantizedTensor, to_affine_quantized
+from .affine_quantized_tensor import AffineQuantizedTensor, to_affine_quantized, to_affine_quantized_static
 
 __all__ = [
     "NF4Tensor",
     "to_nf4",
     "UInt4Tensor"
     "AffineQuantizedTensor",
     "to_affine_quantized",
+    "to_affine_quantized_static",
 ]

torchao/dtypes/affine_quantized_tensor.py

@@ -228,6 +228,49 @@ def from_float(
             dtype=input_float.dtype
         )
 
+    @classmethod
+    def from_float_static(
+        cls,
+        input_float: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        block_size: Tuple[int, ...],
+        target_dtype: torch.dtype,
+        quant_min: Optional[int] = None,
+        quant_max: Optional[int]  = None,
+        zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT,
+        extended_layout: str = "plain",
+        # TODO: this is only for "tensor_core_tiled", need to figure out
+        # the proper API for this arg
+        inner_k_tiles: Optional[int] = None,
+    ):
+        original_shape = input_float.shape
+        if extended_layout == "tensor_core_tiled":
+            orig_out_features, orig_in_features = input_float.shape
+            in_features = find_multiple(orig_in_features, 1024)
+            out_features = find_multiple(orig_out_features, 8)
+            input_float = torch.nn.functional.pad(
+                input_float,
+                (0, in_features - orig_in_features, 0, out_features - orig_out_features),
+            )
+        int_data = quantize_affine(input_float, block_size, scale, zero_point, target_dtype, quant_min, quant_max, zero_point_domain)
+
+        layout_cls_ctr = get_layout_tensor_constructor(extended_layout)
+        # TODO: this is temporary, need to come up with the proper UX
+        if extended_layout == "tensor_core_tiled":
+            layout_tensor = layout_cls_ctr(int_data, scale, zero_point, inner_k_tiles)
+        else:
+            layout_tensor = layout_cls_ctr(int_data, scale, zero_point)
+        return cls(
+            layout_tensor,
+            block_size,
+            original_shape,
+            quant_min,
+            quant_max,
+            zero_point_domain,
+            dtype=input_float.dtype,
+        )
+
     @property
     def extended_layout(self) -> str:
         return self.layout_tensor.extended_layout
@@ -764,3 +807,4 @@ def t(func, *args, **kwargs):
     return return_and_correct_aliasing(func, args, kwargs, new)
 
 to_affine_quantized = AffineQuantizedTensor.from_float
+to_affine_quantized_static = AffineQuantizedTensor.from_float_static

torchao/prototype/calibration_flow/static_quant.py

@@ -0,0 +1,110 @@
+"""
+Demo for static quantization flow
+"""
+import torch
+import copy
+
+# TODO: use the generalized observer for affine qunatization in the future
+from torch.ao.quantization.observer import MinMaxObserver, PerChannelMinMaxObserver
+import torch.nn.functional as F
+from torch import Tensor
+from torchao.dtypes import to_affine_quantized_static
+from torchao.quantization.utils import compute_error
+
+
+class CalibratingLinear(torch.nn.Linear):
+    def __init__(self, in_features: int, out_features: int, act_obs: torch.nn.Module, weight_obs: torch.nn.Module, bias: bool = True, device=None, dtype=None):
+        super().__init__(in_features, out_features, bias, device, dtype)
+        self.act_obs = act_obs
+        self.weight_obs = weight_obs
+
+    def forward(self, input: Tensor):
+        calibrating_input = self.act_obs(input)
+        calibrating_weight = self.weight_obs(self.weight)
+        return F.linear(calibrating_input, calibrating_weight, self.bias)
+
+    @classmethod
+    def from_float(cls, float_linear, act_obs, weight_obs):
+        calibrating_linear = cls(float_linear.in_features, float_linear.out_features, act_obs, weight_obs, False, device=float_linear.weight.device, dtype=float_linear.weight.dtype)
+        calibrating_linear.weight = float_linear.weight
+        calibrating_linear.bias = float_linear.bias
+        return calibrating_linear
+
+class QuantizedLinear(torch.nn.Module):
+    def __init__(self, in_features: int, out_features: int, act_obs: torch.nn.Module, weight_obs: torch.nn.Module, weight: torch.Tensor, bias: torch.Tensor):
+        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])
+        target_dtype = torch.uint8
+        self.qweight = to_affine_quantized_static(weight, weight_scale, weight_zero_point, block_size, target_dtype)
+        self.bias = bias
+
+    def forward(self, input: Tensor):
+        block_size = input.shape
+        target_dtype = torch.uint8
+        qinput = to_affine_quantized_static(input, self.act_scale, self.act_zero_point, block_size, target_dtype)
+        return F.linear(qinput, self.qweight, self.bias)
+
+    @classmethod
+    def from_calibrating(cls, calibrating_linear):
+        quantized_linear = cls(calibrating_linear.in_features, calibrating_linear.out_features, calibrating_linear.act_obs, calibrating_linear.weight_obs, calibrating_linear.weight, calibrating_linear.bias)
+        return quantized_linear
+
+class ToyLinearModel(torch.nn.Module):
+    def __init__(self, m=64, n=32, k=64):
+        super().__init__()
+        self.linear1 = torch.nn.Linear(m, n, bias=False)
+        self.linear2 = torch.nn.Linear(n, k, bias=False)
+
+    def example_inputs(self, batch_size=1, dtype=torch.float32, device="cpu"):
+        return (torch.randn(batch_size, self.linear1.in_features, dtype=dtype, device=device),)
+
+    def forward(self, x):
+        x = self.linear1(x)
+        x = self.linear2(x)
+        return x
+
+dtype = torch.bfloat16
+m = ToyLinearModel(1024, 1024, 1024).eval().to(dtype).to("cuda")
+m_bf16 = copy.deepcopy(m)
+example_inputs = m.example_inputs(dtype=dtype, device="cuda")
+
+m_bf16 = torch.compile(m_bf16, mode='max-autotune')
+
+from torchao.quantization.quant_api import _replace_with_custom_fn_if_matches_filter
+import copy
+
+def to_calibrating_(model, act_obs, weight_obs):
+    _is_linear = lambda m, fqn: isinstance(m, torch.nn.Linear)
+    replacement_fn = lambda m: CalibratingLinear.from_float(m, act_obs, weight_obs)
+    act_obs = copy.deepcopy(act_obs)
+    weight_obs = copy.deepcopy(weight_obs)
+    _replace_with_custom_fn_if_matches_filter(model, replacement_fn, _is_linear)
+
+def to_quantized_(model):
+    _is_calibrating_linear = lambda m, fqn: isinstance(m, CalibratingLinear)
+
+    replacement_fn = lambda m: QuantizedLinear.from_calibrating(m)
+    _replace_with_custom_fn_if_matches_filter(model, replacement_fn, _is_calibrating_linear)
+
+
+# TODO: use the generalized observer for affine qunatization in the future
+act_obs = MinMaxObserver(dtype=torch.uint8, qscheme=torch.per_tensor_affine).to("cuda")
+weight_obs = PerChannelMinMaxObserver(dtype=torch.uint8, qscheme=torch.per_channel_affine).to("cuda")
+
+before_quant = m(*example_inputs)
+
+to_calibrating_(m, act_obs, weight_obs)
+# calibrating / training
+for _ in range(10):
+    m(*example_inputs)
+
+after_obs = m(*example_inputs)
+to_quantized_(m)
+
+print("quantized model:", m)
+after_quant = m(*example_inputs)
+assert compute_error(before_quant, after_quant) > 30
+print("test passed")