[ONNX] Implementation of the SmoothQuant algorithm for the ONNX Backend (#3644)

### Changes

- Implementation of the SmoothQuant algorithm for the ONNX Backend 

### Related tickets

Ref: 164847

### Tests

- tests/onnx/test_smooth_quant.py

Author: andrey-churkin

src/nncf/onnx/graph/model_transformer.py

@@ -16,16 +16,19 @@
 import onnx
 
 import nncf
+from nncf.common.factory import NNCFGraphFactory
 from nncf.common.graph.model_transformer import ModelTransformer
 from nncf.common.graph.transformations.commands import TargetType
 from nncf.common.graph.transformations.layout import TransformationLayout
 from nncf.onnx.graph.node_utils import get_input_edge
+from nncf.onnx.graph.node_utils import get_input_edges_mapping
 from nncf.onnx.graph.onnx_helper import get_children
 from nncf.onnx.graph.onnx_helper import get_children_node_mapping
 from nncf.onnx.graph.onnx_helper import get_edge_dtype
 from nncf.onnx.graph.onnx_helper import get_edge_info_mapping
 from nncf.onnx.graph.onnx_helper import get_name_to_node_map
 from nncf.onnx.graph.onnx_helper import get_node_index
+from nncf.onnx.graph.onnx_helper import get_parents_node_mapping
 from nncf.onnx.graph.onnx_helper import get_tensor
 from nncf.onnx.graph.transformations.commands import ONNXInitializerUpdateCommand
 from nncf.onnx.graph.transformations.commands import ONNXModelExtractionCommand
@@ -55,8 +58,8 @@ def __init__(self, model: onnx.ModelProto, inplace: bool = False):
         self.onnx_model_extractor = onnx.utils.Extractor(inferred_model)
         self._inplace = inplace
 
+    @staticmethod
     def _get_target_edge(
-        self,
         port_id: int,
         node_name: str,
         transform_type: TargetType,
@@ -381,10 +384,25 @@ def _apply_initializer_update_transformations(
         :return: Copy of original model with updated biases.
         """
         name_to_node_map = get_name_to_node_map(model)
+        output_name_to_node_map = get_parents_node_mapping(model)
+
         for transformation in transformations:
             node = name_to_node_map[transformation.target_point.target_node_name]
-            initializer_name = node.input[transformation.target_point.port_id]
-            set_initializer(initializer_name, model, transformation.new_value)
+            # NOTE: An `input_name` is either the name of an initializer or the name of a `Constant` operation output
+            input_name = node.input[transformation.target_point.port_id]
+
+            constant_node = output_name_to_node_map.get(input_name, None)
+
+            if constant_node is None:
+                set_initializer(input_name, model, transformation.new_value)
+            else:
+                for attr in constant_node.attribute:
+                    if attr.name == "value":
+                        array = transformation.new_value.astype(onnx.helper.tensor_dtype_to_np_dtype(attr.t.data_type))
+                        tensor_proto = onnx.numpy_helper.from_array(array)
+                        attr.t.CopyFrom(tensor_proto)
+                        break
+
         return model
 
     def _apply_model_extraction_transformation(self, transformation: ONNXModelExtractionCommand) -> onnx.ModelProto:
@@ -484,16 +502,25 @@ def _apply_multiply_insertion_transformations(
         :returns: Transformed model with Multiply nodes.
         """
         node_name_to_node = get_name_to_node_map(model)
+        # TODO(andrey-churkin): Optimize it
+        graph = NNCFGraphFactory.create(model)
+        input_edges_mapping = get_input_edges_mapping(graph)
 
         for transformation in transformations:
+            port_id = transformation.target_point.port_id
             target_node_name = transformation.target_point.target_node_name
-            target_output_port = transformation.target_point.port_id
-            target_node = node_name_to_node[target_node_name]
-            output_tensor_name = target_node.output[target_output_port]
+            transform_type = transformation.target_point.type
+            output_tensor_name = ONNXModelTransformer._get_target_edge(
+                port_id, target_node_name, transform_type, node_name_to_node, input_edges_mapping
+            )
+
+            # TODO(andrey-churkin): Check type of `transformation.scale_value`
 
             # Create a new initializer for the scale constant
             scale_tensor_name = f"{transformation.multiply_node_name}_scale"
-            scale_tensor = onnx.numpy_helper.from_array(transformation.scale_value, name=scale_tensor_name)
+            scale_tensor = onnx.numpy_helper.from_array(
+                transformation.scale_value.astype(np.float32), name=scale_tensor_name
+            )
             model.graph.initializer.append(scale_tensor)
 
             # Create a new Multiply node
@@ -505,7 +532,8 @@ def _apply_multiply_insertion_transformations(
                 name=transformation.multiply_node_name,
             )
             target_index = get_node_index(model, target_node_name)
-            model.graph.node.insert(target_index + 1, mul_node)
+            insert_index = 0 if target_index is None else target_index + 1
+            model.graph.node.insert(insert_index, mul_node)
 
             for name in transformation.destination_node_names:
                 node = node_name_to_node[name]
@@ -524,6 +552,11 @@ def set_initializer(initializer_name: str, model: onnx.ModelProto, new_value: np
     :param new_value: New value for the initializer tensor.
     """
     initializer = get_tensor(model, initializer_name)
+
+    required_dtype = onnx.helper.tensor_dtype_to_np_dtype(initializer.data_type)
+    if new_value.dtype != required_dtype:
+        new_value = new_value.astype(required_dtype)
+
     new_tensor = onnx.numpy_helper.from_array(new_value, initializer_name)
     initializer.CopyFrom(new_tensor)
 

src/nncf/quantization/algorithms/smooth_quant/algorithm.py

@@ -69,7 +69,7 @@ def __init__(
 
     @property
     def available_backends(self) -> list[BackendType]:
-        return [BackendType.OPENVINO, BackendType.TORCH, BackendType.TORCH_FX]
+        return [BackendType.OPENVINO, BackendType.TORCH, BackendType.TORCH_FX, BackendType.ONNX]
 
     def _set_backend_entity(self, model: TModel) -> None:
         """
@@ -90,6 +90,10 @@ def _set_backend_entity(self, model: TModel) -> None:
             from nncf.quantization.algorithms.smooth_quant.torch_fx_backend import FXSmoothQuantAlgoBackend
 
             self._backend_entity = FXSmoothQuantAlgoBackend()
+        elif model_backend == BackendType.ONNX:
+            from nncf.quantization.algorithms.smooth_quant.onnx_backend import ONNXSmoothQuantAlgoBackend
+
+            self._backend_entity = ONNXSmoothQuantAlgoBackend()
         else:
             msg = f"Cannot return backend-specific entity because {model_backend.value} is not supported!"
             raise nncf.UnsupportedBackendError(msg)

src/nncf/quantization/algorithms/smooth_quant/backend.py

@@ -113,13 +113,12 @@ def get_abs_max_channel_collector(
 
     @staticmethod
     @abstractmethod
-    def get_weight_value(node_with_weight: NNCFNode, model: TModel, port_id: int, nncf_graph: NNCFGraph) -> Tensor:
+    def get_weight_value(node_with_weight: NNCFNode, model: TModel, nncf_graph: NNCFGraph) -> Tensor:
         """
         Returns the weight value for the node with weight.
 
         :param node_with_weight: The node with weight.
         :param model: The model that contains this operation.
-        :param port_id: The input port ID to get weight input.
         :param nncf_graph: NNCFGraph instance.
         :return: The weight value.
         """
@@ -141,7 +140,11 @@ def weight_update_command(
     @staticmethod
     @abstractmethod
     def scale_insertion_command(
-        source_node: NNCFNode, scale_value: TTensor, source_output_port_id: int, nodes: list[NNCFNode]
+        source_node: NNCFNode,
+        scale_value: TTensor,
+        source_output_port_id: int,
+        nodes: list[NNCFNode],
+        scale_node_name: str,
     ) -> TransformationCommand:
         """
         Returns command to insert Smooth Quant node.
@@ -150,6 +153,7 @@ def scale_insertion_command(
         :param scale_value: Smooth Quant value.
         :param source_output_port_id: Output port for source node.
         :param nodes: List of consumers for Smooth node.
+        :param scale_node_name: Scale node name.
         :return: TransformationCommand instance.
         """