mct-nightly 2.4.0.20250629.706__py3-none-any.whl → 2.4.0.20250701.185106__py3-none-any.whl

model_compression_toolkit/core/common/network_editors/actions.py

@@ -20,14 +20,8 @@ from typing import Callable
 from mct_quantizers import QuantizationMethod
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.logger import Logger
-
-
-from model_compression_toolkit.core.common.framework_info import get_fw_info
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
-from model_compression_toolkit.core.common.quantization.quantization_params_fn_selection import \
-    get_activation_quantization_params_fn, get_weights_quantization_params_fn
-from model_compression_toolkit.core.common.quantization.quantization_fn_selection import \
-    get_weights_quantization_fn
+
 
 _EditRule = namedtuple('EditRule', 'filter action')
 
@@ -174,47 +168,6 @@ class ChangeFinalActivationQuantConfigAttr(BaseAction):
                 node.final_activation_quantization_cfg.set_quant_config_attr(parameter_name, parameter_value)
 
 
-class ChangeQuantizationParamFunction(BaseAction):
-    """
-    Class ChangeQuantizationParamFunction to change a node's weights/activations quantization params function.
-    """
-
-    def __init__(self,
-                 attr_name: str = None,
-                 activation_quantization_params_fn: Callable = None,
-                 weights_quantization_params_fn: Callable = None):
-        """
-        Init a ChangeQuantizationParamFunction object.
-
-        Args:
-            attr_name: The weights attribute's name to set the weights quantization params function for (if setting weights params).
-            activation_quantization_params_fn: a params function for a node's activations.
-            weights_quantization_params_fn: a params function for a node's weights.
-        """
-        self.activation_quantization_params_fn = activation_quantization_params_fn
-        self.weights_quantization_params_fn = weights_quantization_params_fn
-        self.attr_name = attr_name
-
-    def apply(self, node: BaseNode, graph):
-        """
-        Change the node's weights/activations quantization params function.
-
-        Args:
-            node: Node object to change its quantization params function.
-            graph: Graph to apply the action on.
-
-        Returns:
-            The node after its quantization params function has been modified.
-        """
-        for nqc in node.candidates_quantization_cfg:
-            if self.activation_quantization_params_fn is not None:
-                nqc.activation_quantization_cfg.set_activation_quantization_params_fn(
-                    self.activation_quantization_params_fn)
-            if self.weights_quantization_params_fn is not None:
-                (nqc.weights_quantization_cfg.get_attr_config(self.attr_name)
-                 .set_weights_quantization_params_fn(self.weights_quantization_params_fn))
-
-
 class ChangeFinalActivationQuantizationMethod(BaseAction):
     """
     Class ChangeFinalActivationQuantizationMethod to change a node's weights/activations quantizer function.
@@ -243,16 +196,6 @@ class ChangeFinalActivationQuantizationMethod(BaseAction):
         """
 
         if self.activation_quantization_method is not None and node.final_activation_quantization_cfg is not None:
-
-            activation_quantization_params_fn = get_activation_quantization_params_fn(
-                self.activation_quantization_method)
-
-            node.final_activation_quantization_cfg.set_activation_quantization_params_fn(
-                activation_quantization_params_fn)
-
-            activation_quantization_fn = get_fw_info().activation_quantizer_mapping.get(self.activation_quantization_method)
-
-            node.final_activation_quantization_cfg.set_activation_quantization_fn(activation_quantization_fn)
             node.final_activation_quantization_cfg.activation_quantization_method = self.activation_quantization_method
 
 
@@ -281,23 +224,12 @@ class ChangeCandidatesActivationQuantizationMethod(BaseAction):
         """
         if self.activation_quantization_method is not None:
             for qc in node.candidates_quantization_cfg:
-                activation_quantization_params_fn = get_activation_quantization_params_fn(
-                    self.activation_quantization_method)
-
-                qc.activation_quantization_cfg.set_activation_quantization_params_fn(activation_quantization_params_fn)
-                activation_quantization_fn = get_fw_info().activation_quantizer_mapping.get(
-                    self.activation_quantization_method)
-
-                if activation_quantization_fn is None:
-                    Logger.critical('Unknown activation quantization method specified.')  # pragma: no cover
-
-                qc.activation_quantization_cfg.set_activation_quantization_fn(activation_quantization_fn)
                 qc.activation_quantization_cfg.activation_quantization_method = self.activation_quantization_method
 
 
 class ChangeFinalWeightsQuantizationMethod(BaseAction):
     """
-    Class ChangeFinalWeightsQuantizationMethod to change a node's weights/activations quantizer function.
+    Class ChangeFinalWeightsQuantizationMethod to change a node's weights/activations quantizer method.
     """
 
     def __init__(self, attr_name: str, weights_quantization_method=None):
@@ -323,21 +255,8 @@ class ChangeFinalWeightsQuantizationMethod(BaseAction):
         """
 
         if self.weights_quantization_method is not None and node.final_weights_quantization_cfg is not None:
-
-            weights_quantization_params_fn = get_weights_quantization_params_fn(self.weights_quantization_method)
-
-            (node.final_weights_quantization_cfg.get_attr_config(self.attr_name)
-             .set_weights_quantization_params_fn(weights_quantization_params_fn))
-
-            weights_quantization_fn = get_weights_quantization_fn(self.weights_quantization_method)
-
-            if weights_quantization_fn is None:
-                Logger.critical('Unknown weights quantization method specified.')  # pragma: no cover
-
-            (node.final_weights_quantization_cfg.get_attr_config(self.attr_name)
-             .set_weights_quantization_fn(weights_quantization_fn))
-            node.final_weights_quantization_cfg.get_attr_config(self.attr_name).weights_quantization_method = \
-                self.weights_quantization_method
+            attr_config = node.final_weights_quantization_cfg.get_attr_config(self.attr_name)
+            attr_config.weights_quantization_method = self.weights_quantization_method
 
 
 class ChangeCandidatesWeightsQuantizationMethod(BaseAction):
@@ -370,18 +289,7 @@ class ChangeCandidatesWeightsQuantizationMethod(BaseAction):
 
         if self.weights_quantization_method is not None:
             for qc in node.candidates_quantization_cfg:
-
-                weights_quantization_params_fn = get_weights_quantization_params_fn(self.weights_quantization_method)
-
                 attr_qc = qc.weights_quantization_cfg.get_attr_config(self.attr_name)
-                attr_qc.set_weights_quantization_params_fn(weights_quantization_params_fn)
-
-                weights_quantization_fn = get_weights_quantization_fn(self.weights_quantization_method)
-
-                if weights_quantization_fn is None:
-                    Logger.critical('Unknown weights quantization method specified.')  # pragma: no cover
-
-                attr_qc.set_weights_quantization_fn(weights_quantization_fn)
                 attr_qc.weights_quantization_method = self.weights_quantization_method
 
 

model_compression_toolkit/core/common/quantization/bit_width_config.py

@@ -19,8 +19,8 @@ from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.matchers.node_matcher import BaseNodeMatcher
 from model_compression_toolkit.logger import Logger
 
-from model_compression_toolkit.core.common.graph.base_node import WeightAttrT
-from model_compression_toolkit.target_platform_capabilities.constants import POS_ATTR
+from model_compression_toolkit.core.common.graph.base_node import WeightAttrT, BaseNode
+from model_compression_toolkit.target_platform_capabilities.constants import POSITIONAL_ATTR
 
 
 @dataclass
@@ -95,7 +95,7 @@ class BitWidthConfig:
         for attr, bit_width, filter in zip (attrs, bit_widths, filters):
             self.manual_weights_bit_width_selection_list += [ManualWeightsBitWidthSelection(filter, bit_width, attr)]
 
-    def get_nodes_to_manipulate_activation_bit_widths(self, graph: Graph) -> Dict:
+    def get_nodes_activation_bit_widths(self, graph: Graph) -> Dict[BaseNode, int]:
         """
         Retrieve nodes from the graph that need their bit-widths for activation changed according to the manual bit-width selections.
 
@@ -108,7 +108,7 @@ class BitWidthConfig:
         activation_nodes_to_change_bit_width = self._construct_node_to_new_activation_bit_mapping(graph)
         return activation_nodes_to_change_bit_width
 
-    def get_nodes_to_manipulate_weights_bit_widths(self, graph: Graph) -> Dict:
+    def get_nodes_weights_bit_widths(self, graph: Graph) -> Dict[BaseNode, Dict[str, int]]:
         """
         Retrieve nodes from the graph that need their bit-widths for weights changed according to the manual bit-width selections.
 
@@ -166,7 +166,7 @@ class BitWidthConfig:
             attrs = BitWidthConfig._expand_to_list_core(filters, attrs)
         return attrs, bit_widths, filters
 
-    def _construct_node_to_new_activation_bit_mapping(self, graph) -> Dict:
+    def _construct_node_to_new_activation_bit_mapping(self, graph) -> Dict[BaseNode, int]:
         """
         Retrieve nodes from the graph that need their activation bit-widths changed according to the manual bit-width selections.
 
@@ -192,7 +192,7 @@ class BitWidthConfig:
                 unit_nodes_to_change_bit_width.update({n: manual_bit_width_selection.bit_width})
         return unit_nodes_to_change_bit_width
 
-    def _construct_node_to_new_weights_bit_mapping(self, graph) -> Dict:
+    def _construct_node_to_new_weights_bit_mapping(self, graph) -> Dict[BaseNode, Dict[str, int]]:
         """
         Retrieve nodes from the graph that need their weights bit-widths changed according to the manual bit-width selections.
 
@@ -212,7 +212,7 @@ class BitWidthConfig:
                     f"to change their bit width to {manual_bit_width_selection.bit_width}.")
 
             for n in filtered_nodes:
-                attr_to_change_bit_width = []
+                attr_to_change_bit_width = {}
                 
                 attrs_str = n.get_node_weights_attributes()
                 if len(attrs_str) == 0:
@@ -225,8 +225,8 @@ class BitWidthConfig:
                             attr.append(attr_str)
                     # this is a positional attribute, so it needs to be handled separately.
                     # Search manual_bit_width_selection's attribute that contain the POS_ATTR string.
-                    elif isinstance(attr_str, int) and POS_ATTR in manual_bit_width_selection.attr:
-                        attr.append(POS_ATTR)
+                    elif isinstance(attr_str, int) and POSITIONAL_ATTR in manual_bit_width_selection.attr:
+                        attr.append(POSITIONAL_ATTR)
                 if len(attr) == 0:
                     Logger.critical(f'The requested attribute {manual_bit_width_selection.attr} to change the bit width for {n} does not exist.')
 
@@ -239,7 +239,7 @@ class BitWidthConfig:
                                 f"Node {n} has an existing manual bit width configuration of {manual_bit_width_selection.attr}." 
                                 f"A new manual configuration request of {manual_bit_width_selection.bit_width} has been received, and the previous value is being overridden.")
 
-                attr_to_change_bit_width.append([manual_bit_width_selection.bit_width, manual_bit_width_selection.attr])
+                attr_to_change_bit_width[manual_bit_width_selection.attr] = manual_bit_width_selection.bit_width
                 unit_nodes_to_change_bit_width.update({n: attr_to_change_bit_width})
 
         return unit_nodes_to_change_bit_width

model_compression_toolkit/core/common/quantization/candidate_node_quantization_config.py

@@ -12,73 +12,133 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Callable, List, Tuple
+import copy
+from dataclasses import dataclass, InitVar
+from typing import Callable, List, Optional
 
-from model_compression_toolkit.core import QuantizationConfig
-from model_compression_toolkit.core.common.framework_info import ChannelAxisMapping
 from model_compression_toolkit.core.common.quantization.node_quantization_config import BaseNodeQuantizationConfig, \
-    NodeWeightsQuantizationConfig, NodeActivationQuantizationConfig
-from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import AttributeQuantizationConfig, \
-    OpQuantizationConfig
-from model_compression_toolkit.logger import Logger
+    NodeWeightsQuantizationConfig, NodeActivationQuantizationConfig, ActivationQuantizationMode
 
 
-##########################################
-# Every node holds a quantization configuration
-# for its weights quantization, and a different quantization
-# configuration for its activation quantization configuration.
-##########################################
-
+@dataclass(eq=True)
 class CandidateNodeQuantizationConfig(BaseNodeQuantizationConfig):
     """
-    Class for representing candidate node configuration, which includes weights and activation configuration combined.
+    Candidate quantization configuration for a node.
     """
+    activation_quantization_cfg: NodeActivationQuantizationConfig
+    # TODO irena: None is passed in several places, need to check if it's handled properly or it's only passed in cases
+    #  that do not affect anything (my guess is it's the second).
+    #  I think in general it makes more sense to set it to None when there are no weights, and maybe when all weights
+    #  are unquantized, and handle it properly everywhere.
+    weights_quantization_cfg: Optional[NodeWeightsQuantizationConfig]
+
+
+# TODO irena: currently all code still looks at candidates_quantization_cfg as previously, so this is just an initial
+#  implementation. For now base config is completely separated from candidates (base config must be equal to one of the
+#  candidates, but we create a separate copy), and updating in place is allowed. Also we require quantization mode to
+#  be identical between all configs.
+@dataclass
+class NodeQuantizationConfig:
+    # quantization config for single precision
+    base_quantization_cfg: CandidateNodeQuantizationConfig
+    # quantization candidate configs for mixed precision
+    candidates_quantization_cfg: List[CandidateNodeQuantizationConfig]
+
+    validate: InitVar[bool] = True
+
+    def update_all(self, update_fn: Callable[[CandidateNodeQuantizationConfig], None], remove_duplicates: bool = True):
+        """
+        Apply update function on the base config and all candidates configs.
+
+        Args:
+            update_fn: function to apply.
+            remove_duplicates: remove duplicate candidates.
+        """
+        if self.base_quantization_cfg:
+            update_fn(self.base_quantization_cfg)
+        for cfg in self.candidates_quantization_cfg:
+            update_fn(cfg)
+        if remove_duplicates:
+            self.remove_duplicates()
 
-    def __init__(self,
-                 qc: QuantizationConfig = None,
-                 op_cfg: OpQuantizationConfig = None,
-                 activation_quantization_cfg: NodeActivationQuantizationConfig = None,
-                 activation_quantization_fn: Callable = None,
-                 activation_quantization_params_fn: Callable = None,
-                 weights_quantization_cfg: NodeWeightsQuantizationConfig = None,
-                 weights_channels_axis: ChannelAxisMapping = None,
-                 node_attrs_list: List[str] = None):
+    def update_activation_quantization_mode(self, mode: ActivationQuantizationMode):
         """
+        Update activation quantization mode for the base config and all candidates configs.
 
         Args:
-            qc: QuantizationConfig to create the node's config from.
-            op_cfg: OpQuantizationConfig of the node with quantizers types to use when creating node quantization configuration.
-            activation_quantization_cfg: An option to pass a NodeActivationQuantizationConfig to create a new config from.
-            activation_quantization_fn: Function to use when quantizing the node's activations.
-            activation_quantization_params_fn: Function to use when computing the threshold for quantizing a node's activations.
-            weights_quantization_cfg: An option to pass a NodeWeightsQuantizationConfig to create a new config from.
-            weights_channels_axis: Axis to quantize a node's weights attribute when quantizing per-channel.
-            node_attrs_list: A list of the node's weights attributes names.
+            mode: quantization mode.
+        """
+        def fn(c):
+            c.activation_quantization_cfg.quant_mode = mode
+
+        self.update_all(fn)
+
+    def disable_weights_quantization(self):
+        """
+        Disable all weights quantization for the base config and all candidates configs.
+        """
+        self.update_all(lambda c: c.weights_quantization_cfg.disable_all_weights_quantization())
+
+    def get_activation_quant_mode(self) -> ActivationQuantizationMode:
+        """
+        Retrieve activation quantization mode.
+
+        Returns:
+            Activation quantization mode.
+
+        Raises:
+            ValueError if not all candidates contain the same mode.
+        """
+        self._validate_consistent_activation_quant_mode()
+        return self.base_quantization_cfg.activation_quantization_cfg.quant_mode
+
+    def remove_duplicates(self):
+        """
+        Remove duplicate candidates. First candidate among duplicates is kept, and the order is preserved.
         """
+        uniq_qcs = []
+        for qc in self.candidates_quantization_cfg:
+            if qc not in uniq_qcs:
+                uniq_qcs.append(qc)
+        self.candidates_quantization_cfg = uniq_qcs
 
-        if activation_quantization_cfg is not None:
-            self.activation_quantization_cfg = activation_quantization_cfg
-        else:
-            if any(v is None for v in (qc, op_cfg, activation_quantization_fn, activation_quantization_params_fn)):  # pragma: no cover
-                Logger.critical(
-                    "Missing required arguments to initialize a node activation quantization configuration. "
-                    "Ensure QuantizationConfig, OpQuantizationConfig, activation quantization function, "
-                    "and parameters function are provided.")
-            self.activation_quantization_cfg = (
-                NodeActivationQuantizationConfig(qc=qc,
-                                                 op_cfg=op_cfg,
-                                                 activation_quantization_fn=activation_quantization_fn,
-                                                 activation_quantization_params_fn=activation_quantization_params_fn))
-
-        if weights_quantization_cfg is not None:
-            self.weights_quantization_cfg = weights_quantization_cfg
-        elif all(v is not None for v in (qc, op_cfg, node_attrs_list)):
-            self.weights_quantization_cfg = NodeWeightsQuantizationConfig(qc=qc,
-                                                                          op_cfg=op_cfg,
-                                                                          weights_channels_axis=weights_channels_axis,
-                                                                          node_attrs_list=node_attrs_list)
-        else:
-            self.weights_quantization_cfg = None
-            Logger.debug("Setting weights quantization config as None during CandidateNodeQuantizationConfig creation."
-                         "Notice, this should happen only for FLN nodes.")
+    def __post_init__(self, validate=True):
+        if validate:
+            if not any(self.base_quantization_cfg == qc for qc in self.candidates_quantization_cfg):
+                raise ValueError('Candidates should contain the base config.')
+            self._validate_consistent_activation_quant_mode()
+            self._validate_consistent_weights_quant_mode()
+        # TODO irena
+        # for now make sure they are separate objects so that one doesnt inadvertently modify the other
+        if any(self.base_quantization_cfg is qc for qc in self.candidates_quantization_cfg):
+            self.base_quantization_cfg = copy.deepcopy(self.base_quantization_cfg)
 
+    def _validate_consistent_activation_quant_mode(self):
+        """
+        Validate that base config and all candidates configs contain identical activation quantization mode.
+
+        Raises:
+            ValueError if activation quantization mode is not consistent.
+        """
+        activation_quant_mode = self.base_quantization_cfg.activation_quantization_cfg.quant_mode
+        if any(qc.activation_quantization_cfg.quant_mode != activation_quant_mode
+               for qc in self.candidates_quantization_cfg):
+            raise ValueError('Quantization candidates with different quantization modes are not currently supported.')
+
+    def _validate_consistent_weights_quant_mode(self):
+        """
+        Validate that base config and all candidates configs contain identical weights quantization mode per attribute,
+        i.e. quantization for each attribute should either be enabled in all configs, or disabled in all configs.
+
+        Raises:
+            ValueError if weights quantization is not consistent.
+        """
+        def get_weights_mode(qc):
+            # in graph fuser weights_quantization_cfg is set to None
+            if qc.weights_quantization_cfg is None:
+                return None
+            return {attr: attr_cfg.enable_weights_quantization for attr, attr_cfg
+                    in qc.weights_quantization_cfg.get_all_weight_attrs_configs().items()}
+        if any(get_weights_mode(self.base_quantization_cfg) != get_weights_mode(qc)
+               for qc in self.candidates_quantization_cfg):
+            raise ValueError('Quantization candidates with different quantization modes are not currently supported.')

model_compression_toolkit/core/common/quantization/filter_nodes_candidates.py

@@ -33,7 +33,7 @@ def filter_nodes_candidates(graph: Graph):
     """
     nodes = list(graph.nodes)
     for n in nodes:
-        n.candidates_quantization_cfg = filter_node_candidates(node=n)
+        n.quantization_cfg.candidates_quantization_cfg = filter_node_candidates(node=n)
 
     return graph