[TorchFX] Do not use torch quantizer in MinMax

Author: daniil-lyakhov

nncf/experimental/quantization/quantizer.py

@@ -11,14 +11,91 @@
 
 from abc import ABC
 from abc import abstractmethod
-from typing import TypeVar
+from enum import Enum
+from typing import Any, Optional, TypeVar
 
 from nncf.common.graph.graph import NNCFGraph
+from nncf.common.quantization.quantizer_setup import QuantizationPointBase
+from nncf.common.quantization.quantizer_setup import QuantizationPointId
 from nncf.common.quantization.quantizer_setup import SingleConfigQuantizerSetup
 
 TModel = TypeVar("TModel")
 
 
+class IntDtype(Enum):
+    INT8 = "INT8"
+    UINT8 = "UINT8"
+
+
+class ExtendedQuantizerSetup(ABC, SingleConfigQuantizerSetup):
+    """
+    Quantizer setup with additional info required to insert
+    quantizers to torch.fx models.
+    """
+
+    @abstractmethod
+    def get_extra_params(self) -> dict[QuantizationPointId, dict[str, Any]]:
+        """
+        Returns extra params
+        """
+
+
+class ExtendedFXQuantizerSetup(ExtendedQuantizerSetup):
+    """
+    Quantizer setup with additional info required to insert
+    quantizers to torch.fx models.
+    """
+
+    QUANTIZER_DTYPE_NAME = "quantizer_dtype"
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._quantization_dtypes: dict[QuantizationPointId, Optional[IntDtype]] = {}
+
+    def add_independent_quantization_point(
+        self, qp: QuantizationPointBase, intermediate_dtype: Optional[IntDtype]
+    ) -> QuantizationPointId:
+        id = super().add_independent_quantization_point(qp)
+        self._quantization_dtypes[id] = intermediate_dtype
+        return id
+
+    def get_extra_params(self) -> dict[int, dict[str, Any]]:
+        return {k: {self.QUANTIZER_DTYPE_NAME: v} for k, v in self._quantization_dtypes.items()}
+
+    def get_state(self) -> dict[str, Any]:
+        """
+        Returns a dictionary with Python data structures (dict, list, tuple, str, int, float, True, False, None) that
+        represents state of the object.
+
+        :return: state of the object
+        """
+        base_state = super().get_state()
+        base_state[self.QUANTIZER_DTYPE_NAME] = {
+            qp_id: dtype.value for qp_id, dtype in self.quantization_points.items()
+        }
+
+    @classmethod
+    def from_state(cls, state: dict[str, Any]) -> "ExtendedFXQuantizerSetup":
+        """
+        Creates the object from its state.
+
+        :param state: Output of `get_state()` method.
+        """
+        state_ = state.copy()
+        dtype_names = state_.pop(cls.QUANTIZER_DTYPE_NAME)
+        super_setup = super().from_state(state_)
+        setup = ExtendedFXQuantizerSetup()
+
+        setup.quantization_points = super_setup.quantization_points
+        setup.unified_scale_groups = super_setup.unified_scale_groups
+        setup.shared_input_operation_set_groups = super_setup.shared_input_operation_set_groups
+        setup._quantization_dtypes = {
+            qp_id: None if name is None else IntDtype[name] for qp_id, name in dtype_names.items()
+        }
+
+        return setup
+
+
 class Quantizer(ABC):
     """
     Quantizer is an interface for the RangeEstimator algorithm
@@ -35,7 +112,7 @@ def transform_prior_quantization(self, model: TModel) -> TModel:
         """
 
     @abstractmethod
-    def get_quantization_setup(self, model: TModel, nncf_graph: NNCFGraph) -> SingleConfigQuantizerSetup:
+    def get_quantization_setup(self, model: TModel, nncf_graph: NNCFGraph) -> ExtendedFXQuantizerSetup:
         """
         Builds SingleConfigQuantizerSetup for the given model.
 

nncf/experimental/torch/fx/quantization/quantize_pt2e.py

@@ -14,7 +14,6 @@
 
 import torch
 import torch.fx
-from torch.ao.quantization.pt2e.duplicate_dq_pass import DuplicateDQPass
 from torch.ao.quantization.pt2e.port_metadata_pass import PortNodeMetaForQDQ
 from torch.ao.quantization.pt2e.utils import _disallow_eval_train
 from torch.ao.quantization.pt2e.utils import _fuse_conv_bn_
@@ -33,6 +32,7 @@
 from nncf.experimental.torch.fx.quantization.quantizer.openvino_quantizer import OpenVINOQuantizer
 from nncf.experimental.torch.fx.quantization.quantizer.torch_ao_adapter import TorchAOQuantizerAdapter
 from nncf.experimental.torch.fx.transformations import QUANTIZE_NODE_TARGETS
+from nncf.experimental.torch.fx.transformations import DuplicateDQPassNoAnnotations
 from nncf.experimental.torch.fx.transformations import compress_post_quantize_transformation
 from nncf.quantization.advanced_parameters import AdvancedBiasCorrectionParameters
 from nncf.quantization.advanced_parameters import AdvancedSmoothQuantParameters
@@ -132,7 +132,7 @@ def quantize_pt2e(
         else:
             constant_fold(quantized_model, _quant_node_constraint)
 
-    pm = PassManager([DuplicateDQPass()])
+    pm = PassManager([DuplicateDQPassNoAnnotations()])
 
     quantized_model = pm(quantized_model).graph_module
     pm = PassManager([PortNodeMetaForQDQ()])

nncf/experimental/torch/fx/quantization/quantizer/openvino_adapter.py

@@ -12,7 +12,7 @@
 import torch.fx
 
 from nncf.common.graph.graph import NNCFGraph
-from nncf.common.quantization.quantizer_setup import SingleConfigQuantizerSetup
+from nncf.experimental.quantization.quantizer import ExtendedFXQuantizerSetup
 from nncf.experimental.quantization.quantizer import Quantizer
 from nncf.experimental.torch.fx.quantization.quantizer.openvino_quantizer import OpenVINOQuantizer
 
@@ -28,5 +28,5 @@ def __init__(self, quantizer: OpenVINOQuantizer):
     def transform_prior_quantization(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
         return self._quantizer.transform_for_annotation(model)
 
-    def get_quantization_setup(self, model: torch.fx.GraphModule, nncf_graph: NNCFGraph) -> SingleConfigQuantizerSetup:
+    def get_quantization_setup(self, model: torch.fx.GraphModule, nncf_graph: NNCFGraph) -> ExtendedFXQuantizerSetup:
         return self._quantizer.get_nncf_quantization_setup(model, nncf_graph)

nncf/experimental/torch/fx/quantization/quantizer/openvino_quantizer.py

@@ -36,6 +36,8 @@
 from nncf.common.quantization.quantizer_setup import SingleConfigQuantizerSetup
 from nncf.common.quantization.structs import QuantizationScheme
 from nncf.common.utils.api_marker import api
+from nncf.experimental.quantization.quantizer import ExtendedFXQuantizerSetup
+from nncf.experimental.quantization.quantizer import IntDtype
 from nncf.experimental.torch.fx.nncf_graph_builder import GraphConverter
 from nncf.experimental.torch.fx.node_utils import get_graph_node_by_name
 from nncf.quantization.advanced_parameters import FP8QuantizationParameters
@@ -135,9 +137,16 @@ def set_ignored_scope(
 
     def get_nncf_quantization_setup(
         self, model: torch.fx.GraphModule, nncf_graph: NNCFGraph
-    ) -> SingleConfigQuantizerSetup:
+    ) -> ExtendedFXQuantizerSetup:
         self._min_max_algo._set_backend_entity(model)
-        return self._min_max_algo.find_quantization_setup(model, nncf_graph)
+        base_setup = self._min_max_algo.find_quantization_setup(model, nncf_graph)
+        dtype_map = {}
+        for id_, qp in base_setup.quantization_points.items():
+            dtype_map[id_] = None if qp.qconfig.mode == QuantizationScheme.SYMMETRIC else IntDtype.UINT8.value
+
+        state = base_setup.get_state()
+        state[ExtendedFXQuantizerSetup.QUANTIZER_DTYPE_NAME] = dtype_map
+        return ExtendedFXQuantizerSetup.from_state(state)
 
     def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
         """

nncf/experimental/torch/fx/quantization/quantizer/torch_ao_adapter.py

@@ -26,10 +26,11 @@
 from nncf.common.quantization.quantizer_setup import ActivationQuantizationInsertionPoint
 from nncf.common.quantization.quantizer_setup import QuantizationPointBase
 from nncf.common.quantization.quantizer_setup import SingleConfigQuantizationPoint
-from nncf.common.quantization.quantizer_setup import SingleConfigQuantizerSetup
 from nncf.common.quantization.quantizer_setup import WeightQuantizationInsertionPoint
 from nncf.common.quantization.structs import QuantizationScheme as QuantizationMode
 from nncf.common.quantization.structs import QuantizerConfig
+from nncf.experimental.quantization.quantizer import ExtendedFXQuantizerSetup
+from nncf.experimental.quantization.quantizer import IntDtype
 from nncf.experimental.quantization.quantizer import Quantizer
 from nncf.experimental.torch.fx.nncf_graph_builder import GraphConverter
 
@@ -47,7 +48,7 @@ def __init__(self, quantizer: TorchAOQuantizer):
     def transform_prior_quantization(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
         return self._quantizer.transform_for_annotation(model)
 
-    def get_quantization_setup(self, model: torch.fx.GraphModule, nncf_graph: NNCFGraph) -> SingleConfigQuantizerSetup:
+    def get_quantization_setup(self, model: torch.fx.GraphModule, nncf_graph: NNCFGraph) -> ExtendedFXQuantizerSetup:
         # Save model and nodes meta before the annotation
         original_meta = model.meta.copy()
         node_name_vs_meta = {}
@@ -116,14 +117,14 @@ def _get_node_args(node: torch.fx.Node) -> tuple[Any, ...]:
         return node.args
 
     @staticmethod
-    def get_quantizer_config_from_annotated_model(annotated: torch.fx.GraphModule) -> SingleConfigQuantizerSetup:
+    def get_quantizer_config_from_annotated_model(annotated: torch.fx.GraphModule) -> ExtendedFXQuantizerSetup:
         """
         Process a torch.fx.GraphModule annotated with quantization specifications
         (e.g., via torch.ao observers) and generates a corresponding NNCF quantization setup object,
         which maps quantization configurations to graph edges.
 
         :param annotated: A torch.fx.GraphModule that has been annotated with Torch quantization observers.
-        :return: A SingleConfigQuantizerSetup containing quantization points derived from the annotated model.
+        :return: A ExtendedFXQuantizerSetup containing quantization points derived from the annotated model.
         """
         edge_or_node_to_qspec = _get_edge_or_node_to_qspec(annotated)
         # Node means all output edges should be quantized.
@@ -142,7 +143,7 @@ def get_quantizer_config_from_annotated_model(annotated: torch.fx.GraphModule) -
                 edge_or_node_to_qspec[edge_or_node], edge_or_node_to_qspec
             )
 
-        q_setup = SingleConfigQuantizerSetup()
+        q_setup = ExtendedFXQuantizerSetup()
         for group_id, edges in group_id_vs_edges.items():
             qspec = group_id_vs_qspec[group_id]
             if qspec is None:
@@ -159,15 +160,15 @@ def get_quantizer_config_from_annotated_model(annotated: torch.fx.GraphModule) -
                 msg = f"Unknown qscheme: {qspec.qscheme}"
                 raise nncf.InternalError(msg)
 
-            signed = qspec.dtype is torch.int8
+            dtype = IntDtype.INT8 if qspec.dtype is torch.int8 else IntDtype.UINT8
             mode = (
                 QuantizationMode.SYMMETRIC
                 if qspec.qscheme in [torch.per_channel_symmetric, torch.per_tensor_symmetric]
                 else QuantizationMode.ASYMMETRIC
             )
             narrow_range = qspec.quant_min % 2 != 0
             qconfig = QuantizerConfig(
-                mode=mode, signedness_to_force=signed, per_channel=per_channel, narrow_range=narrow_range
+                mode=mode, signedness_to_force=False, per_channel=per_channel, narrow_range=narrow_range
             )
 
             joined_edges = defaultdict(list)
@@ -179,7 +180,7 @@ def get_quantizer_config_from_annotated_model(annotated: torch.fx.GraphModule) -
                 qps.extend(TorchAOQuantizerAdapter._get_quantization_points(from_node, to_nodes, annotated, qconfig))
             qp_ids = []
             for qp in qps:
-                qp_ids.append(q_setup.add_independent_quantization_point(qp))
+                qp_ids.append(q_setup.add_independent_quantization_point(qp, dtype))
             if len(qp_ids) > 1:
                 q_setup.register_unified_scale_group(qp_ids)
 

nncf/experimental/torch/fx/transformations.py

@@ -9,13 +9,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import operator
 from copy import copy
 from typing import Any, Callable, Optional, Union
 
 import torch
 import torch.fx
 from torch.ao.quantization.fx.utils import create_getattr_from_value
 from torch.ao.quantization.pt2e.utils import _fuse_conv_bn_
+from torch.fx.node import map_arg
+from torch.fx.passes.infra.pass_base import PassBase
+from torch.fx.passes.infra.pass_base import PassResult
 from torch.quantization.fake_quantize import FakeQuantize
 
 import nncf
@@ -741,3 +745,65 @@ def constraint_fn(node: torch.fx.Node):
         return node.op != "call_function" or node.target not in QUANTIZE_NODE_TARGETS + DEQUANTIZE_NODE_TARGETS
 
     constant_fold(model, constraint_fn=constraint_fn)
+
+
+def _duplicate_dq(gm: torch.fx.GraphModule, dq_node: torch.fx.Node, user: torch.fx.Node):
+    with gm.graph.inserting_after(dq_node):
+        new_node = gm.graph.node_copy(dq_node)
+
+        def maybe_replace_node(n: torch.fx.Node) -> torch.fx.Node:
+            if n == dq_node:
+                return new_node
+            else:
+                return n
+
+        new_args = map_arg(user.args, maybe_replace_node)
+        new_kwargs = map_arg(user.kwargs, maybe_replace_node)
+        user.args = new_args
+        user.kwargs = new_kwargs
+
+
+def _is_sym_size_node(node: torch.fx.Node):
+    return (
+        node.op == "call_function"
+        and node.target == torch.ops.aten.sym_size.default
+        or node.target == torch.ops.aten.sym_numel.default
+        or node.target == torch.ops.aten.sym_numel
+        or node.target == torch.ops.aten.sym_size
+    )
+
+
+def _filter_sym_size_users(node: torch.fx.Node) -> list[torch.fx.Node]:
+    node_users = list(filter((lambda x: (_is_sym_size_node(x) is False)), node.users))
+    return node_users
+
+
+class DuplicateDQPassNoAnnotations(PassBase):
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        for node in graph_module.graph.nodes:
+            if node.op == "call_function" and node.target in DEQUANTIZE_NODE_TARGETS:
+                dq_users = _filter_sym_size_users(node)
+                if len(dq_users) <= 1:
+                    continue
+                # Do not duplicate dq for dynamic quantization
+                # Pattern: choose_qparam - getitem - q - dq
+                q_node = node.args[0]
+                if q_node.op == "call_function" and q_node.target in QUANTIZE_NODE_TARGETS:
+                    getitem_node = q_node.args[1]
+                    if (
+                        isinstance(getitem_node, torch.fx.node.Node)
+                        and getitem_node.op == "call_function"
+                        and getitem_node.target == operator.getitem
+                    ):
+                        choose_qparam_node = getitem_node.args[0]
+                        if (
+                            isinstance(choose_qparam_node, torch.fx.node.Node)
+                            and choose_qparam_node.op == "call_function"
+                            and choose_qparam_node.target == torch.ops.quantized_decomposed.choose_qparams.tensor
+                        ):
+                            continue
+                for user in dq_users:
+                    _duplicate_dq(graph_module, node, user)
+        graph_module.graph.eliminate_dead_code()
+        graph_module.recompile()
+        return PassResult(graph_module, True)