[Torch FX] Map Namespace Names to Metatype

nncf/experimental/torch/fx/nncf_graph_builder.py

@@ -65,22 +65,6 @@ def _get_layer_attributes(
             )
         return None
 
-    def _map_fx_unique_metatypes(node: torch.fx.Node, metatype: om.OperatorMetatype) -> om.OperatorMetatype:
-        """
-        Attempts to retrieve correct subtype for the given node.
-
-        :param node: Given node.
-        :param metatype: Given node metatype.
-        :param model: Target GraphModule instance.
-        :return: Correct FX metatype of the given node if it is exist or the original node metatype otherwise.
-        """
-        if metatype in [om.PTEmbeddingMetatype]:
-            weight_node = node.args[0]
-            if weight_node.op == "get_attr":
-                return om.PTAtenEmbeddingMetatype
-
-        return metatype
-
     @staticmethod
     def get_node_type_and_metatype(node: torch.fx.Node, model: torch.fx.GraphModule) -> Tuple[str, om.OperatorMetatype]:
         """
@@ -90,6 +74,7 @@ def get_node_type_and_metatype(node: torch.fx.Node, model: torch.fx.GraphModule)
         :param model: Given GraphModule.
         :return: Node's type and metatype.
         """
+        node_type_name = None
         if node.op == "placeholder":
             node_type = "input"
             node_metatype = om.PTInputNoopMetatype
@@ -101,9 +86,11 @@ def get_node_type_and_metatype(node: torch.fx.Node, model: torch.fx.GraphModule)
             node_metatype = om.PTConstNoopMetatype
         elif node.op in ("call_function",):
             if hasattr(node.target, "overloadpacket"):
-                node_type = str(node.target.overloadpacket).split(".")[1]
+                node_type = str(node.target.overloadpacket)
+                node_type_name = node_type.split(".")[1]
             elif node.target.__name__ == "getitem":
-                node_type = "__getitem__"
+                node_type = "aten.__getitem__"
+                node_type_name = "__getitem__"
             else:
                 # TODO(dlyakhov): get correct nodes types from this nodes as well
                 node_type = str(node.target)
@@ -118,7 +105,8 @@ def get_node_type_and_metatype(node: torch.fx.Node, model: torch.fx.GraphModule)
             layer_attrs = GraphConverter._get_layer_attributes(node, node_metatype, model)
             node_subtype = node_metatype.determine_subtype(layer_attrs)
             node_metatype = node_subtype or node_metatype
-        return node_type, node_metatype
+        node_type_name = node_type_name or node_type
+        return node_type_name, node_metatype
 
     @staticmethod
     def create_nncf_graph(model: torch.fx.GraphModule) -> PTNNCFGraph:
@@ -135,7 +123,6 @@ def create_nncf_graph(model: torch.fx.GraphModule) -> PTNNCFGraph:
         const_targets_counter = Counter([node.target for node in model.graph.nodes if node.op == "get_attr"])
         for source_node in model.graph.nodes:
             node_type, node_metatype = GraphConverter.get_node_type_and_metatype(source_node, model)
-            node_metatype = GraphConverter._map_fx_unique_metatypes(source_node, node_metatype)
             is_shared_node = source_node.op in ("get_attr",) and (
                 const_targets_counter[source_node.target] > 1 or len(source_node.users) > 1
             )
@@ -148,7 +135,7 @@ def create_nncf_graph(model: torch.fx.GraphModule) -> PTNNCFGraph:
             source_nncf_node = nncf_graph.get_node_by_name(source_node.name)
             for idx, dist_node in enumerate(source_node.users):
                 dist_node_id = nncf_graph.get_node_by_name(dist_node.name).node_id
-                input_port_id, output_port_id, tensor_shape = GraphConverter.get_edge_params(
+                input_port_id, output_port_id, tensor_shape, tensor_dtype = GraphConverter.get_edge_params(
                     model, source_node, source_nncf_node, dist_node, idx
                 )
                 nncf_graph.add_edge_between_nncf_nodes(
@@ -157,7 +144,7 @@ def create_nncf_graph(model: torch.fx.GraphModule) -> PTNNCFGraph:
                     tensor_shape=tensor_shape,
                     input_port_id=input_port_id,
                     output_port_id=output_port_id,
-                    dtype=Dtype.FLOAT,
+                    dtype=tensor_dtype,
                 )
         return nncf_graph
 
@@ -168,7 +155,7 @@ def get_edge_params(
         source_nncf_node: NNCFNode,
         dist_node: torch.fx.Node,
         output_idx: int,
-    ) -> Tuple[int, int, Tuple[int, ...]]:
+    ) -> Tuple[int, int, Tuple[int, ...], Dtype]:
         """
         Retrieves edge params from the given source_node and dist_node pair.
 
@@ -182,8 +169,11 @@ def get_edge_params(
         """
         output_port_id = 0
         tensor_shape = None
+        tensor_dtype = Dtype.FLOAT
         if source_node.op in ("get_attr",):
-            tensor_shape = tuple(get_tensor_constant_from_node(source_node, model).shape)
+            tensor = get_tensor_constant_from_node(source_node, model)
+            tensor_shape = tuple(tensor.shape)
+            tensor_dtype = Dtype.INTEGER if tensor.dtype == torch.int else tensor_dtype
         elif "val" in source_node.meta:
             if source_nncf_node.metatype is om.PTBatchNormMetatype and isinstance(
                 source_node.meta["val"], (tuple, list)
@@ -197,6 +187,7 @@ def get_edge_params(
                 tensor = source_node.meta["val"]
             if isinstance(tensor, torch.Tensor):
                 tensor_shape = tuple(-1 if isinstance(i, torch.SymInt) else i for i in tensor.shape)
+                tensor_dtype = Dtype.INTEGER if tensor.dtype == torch.int else tensor_dtype
             elif isinstance(tensor, torch.SymInt):
                 tensor_shape = (-1,)
 
@@ -205,4 +196,4 @@ def get_edge_params(
             nncf_logger.debug(f"Edge shape between {source_node.name} and {dist_node.name} is unknown.")
 
         input_port_id = dist_node.all_input_nodes.index(source_node)
-        return input_port_id, output_port_id, tensor_shape
+        return input_port_id, output_port_id, tensor_shape, tensor_dtype

nncf/experimental/torch2/function_hook/extractor.py

@@ -176,7 +176,7 @@ def extract_conv(
     if input_node == output_node:
         return conv_module
 
-    if output_node.metatype is not om.PTBatchNormMetatype:
+    if output_node.metatype != om.PT2BatchNormMetatype:
         msg = f"Support only PTBatchNormMetatype as output node, actual: {output_node.metatype}"
         raise nncf.InternalError(msg)
 

nncf/experimental/torch2/function_hook/graph/graph_utils.py

@@ -80,7 +80,11 @@ class FunctionMeta:
 
     @property
     def func_name(self) -> str:
-        return self.func.__name__
+        if self.func.__qualname__.split(".")[0] == "TensorBase":
+            return f"torch.tensor.{self.func.__name__}"
+        elif self.func.__qualname__ == self.func.__name__:
+            return f"torch.nn.functional.{self.func.__name__}"
+        return f"{self.func.__module__}.{self.func.__name__}"
 
 
 @dataclass

nncf/experimental/torch2/function_hook/nncf_graph/nncf_graph_builder.py

@@ -89,9 +89,9 @@ def get_meta_type(node_type: str, meta: Union[ConstMeta, FunctionMeta, InOutMeta
     :param meta: The metadata associated with the node.
     :return: The PTOperatorMetatype object.
     """
-    node_metatype = cast(
-        type[om.PTOperatorMetatype], om.PT_OPERATOR_METATYPES.get_operator_metatype_by_op_name(node_type)
-    )
+    metatype = om.PT_OPERATOR_METATYPES.get_operator_metatype_by_op_name(node_type)
+
+    node_metatype = cast(type[om.PTOperatorMetatype], metatype)
     node_sub_meta_type: Optional[type[om.PTOperatorMetatype]] = None
     if node_metatype.get_subtypes() and isinstance(meta, FunctionMeta):
         node_sub_meta_type = node_metatype.determine_subtype(function_args=meta.args, functions_kwargs=meta.kwargs)
@@ -187,7 +187,7 @@ def convert_to_nncf_graph(nx_graph: nx.MultiDiGraph) -> PTNNCFGraph:
             layer_name=node_name,
             node_metatype=meta_type,
             node_name=node_name,
-            node_type=node_type,
+            node_type=node_type.split(".")[-1],
         )
         map_nx_node_to_nncf_node[node] = nncf_node
 

nncf/quantization/algorithms/min_max/torch_backend.py

@@ -311,11 +311,13 @@ def get_ignored_metatypes(model_type: ModelType, device: TargetDevice) -> List[O
                 om.PTMaxMetatype,
                 om.PTSqueezeMetatype,
                 om.PTLayerNormMetatype,
+                om.PTAtenLayerNormMetatype,
                 om.PTModuleLayerNormMetatype,
                 om.PTGroupNormMetatype,
                 om.PTModuleGroupNormMetatype,
                 # Batchnorm
                 om.PTBatchNormMetatype,
+                om.PT2BatchNormMetatype,
                 om.PTModuleBatchNormMetatype,
                 # Comparison operations
                 om.PTGreaterEqualMetatype,

nncf/quantization/algorithms/min_max/torch_fx_backend.py

@@ -286,6 +286,7 @@ def get_ignored_metatypes(model_type: ModelType, device: TargetDevice) -> List[O
                 om.PTMaxMetatype,
                 om.PTSqueezeMetatype,
                 om.PTLayerNormMetatype,
+                om.PTAtenLayerNormMetatype,
                 om.PTModuleLayerNormMetatype,
                 om.PTGroupNormMetatype,
                 om.PTModuleGroupNormMetatype,

nncf/torch/dynamic_graph/layer_attributes_handlers.py

@@ -39,7 +39,7 @@
     om.PTConvTranspose1dMetatype, om.PTConvTranspose2dMetatype, om.PTConvTranspose3dMetatype
 )
 LINEAR_OP_NAMES = get_all_aliases(om.PTLinearMetatype)
-BATCHNORM_OP_NAMES = get_all_aliases(om.PTBatchNormMetatype)
+BATCHNORM_OP_NAMES = get_all_aliases(om.PTBatchNormMetatype, om.PT2BatchNormMetatype)
 EMBEDDING_OP_NAMES = get_all_aliases(om.PTEmbeddingMetatype, om.PTEmbeddingBagMetatype)
 GROUP_NORM_OP_NAMES = get_all_aliases(om.PTGroupNormMetatype)
 LAYER_NORM_OP_NAMES = get_all_aliases(om.PTLayerNormMetatype)