Make export_int32_bias default to True if encoding_version >= 2.0.0

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

Author: 2 people

TrainingExtensions/onnx/src/python/aimet_onnx/quantsim.py

@@ -1828,7 +1828,8 @@ def _get_weight_and_bias(
         export_model (bool, optional):
             If True, then ONNX model is exported. When False, only encodings are exported.
         export_int32_bias (bool, optional):
-            If true, generate and export int32 bias encoding on the fly (default: `True`)
+            If true, generate and export int32 bias encoding on the fly.
+            Default: `True` if encoding version is 2.0.0 or higher, otherwise `False`.
         encoding_version (str, optional):
             Version of the encoding format to use. (default: {quantsim.encoding_version})
             Supported versions are: {sorted(list(quantsim.VALID_ENCODING_VERSIONS))}
@@ -1848,7 +1849,7 @@ def export(
         filename_prefix: str,
         export_model: bool = True,
         *,
-        export_int32_bias: bool = False,
+        export_int32_bias: Optional[bool] = None,
         encoding_version: Optional[str] = None,
         force_activation_as: Literal["unsigned"]
         | Literal["signed"]
@@ -1870,6 +1871,11 @@ def export(
             )
             warnings.warn(msg, DeprecationWarning, stacklevel=2)
 
+        if export_int32_bias is None:
+            export_int32_bias = version.parse(encoding_version) >= version.parse(
+                "2.0.0"
+            )
+
         with (
             self._concretize_int32_bias_quantizers()
             if export_int32_bias

TrainingExtensions/onnx/test/python/test_quantsim.py

@@ -577,23 +577,26 @@ def test_export_model_2_0_0(self, tmp_path: pathlib.Path, activation_type):
             "version": aimet_onnx.__version__,
         }
         encodings = encodings["encodings"]
-        assert (
-            sorted(e["name"] for e in encodings)
-            == [
+        encoding_keys = set(e["name"] for e in encodings)
+
+        if activation_type == aimet_onnx.int8:
+            assert encoding_keys == {
                 "4",
                 "5",
                 "6",
+                "conv_b",
                 "conv_w",
+                "fc_b",
                 "fc_w",
                 "input",
                 "output",
-            ]
-            if activation_type == aimet_onnx.int8
-            else [
+            }
+        else:
+            assert encoding_keys == {
                 "conv_w",
                 "fc_w",
-            ]
-        )
+            }
+
         # Exported encoding can contain more entry than qc_quantize_op_dict since
         # some grid-preserving op's input/output encodings are auto-generated
         assert set(e["name"] for e in encodings) >= {
@@ -604,7 +607,13 @@ def test_export_model_2_0_0(self, tmp_path: pathlib.Path, activation_type):
 
         # Cross-check with onnx QDQ.
         expected_encodings = _remove_onnx_qdq_nodes(sim.to_onnx_qdq())
-        assert len(encodings) == len(expected_encodings)
+        expected_encoding_keys = set(e["name"] for e in expected_encodings)
+
+        if activation_type == aimet_onnx.int8:
+            assert encoding_keys == expected_encoding_keys | {"conv_b", "fc_b"}
+            encodings = [e for e in encodings if e["name"] in expected_encoding_keys]
+        else:
+            assert encoding_keys == expected_encoding_keys
 
         encodings = sorted(encodings, key=lambda e: e["name"])
         expected_encodings = sorted(expected_encodings, key=lambda e: e["name"])
@@ -1620,7 +1629,13 @@ def callback(session, args):
                     for name, q in sim.qc_quantize_op_dict.items()
                     if q in reconfigured_quantizers
                 ]
-                assert len(mismatched_encodings) == len(reconfigured_tensors)
+
+                if encoding_version in ("0.6.1", "1.0.0"):
+                    assert len(mismatched_encodings) == len(reconfigured_tensors)
+                else:
+                    # mismatched_encodings contains 2 extra entries for int32 bias encodings
+                    assert len(mismatched_encodings) == len(reconfigured_tensors) + 2
+
                 assert np.allclose(
                     out2,
                     out3,

TrainingExtensions/onnx/test/python/test_torch_to_onnx.py

@@ -177,7 +177,12 @@ def test_torch_to_onnx_zero_point_shift(tmp_dir, encoding_version):
     onnx_export_path = os.path.join(tmp_dir, "onnx_export")
     onnx_encoding_path = os.path.join(onnx_export_path, "model.encodings")
     os.makedirs(onnx_export_path, exist_ok=True)
-    onnx_sim.export(onnx_export_path, "model", encoding_version=encoding_version)
+    onnx_sim.export(
+        onnx_export_path,
+        "model",
+        encoding_version=encoding_version,
+        export_int32_bias=False,
+    )
 
     with open(onnx_encoding_path, "r") as f:
         onnx_encodings = json.load(f)

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

@@ -29,6 +29,7 @@
 import json
 import contextlib
 import os
+from packaging import version
 from aimet_torch.v2.nn.true_quant import QuantizedGRU, QuantizedLSTM, QuantizedRNN
 import torch
 import onnx
@@ -1064,7 +1065,8 @@ def __init__(self, sim: QuantizationSimModel):
             Version of the encoding format to use. (default: {quantsim.encoding_version})
             Supported versions are: {sorted(list(quantsim.VALID_ENCODING_VERSIONS))}
         export_int32_bias (bool, optional):
-            If true, generate and export int32 bias encoding on the fly (default: `True`).
+            If true, generate and export int32 bias encoding on the fly.
+            Default: `True` if encoding version is 2.0.0 or higher, otherwise `False`.
         force_activation_as (str, optional):
             Force representing quantized activations as signed or unsigned integers.
             This argument is only applicable for encoding version 2.0.0.
@@ -1093,7 +1095,7 @@ def export(
         f: Union[str, io.BytesIO],
         *,
         encoding_version: Optional[str] = None,
-        export_int32_bias: bool = True,
+        export_int32_bias: Optional[bool] = None,
         force_activation_as: Literal["unsigned"]
         | Literal["signed"]
         | None = "unsigned",
@@ -1107,6 +1109,11 @@ def export(
                 f"Supported versions are: {quantsim.VALID_ENCODING_VERSIONS}"
             )
 
+        if export_int32_bias is None:
+            export_int32_bias = version.parse(encoding_version) >= version.parse(
+                "2.0.0"
+            )
+
         from aimet_torch.onnx import (
             _temporarily_convert_activation_to,
             _check_unsupported_args,