experimental/structs.py -> quantization/structs.py

Author: daniil-lyakhov

src/nncf/common/quantization/quantizer_setup.py

@@ -21,11 +21,11 @@
 from nncf.common.quantization.structs import NonWeightQuantizerId
 from nncf.common.quantization.structs import QuantizationScheme as QuantizationMode
 from nncf.common.quantization.structs import QuantizerConfig
+from nncf.common.quantization.structs import TypedQuantizerConfig
 from nncf.common.quantization.structs import UnifiedScaleType
 from nncf.common.quantization.structs import WeightQuantizerId
 from nncf.common.stateful_classes_registry import CommonStatefulClassesRegistry
 from nncf.config.schemata.defaults import QUANTIZATION_NARROW_RANGE
-from nncf.experimental.quantization.structs import ExtendedQuantizerConfig
 
 QuantizationPointId = int
 
@@ -195,14 +195,14 @@ def from_state(cls, state: dict[str, Any]) -> "SingleConfigQuantizationPoint":
         insertion_point_cls = CommonStatefulClassesRegistry.get_registered_class(insertion_point_cls_name)
         insertion_point = insertion_point_cls.from_state(state[cls._state_names.INSERTION_POINT])  # type: ignore
         qconfig_state = state[cls._state_names.QCONFIG]
-        # Need to instantiate ExtendedQuantizerConfig
+        # Need to instantiate TypedQuantizerConfig
         # to support additional fields used by ExecuTorch-specific quantizer configs.
         # TODO (dlyakhov): Refactor and generalize quantizer config deserialization to cleanly handle both
         # standard and extended formats without relying on manual key comparison (ticket 170078).
         if QuantizerConfig().__dict__.keys() == qconfig_state.keys():
             qconfig = QuantizerConfig.from_state(qconfig_state)
         else:
-            qconfig = ExtendedQuantizerConfig.from_state(qconfig_state)
+            qconfig = TypedQuantizerConfig.from_state(qconfig_state)
 
         kwargs = {
             cls._state_names.INSERTION_POINT: insertion_point,

src/nncf/common/quantization/structs.py

@@ -11,7 +11,7 @@
 
 from copy import deepcopy
 from enum import Enum
-from typing import Any, Optional
+from typing import Any, Literal, Optional
 
 import nncf
 from nncf.common.graph import NNCFNode
@@ -22,6 +22,9 @@
 from nncf.config.schemata.defaults import QUANTIZATION_PER_CHANNEL
 from nncf.parameters import StrEnum
 from nncf.parameters import TargetDevice
+from nncf.tensor.definitions import TensorDataType
+
+IntDtype = Literal[TensorDataType.int8, TensorDataType.uint8]
 
 
 @api()
@@ -421,3 +424,41 @@ def get_params_configured_by_preset(self, quant_group: QuantizerGroup) -> dict[s
         if quant_group == QuantizerGroup.ACTIVATIONS and self == QuantizationPreset.MIXED:
             return {"mode": QuantizationScheme.ASYMMETRIC}
         return {"mode": QuantizationScheme.SYMMETRIC}
+
+
+class TypedQuantizerConfig(QuantizerConfig):
+    """
+    Extended configuration class for quantizers, including destination integer dtype.
+    """
+
+    def __init__(
+        self,
+        num_bits: int = QUANTIZATION_BITS,
+        mode: QuantizationScheme = QuantizationScheme.SYMMETRIC,
+        signedness_to_force: Optional[bool] = None,
+        per_channel: bool = QUANTIZATION_PER_CHANNEL,
+        narrow_range: bool = QUANTIZATION_NARROW_RANGE,
+        dest_dtype: IntDtype = TensorDataType.int8,
+    ):
+        """
+        :param num_bits: Bitwidth of the quantization.
+        :param mode: The mode of quantization (symmetric or asymmetric).
+        :param signedness_to_force: True if the quantizer *must* be signed, False if *must* be unsigned,
+            None if the signed/unsigned attribute should be determined based on the incoming activation
+            statistics during range initialization.
+        :param per_channel: True for per-channel quantization, False for per-tensor.
+        :param narrow_range: True if the range of quantized values should be narrowed as compared to the
+            naive case, False if all 2^`num_bits` quantizations should be used.
+        :param dest_dtype: Target integer data type for quantized values.
+        """
+        super().__init__(num_bits, mode, signedness_to_force, per_channel, narrow_range)
+        self.dest_dtype = dest_dtype
+
+    def __str__(self) -> str:
+        retval = super().__str__()
+        return retval + " DestDtype: {self._dest_dtype}"
+
+    def get_state(self) -> dict[str, Any]:
+        state = super().get_state()
+        state["dest_dtype"] = self.dest_dtype
+        return state

src/nncf/experimental/quantization/structs.py

@@ -29,8 +29,8 @@
 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 TypedQuantizerConfig
 from nncf.experimental.quantization.quantizer import Quantizer
-from nncf.experimental.quantization.structs import ExtendedQuantizerConfig
 from nncf.experimental.torch.fx.nncf_graph_builder import GraphConverter
 from nncf.tensor.definitions import TensorDataType
 
@@ -72,7 +72,7 @@ def _get_quantization_points(
         from_node: torch.fx.Node,
         to_nodes: list[torch.fx.Node],
         annotated_model: torch.fx.GraphModule,
-        qconfig: ExtendedQuantizerConfig,
+        qconfig: TypedQuantizerConfig,
     ) -> list[QuantizationPointBase]:
         """
         Creates quantization points based on the nodes and edges.
@@ -167,7 +167,7 @@ def get_quantizer_config_from_annotated_model(annotated: torch.fx.GraphModule) -
                 else QuantizationMode.ASYMMETRIC
             )
             narrow_range = qspec.quant_max - qspec.quant_min == 254
-            qconfig = ExtendedQuantizerConfig(
+            qconfig = TypedQuantizerConfig(
                 mode=mode,
                 signedness_to_force=False,
                 per_channel=per_channel,

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

@@ -25,9 +25,9 @@
 from nncf.common.quantization.quantizer_propagation.structs import QuantizationTrait
 from nncf.common.quantization.structs import QuantizationScheme
 from nncf.common.quantization.structs import QuantizerConfig
+from nncf.common.quantization.structs import TypedQuantizerConfig
 from nncf.experimental.common.tensor_statistics.collectors import REDUCERS_MAP
 from nncf.experimental.common.tensor_statistics.collectors import TensorReducerBase
-from nncf.experimental.quantization.structs import ExtendedQuantizerConfig
 from nncf.experimental.torch.fx.commands import FXApplyTransformationCommand
 from nncf.experimental.torch.fx.model_utils import get_target_point
 from nncf.experimental.torch.fx.transformations import qdq_insertion_transformation_builder
@@ -195,7 +195,7 @@ def _create_quantizer(
     ) -> FakeQuantize:
         per_channel = quantizer_config.per_channel
         dtype = None
-        if isinstance(quantizer_config, ExtendedQuantizerConfig):
+        if isinstance(quantizer_config, TypedQuantizerConfig):
             dtype = quantizer_config.dest_dtype
 
             if dtype not in [TensorDataType.int8, TensorDataType.uint8]:

src/nncf/quantization/algorithms/min_max/torch_fx_backend.py

@@ -19,7 +19,7 @@
 from nncf.common.quantization.quantizer_setup import SingleConfigQuantizerSetup
 from nncf.common.quantization.quantizer_setup import WeightQuantizationInsertionPoint
 from nncf.common.quantization.structs import QuantizerConfig
-from nncf.experimental.quantization.structs import ExtendedQuantizerConfig
+from nncf.common.quantization.structs import TypedQuantizerConfig
 from nncf.torch.dynamic_graph.context import Scope
 from nncf.torch.graph.transformations.commands import PTTargetPoint
 from tests.cross_fw.shared.serialization import check_serialization
@@ -95,7 +95,7 @@ def test_quantizer_setup_serialization():
     scqp_2 = SingleConfigQuantizationPoint(aqip, qc, directly_quantized_operator_node_names=[str(scope)])
     check_serialization(scqp_2)
 
-    ex_qc = ExtendedQuantizerConfig()
+    ex_qc = TypedQuantizerConfig()
     scqp_ex = SingleConfigQuantizationPoint(aqip, ex_qc, directly_quantized_operator_node_names=[str(scope)])
     check_serialization(scqp_ex)
 

tests/torch/quantization/test_serialize_to_json.py

@@ -17,12 +17,12 @@
 import torch
 
 import nncf
+from nncf.common.quantization.structs import IntDtype
 from nncf.common.quantization.structs import QuantizationScheme as QuantizationMode
 from nncf.common.quantization.structs import QuantizerConfig
 from nncf.common.quantization.structs import QuantizerGroup
+from nncf.common.quantization.structs import TypedQuantizerConfig
 from nncf.experimental.common.tensor_statistics.statistics import MinMaxTensorStatistic
-from nncf.experimental.quantization.structs import ExtendedQuantizerConfig
-from nncf.experimental.quantization.structs import IntDtype
 from nncf.quantization.algorithms.min_max.torch_fx_backend import FXMinMaxAlgoBackend
 from nncf.quantization.fake_quantize import FakeQuantizeParameters
 from nncf.quantization.fake_quantize import calculate_quantizer_parameters
@@ -87,7 +87,7 @@ def test_quantizer_params_sym(case_to_test: CaseQuantParams, dtype: Optional[Int
     narrow_range = case_to_test.narrow_range
     mode = QuantizationMode.SYMMETRIC
     signedness_to_force = None
-    qconfig = ExtendedQuantizerConfig(
+    qconfig = TypedQuantizerConfig(
         num_bits=8,
         mode=mode,
         per_channel=per_ch,
@@ -387,7 +387,7 @@ def test_quantizer_params_asym(case_to_test: CaseQuantParams, ref_zp: Union[int,
     per_ch = case_to_test.per_channel
     narrow_range = case_to_test.narrow_range
     mode = QuantizationMode.ASYMMETRIC
-    qconfig = ExtendedQuantizerConfig(
+    qconfig = TypedQuantizerConfig(
         num_bits=8,
         mode=mode,
         per_channel=per_ch,
@@ -452,7 +452,7 @@ def _check_q_min_q_max(quantizer, signed, narrow_range):
     ],
 )
 def test_extended_q_config_non_supported_dest_dtype(dest_dtype):
-    qconfig = ExtendedQuantizerConfig(dest_dtype=dest_dtype)
+    qconfig = TypedQuantizerConfig(dest_dtype=dest_dtype)
     params = FakeQuantizeParameters(-1.0, 1.0, -1.0, 1.0, 255)
     with pytest.raises(nncf.ParameterNotSupportedError):
         FXMinMaxAlgoBackend._create_quantizer(quantizer_config=qconfig, channel_axis=1, parameters=params)