Allow [b]float16 quantizers in sim.onnx.export

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

@@ -713,7 +713,7 @@ def _create_int32_bias_quantizer(self, input, _): # pylint: disable=redefined-bu
 
         if len(input) == 1:
             input, = input
-            if self.input_quantizers[0]:
+            if isinstance(self.input_quantizers[0], AffineQuantizerBase):
                 input_scale = self.input_quantizers[0].get_scale()
             elif isinstance(input, QuantizedTensorBase) and isinstance(input.encoding, AffineEncoding):
                 input_scale = input.encoding.scale

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

@@ -147,4 +147,18 @@ def to_qnn_encoding_dict(self, encoding_version=None) -> Union[List, Dict]:
         if encoding_version == '1.0.0':
             return {'dtype': 'FLOAT', 'bw': self.bitwidth, 'enc_type': EncodingType.PER_TENSOR.name}
 
+        if encoding_version == "2.0.0.beta":
+            if self.exponent_bits == 5 and self.mantissa_bits == 10:
+                # float16
+                return {}
+
+            if self.exponent_bits == 8 and self.mantissa_bits == 7:
+                # bfloat16
+                return {}
+
+            raise NotImplementedError(
+                "Floating point encoding export only supports [b]float16; "
+                f"got exponent_bits={self.exponent_bits}, mantissa_bits={self.mantissa_bits}"
+            )
+
         raise AssertionError(f'Export encoding version {encoding_version} not supported.')

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

@@ -61,7 +61,6 @@
     save_checkpoint,
     load_checkpoint,
     check_accumulator_overflow,
-    _QuantizedModuleProtocol,
 )
 from aimet_torch.v2 import nn as aimet_nn
 from aimet_torch.v2.nn import BaseQuantizationMixin, QuantizationMixin, UnknownModuleError
@@ -70,6 +69,7 @@
 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.float import FloatQuantizeDequantize
 from aimet_torch.v2.quantization.encoding_analyzer import PercentileEncodingAnalyzer
 from aimet_torch.v2.utils import patch_attr
 from aimet_torch import utils
@@ -616,9 +616,11 @@ def _concretize_int32_bias_quantizers(self, args):
                     # In this case, we honor the custom bias quantizer defined by the user
                     continue
 
-                # pylint: disable=protected-access
-                handle = qmodule.register_forward_hook(type(qmodule)._create_int32_bias_quantizer)
-                handles.append(handle)
+                if "weight" in qmodule.param_quantizers and \
+                        isinstance(qmodule.param_quantizers["weight"], AffineQuantizerBase):
+                    # pylint: disable=protected-access
+                    handle = qmodule.register_forward_hook(type(qmodule)._create_int32_bias_quantizer)
+                    handles.append(handle)
             try:
                 self.model(*args)
             finally:
@@ -682,6 +684,41 @@ def _temporarily_unfold_param_quantizers(sim: QuantizationSimModel):
             qmodule._fold_param_quantizers()
 
 
+@contextlib.contextmanager
+def _remove_fp16_quantizers(sim: QuantizationSimModel):
+    """
+    Temporarily remove [b]float16 quantizers for sim.onnx.export,
+    as sim.onnx.export does NOT support exporting [b]float16 quantizers.
+    """
+    original_containers = {}
+
+    try:
+        for qmodule in sim.qmodules():
+            for name, qtzr in qmodule.param_quantizers.items():
+                if isinstance(qtzr, FloatQuantizeDequantize) and \
+                        (qtzr.is_float16() or qtzr.is_bfloat16()):
+                    original_containers[(qmodule.param_quantizers, name)] = qtzr
+                    qmodule.param_quantizers[name] = None
+
+            for i, qtzr in enumerate(qmodule.input_quantizers):
+                if isinstance(qtzr, FloatQuantizeDequantize) and \
+                        (qtzr.is_float16() or qtzr.is_bfloat16()):
+                    original_containers[(qmodule.input_quantizers, i)] = qtzr
+                    qmodule.input_quantizers[i] = None
+
+            for i, qtzr in enumerate(qmodule.output_quantizers):
+                if isinstance(qtzr, FloatQuantizeDequantize) and \
+                        (qtzr.is_float16() or qtzr.is_bfloat16()):
+                    original_containers[(qmodule.output_quantizers, i)] = qtzr
+                    qmodule.output_quantizers[i] = None
+
+        yield
+
+    finally:
+        for (container, key), qtzr in original_containers.items():
+            container[key] = qtzr
+
+
 class _QuantizationSimOnnxExport:
     """
     Helper class for exporting quantized models to ONNX format.
@@ -708,14 +745,13 @@ def export(self,
         :param f: file object or path where to store exported ONNX mode
         """
         # pylint: disable=too-many-locals, too-many-branches, protected-access
-        if self._has_non_affine_quantizer(self.sim.model):
-            raise RuntimeError("Export using onnx only export only supports affine quantizers. "
-                               "Other quantizer types are not supported.")
+        self._check_unsupported_quantizers(self.sim.model)
 
         with tempfile.TemporaryDirectory() as tmp_dir:
             with _temporarily_unfold_param_quantizers(self.sim), \
                     self.sim._concretize_int32_bias_quantizers(args), \
-                    self.sim._apply_qdq_to_model_parameters(self.sim.model):
+                    self.sim._apply_qdq_to_model_parameters(self.sim.model), \
+                    _remove_fp16_quantizers(self.sim):
                 tmp_onnx_path = os.path.join(tmp_dir, "quantized_model.onnx")
                 export(self.sim.model, args, tmp_onnx_path, *posargs, **kwargs)
                 onnx_model = onnx.load(tmp_onnx_path)
@@ -745,30 +781,31 @@ def export(self,
                 "encodings": [
                     {"name": name, **qnn_encoding}
                     for name, qnn_encoding in qnn_encodings.items()
+                    if qnn_encoding
                 ]
             })
         else:
             if quantsim.encoding_version >= "1.0.0":
                 param_encodings = [
                     {"name": name, **qnn_encoding}
                     for name, qnn_encoding in qnn_encodings.items()
-                    if name in param_names
+                    if qnn_encoding and name in param_names
                 ]
                 activation_encodings = [
                     {"name": name, **qnn_encoding}
                     for name, qnn_encoding in qnn_encodings.items()
-                    if name not in param_names
+                    if qnn_encoding and name not in param_names
                 ]
             else:
                 param_encodings = {
                     name: qnn_encoding
                     for name, qnn_encoding in qnn_encodings.items()
-                    if name in param_names
+                    if qnn_encoding and name in param_names
                 }
                 activation_encodings = {
                     name: qnn_encoding
                     for name, qnn_encoding in qnn_encodings.items()
-                    if name not in param_names
+                    if qnn_encoding and name not in param_names
                 }
 
             encodings_dict.update({
@@ -787,15 +824,15 @@ def export(self,
             json.dump(encodings_dict, encoding_file, indent=2)
 
     @staticmethod
-    def _has_non_affine_quantizer(module: torch.nn.Module):
-        for submodule in module.modules():
-            if isinstance(submodule, _QuantizedModuleProtocol):
-                for quantizer in itertools.chain(submodule.input_quantizers,
-                                                 submodule.output_quantizers,
-                                                 submodule.param_quantizers.values()):
-                    if quantizer and not isinstance(quantizer, AffineQuantizerBase):
-                        return True
-        return False
+    def _check_unsupported_quantizers(module: torch.nn.Module):
+        for qtzr in module.modules():
+            if isinstance(qtzr, FloatQuantizeDequantize):
+                if not qtzr.is_float16() and not qtzr.is_bfloat16():
+                    msg = " ".join([
+                        "sim.onnx.export doesn't support exporting floating point encodings",
+                        f"except [b]float16. Got {qtzr.bitwidth}-bit float encoding",
+                    ])
+                    raise RuntimeError(msg)
 
 
 @deprecated("""

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

@@ -274,17 +274,26 @@ def test_export_torchvision_models(model_factory, input_shape):
 
 @torch.no_grad()
 @pytest.mark.parametrize("encoding_version", ["0.6.1", "1.0.0", "2.0.0.beta"])
-@pytest.mark.parametrize("lpbq", (False, True))
-@pytest.mark.parametrize("fold_param_quantizers", (False, True))
-def test_quantsim_export_resnet18(encoding_version, lpbq: bool, fold_param_quantizers: bool):
+@pytest.mark.parametrize("lpbq", [False, True])
+@pytest.mark.parametrize("fold_param_quantizers", [False, True])
+@pytest.mark.parametrize(
+    "weight_dtype,  activation_dtype", [
+    (torch.int8,    torch.uint8),
+    (torch.int8,    torch.float16),
+    (torch.float16, torch.float16),
+])
+def test_quantsim_export_resnet18(encoding_version, lpbq: bool, fold_param_quantizers: bool,
+                                  weight_dtype: torch.dtype, activation_dtype: torch.dtype):
     """
     When: Export quantized torchvision model using quantsim.export
     """
     x = torch.randn(1, 3, 224, 224)
     model = resnet18().eval()
     model = prepare_model(model)
     fold_all_batch_norms(model, None, x)
-    sim = QuantizationSimModel(model, x, config_file=get_path_for_per_channel_config())
+    sim = QuantizationSimModel(model, x,
+                               default_param_bw=weight_dtype.itemsize * 8,
+                               default_output_bw=activation_dtype.itemsize * 8)
 
     if lpbq:
         set_grouped_blockwise_quantization_for_weights(sim,
@@ -294,6 +303,26 @@ def test_quantsim_export_resnet18(encoding_version, lpbq: bool, fold_param_quant
                                                        decompressed_bw=8,
                                                        block_size=64)
 
+    if weight_dtype.is_floating_point:
+        for qmodule in sim.qmodules():
+            for name, qtzr in qmodule.param_quantizers.items():
+                if not qtzr:
+                    continue
+                qmodule.param_quantizers[name] = Q.float.FloatQuantizeDequantize(dtype=weight_dtype)
+
+    if activation_dtype.is_floating_point:
+        for qmodule in sim.qmodules():
+            for i, qtzr in enumerate(qmodule.input_quantizers):
+                if not qtzr:
+                    continue
+                qmodule.input_quantizers[i] = Q.float.FloatQuantizeDequantize(dtype=activation_dtype)
+
+        for qmodule in sim.qmodules():
+            for i, qtzr in enumerate(qmodule.output_quantizers):
+                if not qtzr:
+                    continue
+                qmodule.output_quantizers[i] = Q.float.FloatQuantizeDequantize(dtype=activation_dtype)
+
     sim.compute_encodings(lambda model: model(x))
 
     # Compute original pytorch model output with qdq weights
@@ -302,19 +331,20 @@ def test_quantsim_export_resnet18(encoding_version, lpbq: bool, fold_param_quant
             f"{module_name}.{param_name}": qtzr.get_encodings().to_qnn_encoding_dict(encoding_version)
             for module_name, qmodule in sim.named_qmodules()
             for param_name, qtzr in qmodule.param_quantizers.items()
+            if isinstance(qtzr, Q.affine.AffineQuantizerBase)
         }
         expected_activation_encodings = {}
         expected_activation_encodings.update({
             f"{module_name}.input_quantizers.{i}": qtzr.get_encodings().to_qnn_encoding_dict(encoding_version)
             for module_name, qmodule in sim.named_qmodules()
             for i, qtzr in enumerate(qmodule.input_quantizers)
-            if qtzr is not None
+            if isinstance(qtzr, Q.affine.AffineQuantizerBase)
         })
         expected_activation_encodings.update({
             f"{module_name}.output_quantizers.{i}": qtzr.get_encodings().to_qnn_encoding_dict(encoding_version)
             for module_name, qmodule in sim.named_qmodules()
             for i, qtzr in enumerate(qmodule.output_quantizers)
-            if qtzr is not None
+            if isinstance(qtzr, Q.affine.AffineQuantizerBase)
         })
 
         with remove_activation_quantizers(sim.model):
@@ -407,6 +437,6 @@ def test_quantsim_export_resnet18(encoding_version, lpbq: bool, fold_param_quant
               the original pytorch model with qdq weights
         """
         sess = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
-        out, = sess.run(None, {onnx_model.graph.input[0].name: x.numpy()})
+        out, = sess.run(None, {"input": x.numpy()})
 
         assert torch.allclose(torch.from_numpy(out), expected_out, atol=1e-5)