Fix check in model collector for quantized FLN node (#1462)

To decide whether to collect statistics for a node inside FLN, it is not enough to check its quantization mode, but also needs to verify if it is supposed to be quantized based on the noe config in the fusing information.

Author: ofirgo

model_compression_toolkit/core/common/fusion/fusing_info.py

@@ -39,13 +39,13 @@ class FusingInfo:
     - 'fusing_patterns': The patterns to generate the fused operators from.
     - 'manual_fused_ops': List of sequence of node names to handle as fused ops (even if they are not part of the fusing patterns).
     - `fusing_data`: A dictionary mapping fused operation IDs to lists of nodes that belong to that operation.
-    - `node_to_fused_node_map`: A dictionary mapping each node name to the ID of the fused operation it belongs to.
+    - `node_name_to_fused_op_id`: A dictionary mapping each node name to the ID of the fused operation it belongs to.
 
     """
     fusing_patterns: List[list[any]] = None
     manual_fused_ops: List[List[str]] = None
     fusing_data: Dict[str, Tuple['BaseNode']] = field(default_factory=dict)
-    node_to_fused_node_map: Dict[str, str] = field(init=False, default_factory=dict)
+    node_name_to_fused_op_id: Dict[str, str] = field(init=False, default_factory=dict)
     fused_op_id_to_quant_config: Dict[str, OpQuantizationConfig] = field(default_factory=dict)
 
     def __post_init__(self):
@@ -64,10 +64,10 @@ def _init_node_mapping(self) -> None:
         """
         Init the node-to-fused-node mapping based on the initial fusing data.
         """
-        self.node_to_fused_node_map.clear()
+        self.node_name_to_fused_op_id.clear()
         for op_id, nodes in self.fusing_data.items():
             for node in nodes:
-                self.node_to_fused_node_map[node.name] = op_id
+                self.node_name_to_fused_op_id[node.name] = op_id
 
     def get_manual_nodes_to_fuse(self) -> List[List[str]]:
         """
@@ -115,7 +115,7 @@ def add_fused_operation(self, op_id: str, nodes: Tuple['BaseNode']) -> None:
         self.fusing_data[op_id] = nodes
         # Update the mapping for these nodes
         for node in nodes:
-            self.node_to_fused_node_map[node.name] = op_id
+            self.node_name_to_fused_op_id[node.name] = op_id
 
         # Update the quantization config mapping for this operation
         if self.fusing_patterns is not None:
@@ -152,7 +152,7 @@ def remove_fused_operation(self, op_id: str) -> None:
             self._manual_fused_ops.remove(node_names)
 
         for node in nodes:
-            self.node_to_fused_node_map.pop(node.name, None)
+            self.node_name_to_fused_op_id.pop(node.name, None)
         del self.fusing_data[op_id]
         self.fused_op_id_to_quant_config.pop(op_id, None)
 
@@ -166,7 +166,7 @@ def get_fused_op_id_for_node(self, node_name: str) -> Optional[str]:
         Returns:
             The name of the fused node containing this node, or None if not fused.
         """
-        return self.node_to_fused_node_map.get(node_name)
+        return self.node_name_to_fused_op_id.get(node_name)
 
     def get_node_to_fused_node_map(self) -> Dict[str, str]:
         """
@@ -175,7 +175,7 @@ def get_node_to_fused_node_map(self) -> Dict[str, str]:
         Returns:
             A dictionary mapping each original node name to its fused node name.
         """
-        return self.node_to_fused_node_map.copy()
+        return self.node_name_to_fused_op_id.copy()
 
     def get_fusing_quantization_config_map(self) -> Dict[str, OpQuantizationConfig]:
         """
@@ -198,10 +198,12 @@ def get_fused_nodes(self, op_id: str) -> Optional[List['BaseNode']]:
         """
         return self.fusing_data.get(op_id)
 
-    def get_nodes_to_disable_activation_quantization(self) -> List['BaseNode']:
+    def get_inner_fln_nodes(self) -> List['BaseNode']:
         """
-        Returns a list of the nodes that their activation quantization is disabled due to fusing.
+        Returns a list of the nodes that are part but not the last node of an FLN.
         """
+        # TODO: the order of the nodes is not gurenteed when returned as dict from get_all_fused_operations -
+        #  then, removing the last one can cause issues
         return [node for nodes in self.get_all_fused_operations().values() for node in nodes[:-1]]
 
     def get_fused_op_quantization_config(self, op_id: str) -> OpQuantizationConfig:
@@ -228,6 +230,22 @@ def is_node_in_fused_op(self, node: 'BaseNode') -> bool:
         """
         return any(node in nodes for nodes in self.fusing_data.values())
 
+    def is_quantized_node_in_fln(self, node: 'BaseNode') -> bool:
+        """
+        Check whether a node inside an FLN and should be quantized.
+
+        Args:
+            node (BaseNode): The node to check.
+
+        Returns:
+            bool: True if the node is in any fused operation and should be quantized.
+        """
+        if self.is_node_in_fused_op(node):
+            node_q_cfg = self.fused_op_id_to_quant_config[self.node_name_to_fused_op_id[node.name]]
+            return node_q_cfg is not None and node_q_cfg.enable_activation_quantization
+
+        return False
+
     def get_all_fused_operations(self) -> Dict[str, Tuple['BaseNode']]:
         """
         Retrieve fused information.
@@ -340,7 +358,7 @@ def __repr__(self) -> str:
             for op_id, nodes in self.fusing_data.items()
         )
         mapping_repr = ", ".join(
-            f"{node} -> {op_id}" for node, op_id in self.node_to_fused_node_map.items()
+            f"{node} -> {op_id}" for node, op_id in self.node_name_to_fused_op_id.items()
         )
         return (
             f"FusingInfo(\n"

model_compression_toolkit/core/common/graph/base_graph.py

@@ -908,7 +908,7 @@ def disable_fused_nodes_activation_quantization(self):
         Disable activation quantization for all nodes in fused operations,
         except for the last node in each fused group.
         """
-        nodes_to_disable = self.fusing_info.get_nodes_to_disable_activation_quantization()
+        nodes_to_disable = self.fusing_info.get_inner_fln_nodes()
         for node in nodes_to_disable:
             for qc in node.candidates_quantization_cfg:
                 qc.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.FLN_QUANT

model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py

@@ -677,7 +677,7 @@ def _get_target_activation_nodes(self,
         elif target_criterion == TargetInclusionCriterion.AnyQuantizedNonFused:
             nodes = [n for n in nodes if n.is_activation_quantization_enabled() or n.is_quantization_preserving()]
             # remove fused nodes (due to SNC, where the non-linear is quantized, even though it should not be quantized)
-            nodes = [n for n in nodes if n not in self.graph.fusing_info.get_nodes_to_disable_activation_quantization()]
+            nodes = [n for n in nodes if n not in self.graph.fusing_info.get_inner_fln_nodes()]
         elif target_criterion == TargetInclusionCriterion.QNonConfigurable:
             nodes = [n for n in nodes if n.is_activation_quantization_enabled() and not n.has_configurable_activation()]
         elif target_criterion != TargetInclusionCriterion.Any:    # pragma: no cover

model_compression_toolkit/core/common/model_collector.py

@@ -30,7 +30,8 @@
 
 
 def create_stats_collector_for_node(node: common.BaseNode,
-                                    fw_info: FrameworkInfo) -> BaseStatsCollector:
+                                    fw_info: FrameworkInfo,
+                                    quant_node_in_fln: bool) -> BaseStatsCollector:
     """
     Gets a node and a groups list and create and return a statistics collector for a node
     according to whether its statistics should be collected and the prior information we
@@ -44,7 +45,7 @@ def create_stats_collector_for_node(node: common.BaseNode,
         Statistics collector for statistics collection for the node.
     """
 
-    if node.is_activation_quantization_enabled() or node.is_fln_quantization():
+    if node.is_activation_quantization_enabled() or quant_node_in_fln:
         min_output = getattr(node.prior_info, 'min_output', None)
         max_output = getattr(node.prior_info, 'max_output', None)
         stats_collector = common.StatsCollector(out_channel_axis=fw_info.out_channel_axis_mapping.get(node.type),
@@ -160,7 +161,8 @@ def __init__(self, graph: Graph,
 
         # Assign statistics collectors to nodes
         for n in graph.get_topo_sorted_nodes():
-            sc = create_stats_collector_for_node(n, fw_info=fw_info)  # Get static collector for the node
+            quant_node_in_fln = n.is_fln_quantization() and graph.fusing_info.is_quantized_node_in_fln(n)
+            sc = create_stats_collector_for_node(n, fw_info=fw_info, quant_node_in_fln=quant_node_in_fln)  # Get static collector for the node
             # If we use bias correction, and the node has kernel weights to quantize, we need to make sure
             # its previous nodes' tensors are consistent with this node.
             kernel_attr = fw_info.get_kernel_op_attributes(n.type)[0]

tests_pytest/_fw_tests_common_base/base_ru_integration_test.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 import abc
-import copy
 from typing import  Generator
 from unittest.mock import Mock
 
@@ -245,7 +244,7 @@ def test_snc_fusing(self):
                                  f"{layer_names['act2']}": f"FusedNode_{layer_names['conv2']}_{layer_names['act2']}",
                                  f"{layer_names['conv3']}": f"FusedNode_{layer_names['conv3']}_{layer_names['act3']}",
                                  f"{layer_names['act3']}": f"FusedNode_{layer_names['conv3']}_{layer_names['act3']}"}
-        assert graph.fusing_info.node_to_fused_node_map == expected_node_fuse_op
+        assert graph.fusing_info.node_name_to_fused_op_id == expected_node_fuse_op
 
         ru_calculator = ResourceUtilizationCalculator(graph, self.fw_impl, self.fw_info)
         ru_before_snc, detailed_ru_before_snc = ru_calculator.compute_resource_utilization(
@@ -277,7 +276,7 @@ def test_snc_fusing(self):
                                  f"{layer_names['conv3']}_pre_pad": f"FusedNode_{layer_names['conv3']}_pre_pad_{layer_names['conv3']}_{layer_names['act3']}",
                                  f"{layer_names['conv3']}": f"FusedNode_{layer_names['conv3']}_pre_pad_{layer_names['conv3']}_{layer_names['act3']}",
                                  f"{layer_names['act3']}": f"FusedNode_{layer_names['conv3']}_pre_pad_{layer_names['conv3']}_{layer_names['act3']}"}
-        assert graph.fusing_info.node_to_fused_node_map == expected_node_fuse_op
+        assert graph.fusing_info.node_name_to_fused_op_id == expected_node_fuse_op
 
         linear_w_min_nbit, linear_a_min_nbit, default_w_nbits, default_a_nbit, binary_out_a_bit = nbits
 

tests_pytest/_fw_tests_common_base/fusing/base_fusing_info_generator_test.py

@@ -118,7 +118,7 @@ def fused_graph(self, graph_with_fusion_metadata):
 
     def test_expected_fusing_info(self, graph_with_fusion_metadata):
         actual_fi = graph_with_fusion_metadata.fusing_info
-        assert self.expected_fi.node_to_fused_node_map == actual_fi.node_to_fused_node_map
+        assert self.expected_fi.node_name_to_fused_op_id == actual_fi.node_name_to_fused_op_id
 
     def test_expected_fused_graph(self, fused_graph):
         expected_fused_nodes = self.expected_fi.fusing_data

tests_pytest/_fw_tests_common_base/fusing/base_graph_with_fusing_metadata_test.py

@@ -106,7 +106,7 @@ def test_expected_fusing_info(self, graph_with_fusion_metadata):
         actual_fi = graph_with_fusion_metadata.fusing_info
         assert len(actual_fi.get_all_fused_operations()) == 2
         assert sorted(actual_fi.get_all_fused_operations().keys()) == ['FusedNode_conv_relu', 'FusedNode_linear_softmax']
-        assert actual_fi.node_to_fused_node_map == {'conv': 'FusedNode_conv_relu',
+        assert actual_fi.node_name_to_fused_op_id == {'conv': 'FusedNode_conv_relu',
                                                     'relu': 'FusedNode_conv_relu',
                                                     'linear': 'FusedNode_linear_softmax',
                                                     'softmax': 'FusedNode_linear_softmax'}

tests_pytest/common_tests/unit_tests/test_model_collector.py

@@ -59,7 +59,7 @@ def test_create_stats_collector_for_node_activation_enabled(self, fw_info_mock):
         node.type = DummyLayer
         node.prior_info = Mock(min_output=-1, max_output=9)
 
-        collector = create_stats_collector_for_node(node, fw_info_mock)
+        collector = create_stats_collector_for_node(node, fw_info_mock, False)
         assert isinstance(collector, StatsCollector)
 
     def test_create_stats_collector_for_node_activation_disabled(self, fw_info_mock):
@@ -74,7 +74,7 @@ def test_create_stats_collector_for_node_activation_disabled(self, fw_info_mock)
         # Even if prior_info exists, it should not be used.
         node.prior_info = Mock(min_output=None, max_output=None)
 
-        collector = create_stats_collector_for_node(node, fw_info_mock)
+        collector = create_stats_collector_for_node(node, fw_info_mock, False)
         assert isinstance(collector, NoStatsCollector)
 
     def test_create_stats_collector_for_node_fln_activation_enabled(self, fw_info_mock):
@@ -91,7 +91,7 @@ def test_create_stats_collector_for_node_fln_activation_enabled(self, fw_info_mo
         node.type = DummyLayer
         node.prior_info = Mock(min_output=-8, max_output=9)
 
-        collector = create_stats_collector_for_node(node, fw_info_mock)
+        collector = create_stats_collector_for_node(node, fw_info_mock, True)
         assert isinstance(collector, StatsCollector)
         assert collector.mpcc.init_min_value == -8
         assert collector.mpcc.init_max_value == 9
@@ -108,7 +108,7 @@ def test_create_stats_collector_for_node_fln_activation_enabled_no_prior_info(se
         node.type = DummyLayer
         node.prior_info = None
 
-        collector = create_stats_collector_for_node(node, fw_info_mock)
+        collector = create_stats_collector_for_node(node, fw_info_mock, True)
         assert isinstance(collector, StatsCollector)
         assert collector.mpcc.init_min_value is None
         assert collector.mpcc.init_max_value is None
@@ -125,7 +125,7 @@ def test_create_stats_collector_for_node_fln_activation_enabled_no_prior_info_mo
         node.type = DummyLayer
         node.prior_info = Mock(min_output=None, max_output=None)
 
-        collector = create_stats_collector_for_node(node, fw_info_mock)
+        collector = create_stats_collector_for_node(node, fw_info_mock, True)
         assert isinstance(collector, StatsCollector)
         assert collector.mpcc.init_min_value is None
         assert collector.mpcc.init_max_value is None
@@ -264,6 +264,7 @@ def test_assigns_stats_collectors_to_fln_quantization_nodes(self, fw_impl_mock,
                       output_nodes=[OutTensor(node3, 0)],
                       edge_list=[Edge(node1, node2, 0, 0), Edge(node2, node3, 0, 0)])
         graph.set_out_stats_collector_to_node = Mock(wraps=graph.set_out_stats_collector_to_node)
+        graph.fusing_info.is_quantized_node_in_fln = lambda n: n.is_fln_quantization
 
         # Simulate kernel attribute retrieval.
         fw_info_mock.get_kernel_op_attributes.return_value = [None]