Fuse supergroups to onnx function nodes in QuantSim init

Signed-off-by: Michael Tuttle <mtuttle@qti.qualcomm.com>

Author: michaelgtuttle

TrainingExtensions/onnx/src/python/aimet_onnx/experimental/adascale/model_converter.py

@@ -81,7 +81,7 @@ def get_pt_block(
         output_names,
     )
     subgraph_model = onnx_ir.Model(
-        subgraph, ir_version=model.ir_version, functions=model.functions
+        subgraph, ir_version=model.ir_version, functions=list(model.functions.values())
     )
     ir_utils.remove_aimet_quantizers(subgraph_model)
     ir_utils.inline_all_supergroups(subgraph_model)

TrainingExtensions/onnx/src/python/aimet_onnx/graph_passes/fusions/fusion.py

@@ -33,7 +33,7 @@ def fuse_supergroups(
     unknown_patterns = [p for p in patterns if p not in FUSION_PASS_REGISTRY]
     if unknown_patterns:
         raise ValueError(
-            f"Unknown pattern names: {unknown_patterns}. "
+            f"Graph pass requested but not found: {unknown_patterns}. "
             f"Available patterns: {list(FUSION_PASS_REGISTRY.passes.keys())}"
         )
 

TrainingExtensions/onnx/src/python/aimet_onnx/graph_passes/passes/common_patterns.py

@@ -16,6 +16,8 @@
 
 def match_rms_norm_pattern(op: Op, model: ModelProto) -> List[Op]:
     """Common pattern for RMSNormalization which can be re-used"""
+    if op.type == "RMSNormalization":
+        return [op]
     # Match Mul(x, x) or Pow(x, 2)
     match = match_pow_2_pattern(op, model)
     if not match or len(op.output_ops) != 1:

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

@@ -152,6 +152,8 @@
 ]
 _tie_qtzrs = True
 
+_fuse_supergroups = True
+
 data_types_to_quantize = [np.float32, np.float16]
 
 _DEPRECATED_ARGS = {
@@ -443,20 +445,13 @@ def __init__(
                 "Quantizing models with BFLOAT16 tensors is not supported"
             )
 
-        self.model = model
         self._op_domain = op_domain
         self.providers = providers
 
         if not dummy_input:
-            dummy_input = make_dummy_input(self.model.model)
+            dummy_input = make_dummy_input(model.model)
 
         self.qc_quantize_op_dict = {}
-        self.connected_graph = ConnectedGraph(self.model)
-        if _has_unfolded_batchnorms(self.model.model, self.connected_graph):
-            logger.warning(
-                "Model contains unfolded BatchNormalization layers. To accurately simulate quantization behavior, "
-                "please call aimet_onnx.batch_norm_fold.fold_all_batch_norms_to_weight(model) before creating QuantizationSimModel."
-            )
         self._quant_scheme = quant_scheme
         self._param_type = param_type
         self._activation_type = activation_type
@@ -473,6 +468,23 @@ def __init__(
         for lib in user_onnx_libs or []:
             self._ort_session_options.register_custom_ops_library(lib)
 
+        quantsim_configurator = QuantSimConfigurator(
+            config_file,
+            self._param_type,
+            self._activation_type,
+        )
+        if _fuse_supergroups:
+            model = quantsim_configurator.apply_fusions(model)
+
+        self.model = model
+        self.connected_graph = ConnectedGraph(self.model)
+
+        if _has_unfolded_batchnorms(self.model.model, self.connected_graph):
+            logger.warning(
+                "Model contains unfolded BatchNormalization layers. To accurately simulate quantization behavior, "
+                "please call aimet_onnx.batch_norm_fold.fold_all_batch_norms_to_weight(model) before creating QuantizationSimModel."
+            )
+
         # Get names of parameters and activations to quantize
         self._get_param_names()
         self._get_activations_to_quantize(dummy_input)
@@ -485,7 +497,14 @@ def __init__(
         }
 
         # Apply configurations based on provided config file.
-        quantsim_configurator = self._add_configuration_(config_file)
+        quantsim_configurator.configure_quantizers(
+            self.model,
+            self.connected_graph,
+            self.qc_quantize_op_dict,
+            self.param_names,
+            self.activation_names,
+            self.input_quantizers_name,
+        )
         self._hw_version = quantsim_configurator._get_hw_version()
         self._supported_kernels = quantsim_configurator.get_supported_kernels()
         self._op_to_supported_kernel = (
@@ -663,28 +682,6 @@ def get_supported_kernels(self) -> Dict:
         """
         return self._supported_kernels
 
-    def _add_configuration_(self, config_file: str):
-        """
-        Add configuration based on config file
-
-        :param config_file: Path to Configuration file for model quantizers
-        """
-        quantsim_configurator = QuantSimConfigurator(
-            self.model,
-            self.connected_graph,
-            config_file,
-            self._param_type,
-            self._activation_type,
-        )
-        quantsim_configurator.configure_quantizers(
-            self.qc_quantize_op_dict,
-            self.param_names,
-            self.activation_names,
-            self.input_quantizers_name,
-        )
-
-        return quantsim_configurator
-
     def _get_param_names(self):
         """
         Get the names of params

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

@@ -10,7 +10,9 @@
 
 import numpy as np
 import onnx
+import onnx_ir
 from packaging import version
+from onnxruntime.quantization.onnx_quantizer import ONNXModel
 
 from aimet_onnx.common.defs import QuantizationDataType, qtype
 from aimet_onnx.common.graph_searcher import (
@@ -41,6 +43,7 @@
 from aimet_onnx.utils import get_product_name_from_quantized_name
 from aimet_onnx.qc_quantize_op import OpMode, QcQuantizeOp
 from aimet_onnx.graph_passes.pass_registry import apply_graph_passes, find_all_matches
+from aimet_onnx.graph_passes.fusions import fuse_supergroups
 
 # pylint: disable=no-name-in-module, ungrouped-imports
 if version.parse(onnx.__version__) >= version.parse("1.14.0"):
@@ -86,16 +89,14 @@ class QuantSimConfigurator(AimetCommonQuantSimConfigurator):
 
     def __init__(
         self,
-        model: ModelProto,
-        conn_graph: ConnectedGraph,
         config_file: str,
         param_type: qtype,
         activation_type: qtype,
     ):
         super().__init__(config_file, param_type, activation_type)
 
-        self._model = model
-        self._conn_graph = conn_graph
+        self._model = None
+        self._conn_graph = None
         self._quant_ops_dict = {}
         self._param_names = {}
         self._activation_names = {}
@@ -121,6 +122,8 @@ def get_supported_kernels(self) -> Dict:
 
     def configure_quantizers(
         self,
+        model: ONNXModel,
+        conn_graph: ConnectedGraph,
         quant_ops_dict: Dict,
         param_names: List[str],
         activation_names: List[str],
@@ -129,6 +132,8 @@ def configure_quantizers(
         """
         Configures quantizers based on config file
         """
+        self._model = model
+        self._conn_graph = conn_graph
         self._quant_ops_dict = quant_ops_dict
         self._param_names = param_names
         self._activation_names = activation_names
@@ -150,12 +155,21 @@ def configure_quantizers(
         # Run supergroup passes if specified in config
         supergroup_pass_list = self._get_supergroup_pass_list()
         apply_graph_passes(
-            self._model.model,
-            self._conn_graph,
+            model.model,
+            conn_graph,
             self._quant_ops_dict,
             supergroup_pass_list,
         )
 
+    def apply_fusions(self, model: ONNXModel):
+        """
+        Fuse supergroups into onnx functions based on config file supergroup list
+        """
+        fusion_patterns = self._get_supergroup_pass_list()
+        ir_model = onnx_ir.from_proto(model.model)
+        fused_model = fuse_supergroups(ir_model, fusion_patterns)
+        return ONNXModel(onnx_ir.to_proto(fused_model))
+
     def _map_quantizers_to_ops(self) -> Dict[str, OpToQuantizers]:
         """
         Creates a dict where key is the name of the op and value is OpToQuantizers which comprises of input quantizers

TrainingExtensions/onnx/test/python/test_enpu_config.py

@@ -226,10 +226,8 @@ def test_batchnorm_params(
             *(
                 nn.LayerNorm(
                     normalized_shape=10,
-                    elementwise_affine=elementwise_affine,
                     bias=bias,
                 )
-                for elementwise_affine in (True, False)
                 for bias in (True, False)
             ),
             *(
@@ -301,7 +299,6 @@ def test_supergroup(self, tmp_path: Path, model: nn.Module, tie_quantizers: bool
         intermediate_activations = (
             sim.qc_quantize_op_dict.keys() - param_names - input_names - output_names
         )
-        assert intermediate_activations  # Sanity check
 
         for name in intermediate_activations:
             qtzr = sim.qc_quantize_op_dict[name]

TrainingExtensions/onnx/test/python/test_find_transformer_blocks.py

@@ -21,10 +21,10 @@ def verify_find_blocks(sim, model_type):
     end_points_names = [(op1.name, op2.name) for op1, op2 in end_points]
     assert end_points_names == [
         (
-            "/model/model/layers.0/input_layernorm/Pow",
-            "/model/model/layers.1/input_layernorm/Pow",
+            "/model/model/layers.0/input_layernorm",
+            "/model/model/layers.1/input_layernorm",
         ),
-        ("/model/model/layers.1/input_layernorm/Pow", "/model/model/norm/Pow"),
+        ("/model/model/layers.1/input_layernorm", "/model/model/norm"),
     ]
     conv_linear_blocks = get_conv_linear_layers_decoder_block(sim, end_points)
     conv_linear_blocks_names = []
@@ -89,10 +89,10 @@ def test_get_decoder_blocks_phi(add_genai_tests_path):
     end_points_names = [(op1.name, op2.name) for op1, op2 in end_points]
     assert end_points_names == [
         (
-            "/model/model/layers.0/input_layernorm/Pow",
-            "/model/model/layers.1/input_layernorm/Pow",
+            "/model/model/layers.0/input_layernorm",
+            "/model/model/layers.1/input_layernorm",
         ),
-        ("/model/model/layers.1/input_layernorm/Pow", "/model/model/norm/Pow"),
+        ("/model/model/layers.1/input_layernorm", "/model/model/norm"),
     ]
     conv_linear_blocks = get_conv_linear_layers_decoder_block(
         collection.backbone, end_points

TrainingExtensions/onnx/test/python/test_graph_passes/test_fusions.py

@@ -660,7 +660,7 @@ def test_unknown_pattern_raises_error(self, tmp_path):
         model_proto = create_layernorm_model(tmp_path)
         model = onnx_ir.from_proto(model_proto)
 
-        with pytest.raises(ValueError, match="Unknown pattern names"):
+        with pytest.raises(ValueError, match="Graph pass requested but not found"):
             fuse_supergroups(model, patterns=["UnknownPattern"])
 
     @pytest.mark.parametrize(

TrainingExtensions/onnx/test/python/test_graph_passes/test_graph_registry.py

@@ -17,7 +17,7 @@
 import tempfile
 
 from ..models.models_for_tests import build_dummy_model
-from ..utils import tmp_dir
+from ..utils import tmp_dir, patch_fuse_supergroups
 
 
 def _generate_quantsim_config(supergroup_pass_name: str, file_path: str) -> dict:
@@ -62,14 +62,15 @@ def test_register_and_apply_graph_pass(tmp_dir):
 
     config_file = str(os.path.join(tmp_dir, "quantsim_config.json"))
     _generate_quantsim_config("DummyTestGraphPass", config_file)
-    sim = QuantizationSimModel(
-        model,
-        input_data,
-        quant_scheme=QuantScheme.post_training_tf,
-        default_param_bw=8,
-        default_activation_bw=8,
-        config_file=config_file,
-    )
+    with patch_fuse_supergroups(False):
+        sim = QuantizationSimModel(
+            model,
+            input_data,
+            quant_scheme=QuantScheme.post_training_tf,
+            default_param_bw=8,
+            default_activation_bw=8,
+            config_file=config_file,
+        )
 
     graph = ConnectedGraph(model)
     disable_quantizers = set(

TrainingExtensions/onnx/test/python/test_graph_passes/test_layernorm.py

@@ -9,6 +9,7 @@
 import pytest
 
 from ..models.test_models import layernorm_model
+from ..utils import patch_fuse_supergroups
 from .utils import assert_on_const_quantizers, assert_on_output_quantizers
 
 
@@ -20,14 +21,15 @@ def test_layer_norm(elementwise_affine, bias):
     graph = ConnectedGraph(model)
 
     input_data = {"x": np.random.rand(1, 3, dim, dim).astype(np.float32)}
-    sim = QuantizationSimModel(
-        model,
-        input_data,
-        quant_scheme=QuantScheme.post_training_tf,
-        default_param_bw=8,
-        default_activation_bw=8,
-        config_file="htp_v81",
-    )
+    with patch_fuse_supergroups(False):
+        sim = QuantizationSimModel(
+            model,
+            input_data,
+            quant_scheme=QuantScheme.post_training_tf,
+            default_param_bw=8,
+            default_activation_bw=8,
+            config_file="htp_v81",
+        )
 
     all_ops = graph.ordered_ops
     # Check if quantization is disabled for LayerNormalization intermediate op outputs
@@ -52,14 +54,15 @@ def test_layer_norm_intermediate():
     graph = ConnectedGraph(model)
 
     input_data = {"x": np.random.rand(1, 3, dim, dim).astype(np.float32)}
-    sim = QuantizationSimModel(
-        model,
-        input_data,
-        quant_scheme=QuantScheme.post_training_tf,
-        default_param_bw=8,
-        default_activation_bw=8,
-        config_file="htp_v81",
-    )
+    with patch_fuse_supergroups(False):
+        sim = QuantizationSimModel(
+            model,
+            input_data,
+            quant_scheme=QuantScheme.post_training_tf,
+            default_param_bw=8,
+            default_activation_bw=8,
+            config_file="htp_v81",
+        )
 
     all_ops = graph.ordered_ops
     # Check if quantization is disabled for LayerNormalization intermediate op outputs