Add fold_param_quantizers as a public feature

Signed-off-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>
Co-authored-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>

Author: 2 people

TrainingExtensions/torch/src/python/aimet_torch/v2/nn/base.py

@@ -46,14 +46,19 @@
 from torch import nn
 
 from aimet_torch.utils import is_vector_encoding
-from aimet_torch.v2.quantization.affine.encoding import VectorEncoding, AffineEncoding
+from aimet_torch.v2.quantization.affine.encoding import (
+    AffineEncoding,
+    GroupedBlockEncoding,
+    VectorEncoding,
+)
 from aimet_torch.v2.quantization.affine import (
     AffineQuantizerBase,
     GroupedBlockQuantizeDequantize,
     QuantizeDequantize,
 )
+from aimet_torch.v2.quantization.float import FloatEncoding, FloatQuantizeDequantize
 
-from aimet_torch.v2.quantization.tensor import QuantizedTensorBase
+from aimet_torch.v2.quantization.tensor import QuantizedTensorBase, DequantizedTensor
 from aimet_torch.v2.quantization.base import QuantizerBase
 from aimet_torch.v2.utils import (
     patch_attr,
@@ -750,10 +755,39 @@ def _create_int32_bias_quantizer(self, input, _): # pylint: disable=redefined-bu
     def _derive_bias_scale(self, input_scale: Optional[torch.Tensor], weight_scale: Optional[torch.Tensor]):
         raise NotImplementedError
 
-    def _fold_param_quantizers(self):
+    def fold_param_quantizers(self):
         """
-        Fold param quantizers into parameters to speed up inference.
+        Fold parameter quantizers into their associated parameters to accelerate inference.
+
+        Example:
+
+          >>> qlinear = QuantizedLinear(10, 10)
+          >>> qlinear.param_quantizers["weight"] = QuantizeDequantize((), -128, 127, symmetric=True)
+          >>> type(qlinear.weight)
+          <class 'torch.nn.parameter.Parameter'>
+          >>> qlinear
+          QuantizedLinear(
+            in_features=10, out_features=10, bias=True
+            (param_quantizers): ModuleDict(
+              (weight): QuantizeDequantize(shape=(), qmin=-128, qmax=127, symmetric=True)
+              (bias): None
+            )
+          )
+          >>> qlinear.fold_param_quantizers()
+          >>> type(qlinear.weight)
+          <class 'aimet_torch.v2.quantization.tensor.DequantizedTensor'>
+          >>> qlinear
+          QuantizedLinear(
+            in_features=10, out_features=10, bias=True
+            (param_quantizers): ModuleDict(
+              (weight): None
+              (bias): None
+            )
+          )
         """
+        return self._fold_param_quantizers()
+
+    def _fold_param_quantizers(self):
         self._compute_param_encodings(overwrite=False)
 
         for param_name, param_qtzr in self.param_quantizers.items():
@@ -765,6 +799,33 @@ def _fold_param_quantizers(self):
             setattr(self, param_name, torch.nn.Parameter(qdq_param, requires_grad=param.requires_grad))
             self.param_quantizers[param_name] = None
 
+    def _unfold_param_quantizers(self):
+        """
+        Re-instantiate param quantizers for ease of export
+        """
+        for param_name, qdq_param in self.named_parameters():
+            if not isinstance(qdq_param, DequantizedTensor):
+                continue
+
+            if qdq_param.encoding is None:
+                continue
+
+            if isinstance(qdq_param.encoding, GroupedBlockEncoding):
+                param_qtzr = GroupedBlockQuantizeDequantize.from_encodings(qdq_param.encoding)
+            elif isinstance(qdq_param.encoding, AffineEncoding):
+                param_qtzr = QuantizeDequantize.from_encodings(qdq_param.encoding)
+            elif isinstance(qdq_param.encoding, FloatEncoding):
+                param_qtzr = FloatQuantizeDequantize.from_encodings(qdq_param.encoding)
+            else:
+                raise ValueError
+
+            if not param_qtzr:
+                continue
+
+            param = qdq_param.as_subclass(torch.Tensor)
+            setattr(self, param_name, torch.nn.Parameter(param, requires_grad=param.requires_grad))
+            self.param_quantizers[param_name] = param_qtzr
+
 
 def _remove_quantizers(quantizers, keys):
     orig_quantizers = {key: quantizers[key] for key in keys}

TrainingExtensions/torch/src/python/aimet_torch/v2/quantization/affine/quantizer.py

@@ -232,6 +232,21 @@ def get_encodings(self) -> Optional[AffineEncoding]:
                                   self.qmin, self.qmax, self._symmetric, self.block_size)
         return None
 
+    @classmethod
+    def from_encodings(cls, encodings: AffineEncoding) -> "AffineQuantizerBase":
+        if not isinstance(encodings, AffineEncoding):
+            raise TypeError(f"Expected {AffineEncoding}; got {type(encodings)}")
+
+        qtzr = cls(shape=encodings.scale.shape,
+                   qmin=encodings.qmin,
+                   qmax=encodings.qmax,
+                   symmetric=encodings.symmetry,
+                   block_size=encodings.block_size)
+
+        qtzr.set_range(encodings.min, encodings.max)
+
+        return qtzr
+
     @torch.no_grad()
     def get_legacy_encodings(self) -> Optional[List[Dict]]:
         """
@@ -863,3 +878,19 @@ def get_encodings(self) -> Optional[GroupedBlockEncoding]:
                                         decompressed_bw=self.decompressed_bw,
                                         per_channel_scale=per_channel_scale)
         return None
+
+    @classmethod
+    def from_encodings(cls, encodings: GroupedBlockEncoding) -> "GroupedBlockQuantizeDequantize":
+        if not isinstance(encodings, GroupedBlockEncoding):
+            raise TypeError(f"Expected {GroupedBlockEncoding}; got {type(encodings)}")
+
+        qtzr = cls(shape=encodings.scale.shape,
+                   bitwidth=encodings.bitwidth,
+                   symmetric=encodings.symmetry,
+                   decompressed_bw=encodings.decompressed_bw,
+                   block_size=encodings.block_size,
+                   block_grouping=encodings.block_grouping)
+
+        qtzr.set_range(encodings.min, encodings.max)
+
+        return qtzr

TrainingExtensions/torch/src/python/aimet_torch/v2/quantization/base/quantizer.py

@@ -124,6 +124,13 @@ def set_encodings(self, encodings: EncodingBase):
         """
         raise NotImplementedError
 
+    @classmethod
+    @abc.abstractmethod
+    def from_encodings(cls, encodings: EncodingBase) -> "QuantizerBase":
+        """
+        Create quantizer object from encoding object
+        """
+
     def register_quantization_parameter(self, name: str, param: nn.Parameter):
         """
         Register quantization parameter.

TrainingExtensions/torch/src/python/aimet_torch/v2/quantization/float/quantizer.py

@@ -246,6 +246,19 @@ def get_encodings(self) -> Optional[FloatEncoding]:
             return FloatEncoding(self.mantissa_bits, self.exponent_bits, self.maxval)
         return None
 
+    @classmethod
+    def from_encodings(cls, encodings: FloatEncoding) -> "FloatQuantizeDequantize":
+        if not isinstance(encodings, FloatEncoding):
+            raise TypeError(f"Expected {FloatEncoding}; got {type(encodings)}")
+
+        qtzr = cls(exponent_bits=encodings.exponent_bits,
+                   mantissa_bits=encodings.mantissa_bits)
+
+        if encodings.maxval is not None:
+            qtzr.maxval.copy_(encodings.maxval)
+
+        return qtzr
+
     @contextlib.contextmanager
     def compute_encodings(self):
         """

TrainingExtensions/torch/src/python/aimet_torch/v2/quantsim/quantsim.py

@@ -67,6 +67,7 @@
 from aimet_torch.v2.nn import BaseQuantizationMixin, QuantizationMixin, UnknownModuleError
 from aimet_torch.v2.nn.fake_quant import _legacy_impl
 from aimet_torch.v2._builder import _V2LazyQuantizeWrapper
+from aimet_torch.v2.quantization import DequantizedTensor
 from aimet_torch.v2.quantization.base import QuantizerBase
 from aimet_torch.v2.quantization.affine import AffineQuantizerBase
 from aimet_torch.v2.quantization.encoding_analyzer import PercentileEncodingAnalyzer
@@ -628,6 +629,58 @@ def _concretize_int32_bias_quantizers(self, args):
             for qmodule, qtzr in orig_bias_quantizers.items():
                 qmodule.param_quantizers["bias"] = qtzr
 
+    def fold_param_quantizers(self):
+        """
+        Fold parameter quantizers into their associated parameters to accelerate inference.
+
+        Example:
+
+          >>> sim = QuantizationSimModel(...)
+          >>> type(sim.model[0].weight)
+          <class 'torch.nn.parameter.Parameter'>
+          >>> sim.model[0]
+          QuantizedLinear(
+            in_features=10, out_features=10, bias=True
+            (param_quantizers): ModuleDict(
+              (weight): QuantizeDequantize(shape=(), qmin=-128, qmax=127, symmetric=True)
+              (bias): None
+            )
+          )
+          >>> sim.fold_param_quantizers()
+          >>> type(sim.model[0].weight)
+          <class 'aimet_torch.v2.quantization.tensor.DequantizedTensor'>
+          >>> sim.model[0]
+          QuantizedLinear(
+            in_features=10, out_features=10, bias=True
+            (param_quantizers): ModuleDict(
+              (weight): None
+              (bias): None
+            )
+          )
+        """
+        for qmodule in self.qmodules():
+            qmodule.fold_param_quantizers()
+
+
+@contextlib.contextmanager
+def _temporarily_unfold_param_quantizers(sim: QuantizationSimModel):
+    # pylint: disable=protected-access
+    """
+    Temporarily re-instantiate param quantizers for ease of export
+    """
+    modules_with_folded_parameters = [
+        qmodule for qmodule in sim.qmodules()
+        if any(isinstance(param, DequantizedTensor) for param in qmodule.parameters())
+    ]
+
+    try:
+        for qmodule in modules_with_folded_parameters:
+            qmodule._unfold_param_quantizers()
+        yield
+    finally:
+        for qmodule in modules_with_folded_parameters:
+            qmodule._fold_param_quantizers()
+
 
 class _QuantizationSimOnnxExport:
     """
@@ -660,7 +713,8 @@ def export(self,
                                "Other quantizer types are not supported.")
 
         with tempfile.TemporaryDirectory() as tmp_dir:
-            with self.sim._concretize_int32_bias_quantizers(args), \
+            with _temporarily_unfold_param_quantizers(self.sim), \
+                    self.sim._concretize_int32_bias_quantizers(args), \
                     self.sim._apply_qdq_to_model_parameters(self.sim.model):
                 tmp_onnx_path = os.path.join(tmp_dir, "quantized_model.onnx")
                 export(self.sim.model, args, tmp_onnx_path, *posargs, **kwargs)

TrainingExtensions/torch/test/python/v2/experimental/test_onnx.py

@@ -275,7 +275,8 @@ def test_export_torchvision_models(model_factory, input_shape):
 @torch.no_grad()
 @pytest.mark.parametrize("encoding_version", ["0.6.1", "1.0.0"])
 @pytest.mark.parametrize("lpbq", (False, True))
-def test_quantsim_export_resnet18(encoding_version, lpbq: bool):
+@pytest.mark.parametrize("fold_param_quantizers", (False, True))
+def test_quantsim_export_resnet18(encoding_version, lpbq: bool, fold_param_quantizers: bool):
     """
     When: Export quantized torchvision model using quantsim.export
     """
@@ -319,6 +320,9 @@ def test_quantsim_export_resnet18(encoding_version, lpbq: bool):
         with remove_activation_quantizers(sim.model):
             expected_out = sim.model(x)
 
+    if fold_param_quantizers:
+        sim.fold_param_quantizers()
+
     with tempfile.TemporaryDirectory() as dirname:
         onnx_path = os.path.join(dirname, "torchvision_model.onnx")
         encodings_path = os.path.join(dirname, "torchvision_model.encodings")

TrainingExtensions/torch/test/python/v2/quantsim/test_quantsim.py

@@ -46,6 +46,7 @@
 from aimet_common.defs import QuantizationDataType, QuantScheme
 from aimet_torch import onnx_utils
 from aimet_torch.v2.quantsim import QuantizationSimModel, load_encodings_to_sim
+from aimet_torch.v2.quantization import DequantizedTensor
 from aimet_torch.v2.quantization.encoding_analyzer import PercentileEncodingAnalyzer
 from aimet_torch.v2.quantization.base import QuantizerBase
 from aimet_torch.v2.quantization.affine import AffineQuantizerBase, GroupedBlockQuantizeDequantize, QuantizeDequantize
@@ -1242,6 +1243,41 @@ def test_compute_encodings_optional_arg(self):
                 assert torch.equal(qtzr_a.get_min(), qtzr_b.get_min())
                 assert torch.equal(qtzr_a.get_max(), qtzr_b.get_max())
 
+    @pytest.mark.parametrize("data_type", [QuantizationDataType.int, QuantizationDataType.float])
+    def test_fold_param_quantizers(self, tmpdir, data_type):
+        model = torch.nn.Sequential(
+            torch.nn.Linear(10, 10),
+        )
+        x = torch.randn(10, 10)
+        sim = QuantizationSimModel(model, x,
+                                   default_param_bw=16,
+                                   default_output_bw=16,
+                                   default_data_type=data_type)
+        sim.compute_encodings(lambda model: model(x))
+
+        sim.export(tmpdir, "before_fold", x)
+
+        """
+        When: Call fold_param_quantizers()
+        Then: 1. All param quantizers should be folded to the parameter
+              2. Export artifact of sim.export() should not be affected
+        """
+        sim.fold_param_quantizers()
+        assert sim.model[0].param_quantizers["weight"] is None
+        assert isinstance(sim.model[0].weight, DequantizedTensor)
+
+        sim.export(tmpdir, "after_fold", x)
+
+        with open(os.path.join(tmpdir, "before_fold.encodings")) as f:
+            encodings_before_fold = json.load(f)
+        with open(os.path.join(tmpdir, "after_fold.encodings")) as f:
+            encodings_after_fold = json.load(f)
+
+        assert encodings_before_fold == encodings_after_fold
+
+        # trivial sanity check
+        assert [enc["name"] for enc in encodings_before_fold["param_encodings"]] == ["0.weight"]
+
 
 class TestQuantsimUtilities: