mct-nightly 2.3.0.20250415.557__tar.gz → 2.3.0.20250417.547__tar.gz

{mct_nightly-2.3.0.20250415.557 → mct_nightly-2.3.0.20250417.547}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mct-nightly
-Version: 2.3.0.20250415.557
+Version: 2.3.0.20250417.547
 Summary: A Model Compression Toolkit for neural networks
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License

{mct_nightly-2.3.0.20250415.557 → mct_nightly-2.3.0.20250417.547}/mct_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mct-nightly
-Version: 2.3.0.20250415.557
+Version: 2.3.0.20250417.547
 Summary: A Model Compression Toolkit for neural networks
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License

{mct_nightly-2.3.0.20250415.557 → mct_nightly-2.3.0.20250417.547}/model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "2.3.0.20250415.000557"
+__version__ = "2.3.0.20250417.000547"

{mct_nightly-2.3.0.20250415.557 → mct_nightly-2.3.0.20250417.547}/model_compression_toolkit/core/common/framework_info.py

@@ -22,6 +22,12 @@ from mct_quantizers import QuantizationMethod
 from model_compression_toolkit.defaultdict import DefaultDict
 
 
+# Default value to use for ops without kernel.
+# This is a weird default, but it's used all over the place, so for now only extract it to const so that it can be
+# referenced by variable instead of hard-coded.
+DEFAULT_KERNEL_ATTRIBUTES = [None]
+
+
 class ChannelAxis(Enum):
     """
 

{mct_nightly-2.3.0.20250415.557 → mct_nightly-2.3.0.20250417.547}/model_compression_toolkit/core/common/graph/base_graph.py

@@ -16,7 +16,7 @@ from collections import namedtuple
 
 from copy import copy, deepcopy
 from functools import wraps
-from typing import List, Tuple, Any, Callable
+from typing import List, Tuple, Any, Callable, Dict
 
 import networkx as nx
 import numpy as np
@@ -684,7 +684,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
                 sorted_configurable_nodes.append(n)
         return sorted_configurable_nodes
 
-    def get_min_candidates_config(self, fw_info: FrameworkInfo) -> List[int]:
+    def get_min_candidates_config(self, fw_info: FrameworkInfo) -> Dict[BaseNode, int]:
         """
         Builds a minimal configuration.
         Note: we assume that a minimal configuration exists, i.e., each configurable node has exactly one candidate
@@ -694,18 +694,13 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         Args:
             fw_info: fw_info: FrameworkInfo object with information about the specific framework's model.
 
-        Returns: A list of candidate for each node (list on indices)
+        Returns:
+            A dict from layer to an index of its minimal candidate.
         """
-
         conf_sorted_nodes = self.get_configurable_sorted_nodes(fw_info)
-        min_cfg_candidates = [n.find_min_candidates_indices() for n in conf_sorted_nodes]  # list of lists of indices
-
-        assert all([len(lst) == 1 for lst in min_cfg_candidates]), \
-            f"A minimal config candidate must be defined, but some node have multiple potential minimal candidates"
-
-        return [lst[0] for lst in min_cfg_candidates]
+        return {n: n.find_min_candidate_index() for n in conf_sorted_nodes}
 
-    def get_max_candidates_config(self, fw_info: FrameworkInfo) -> List[int]:
+    def get_max_candidates_config(self, fw_info: FrameworkInfo) -> Dict[BaseNode, int]:
         """
         Builds a maximal configuration.
         Note: we assume that a maximal configuration exists, i.e., each configurable node has exactly one candidate
@@ -715,16 +710,11 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         Args:
             fw_info: fw_info: FrameworkInfo object with information about the specific framework's model.
 
-        Returns: A list of candidate for each node (list on indices)
+        Returns:
+            A dict from layer to an index of its maximal candidate.
         """
-
         conf_sorted_nodes = self.get_configurable_sorted_nodes(fw_info)
-        max_cfg_candidates = [n.find_max_candidates_indices() for n in conf_sorted_nodes]  # list of lists of indices
-
-        assert all([len(lst) == 1 for lst in max_cfg_candidates]), \
-            f"A maximal config candidate must be defined, but some node have multiple potential maximal candidates"
-
-        return [lst[0] for lst in max_cfg_candidates]
+        return {n: n.find_max_candidate_index() for n in conf_sorted_nodes}
 
     def get_final_weights_config(self, fw_info: FrameworkInfo) -> List[Tuple[BaseNode, int]]:
         """

{mct_nightly-2.3.0.20250415.557 → mct_nightly-2.3.0.20250417.547}/model_compression_toolkit/core/common/graph/base_node.py

@@ -484,49 +484,35 @@ class BaseNode:
         # for scalar shape (None,) prod returns 1
         return sum([np.prod([x for x in output_shape if x is not None]) for output_shape in output_shapes])
 
-    def find_min_candidates_indices(self) -> List[int]:
+    def find_min_candidate_index(self) -> int:
         """
-        Returns a list with potential minimal candidates.
-        A potential minimal candidate is a candidate which its weights_n_bits and activation_n_bits pair is
-        on the Pareto Front, i.e., there is no other candidate that its n_bits pair exceeds in both entries.
-
-        Returns: A list of indices of potential minimal candidates.
-
-        """
-
-        # We assume that the candidates are sorted according to weights_n_bits first and activation_n_bits second
-        # First, we add the last candidate to the set of minimal candidates (candidate, index)
-        first_min = (len(self.candidates_quantization_cfg) - 1,
-                     self.candidates_quantization_cfg[-1].activation_quantization_cfg.activation_n_bits)
-        min_candidates = [first_min]
-
-        # Iterate over all other candidates, and add ones with higher weights_n_bits but smaller activation_n_bits
-        for i, c in reversed(list(enumerate(self.candidates_quantization_cfg))):
-            if c.activation_quantization_cfg.activation_n_bits < first_min[1]:
-                min_candidates.append((i, c))
-
-        return [i for i, a_n_bits in min_candidates]
-
-    def find_max_candidates_indices(self) -> List[int]:
+        Returns:
+            The index of the minimal bit-width candidate.
         """
-        Returns a list with potential maximal candidates.
-        A potential maximal candidate is a candidate which its weights_n_bits and activation_n_bits pair is
-        on the Pareto Front, i.e., there is no other candidates that its n_bits pair is lower in both entries.
+        aw_nbits = [(c.activation_quantization_cfg.activation_n_bits,
+                     *[v.weights_n_bits for v in c.weights_quantization_cfg.get_all_weight_attrs_configs().values()])
+                    for c in self.candidates_quantization_cfg]
+        min_nbits = min(aw_nbits)
+        min_ind = [i for i, nb in enumerate(aw_nbits) if min_nbits == nb]
+        # check that no other candidate has a lower nbit for any weight
+        if len(min_ind) > 1 or any(nb[i] < min_nbits[i] for i in range(len(min_nbits)) for nb in aw_nbits):
+            raise ValueError('Expected exactly one candidate with min activation and min weights.')
+        return min_ind[0]
 
-        Returns: A list of indices of potential maximal candidates.
+    def find_max_candidate_index(self) -> int:
         """
-
-        # We assume that the candidates are sorted according to weights_n_bits first and activation_n_bits second
-        # First, we add the first candidate to the set of maximal candidates (candidate, index)
-        first_max = (0, self.candidates_quantization_cfg[0].activation_quantization_cfg.activation_n_bits)
-        max_candidates = [first_max]
-
-        # Iterate over all other candidates, and add ones with higher weights_n_bits but smaller activation_n_bits
-        for i, c in enumerate(self.candidates_quantization_cfg):
-            if c.activation_quantization_cfg.activation_n_bits > first_max[1]:
-                max_candidates.append((i, c))
-
-        return [i for i, a_n_bits in max_candidates]
+        Returns:
+            The index of the maximal bit-width candidate.
+        """
+        aw_nbits = [(c.activation_quantization_cfg.activation_n_bits,
+                     *[v.weights_n_bits for v in c.weights_quantization_cfg.get_all_weight_attrs_configs().values()])
+                    for c in self.candidates_quantization_cfg]
+        max_nbits = max(aw_nbits)
+        max_ind = [i for i, nb in enumerate(aw_nbits) if max_nbits == nb]
+        # check that no other candidate has a higher nbit for any weight
+        if len(max_ind) > 1 or any(nb[i] > max_nbits[i] for i in range(len(max_nbits)) for nb in aw_nbits):
+            raise ValueError('Expected exactly one candidate with max activation and max weights.')
+        return max_ind[0]
 
     def get_unique_weights_candidates(self, attr: str) -> List[Any]:
         """

{mct_nightly-2.3.0.20250415.557 → mct_nightly-2.3.0.20250417.547}/model_compression_toolkit/core/common/mixed_precision/mixed_precision_ru_helper.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import List, Set, Dict, Tuple
+from typing import Set, Dict, Tuple
 
 import numpy as np
 
 from model_compression_toolkit.core import FrameworkInfo
-from model_compression_toolkit.core.common import Graph
+from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import \
     RUTarget
@@ -36,7 +36,7 @@ class MixedPrecisionRUHelper:
         self.fw_impl = fw_impl
         self.ru_calculator = ResourceUtilizationCalculator(graph, fw_impl, fw_info)
 
-    def compute_utilization(self, ru_targets: Set[RUTarget], mp_cfg: List[int]) -> Dict[RUTarget, np.ndarray]:
+    def compute_utilization(self, ru_targets: Set[RUTarget], mp_cfg: Dict[BaseNode, int]) -> Dict[RUTarget, np.ndarray]:
         """
         Compute utilization of requested targets for a specific configuration:
           for weights and bops - total utilization,
@@ -74,7 +74,7 @@ class MixedPrecisionRUHelper:
                                                  f'Requested {ru_targets}')
         return ru_dict
 
-    def get_quantization_candidates(self, mp_cfg) \
+    def get_quantization_candidates(self, mp_cfg: Dict[BaseNode, int]) \
             -> Tuple[Dict[str, NodeActivationQuantizationConfig], Dict[str, NodeWeightsQuantizationConfig]]:
         """
         Retrieve quantization candidates objects for weights and activations from the configuration list.
@@ -86,8 +86,7 @@ class MixedPrecisionRUHelper:
             A mapping between nodes to weights quantization config, and a mapping between nodes and activation
             quantization config.
         """
-        mp_nodes = self.graph.get_configurable_sorted_nodes(self.fw_info)
-        node_qcs = {n: n.candidates_quantization_cfg[mp_cfg[i]] for i, n in enumerate(mp_nodes)}
+        node_qcs = {n: n.candidates_quantization_cfg[candidate_idx] for n, candidate_idx in mp_cfg.items()}
         act_qcs = {n.name: cfg.activation_quantization_cfg for n, cfg in node_qcs.items()}
         w_qcs = {n.name: cfg.weights_quantization_cfg for n, cfg in node_qcs.items()}
         return act_qcs, w_qcs

{mct_nightly-2.3.0.20250415.557 → mct_nightly-2.3.0.20250417.547}/model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py

@@ -14,10 +14,10 @@
 # ==============================================================================
 
 from enum import Enum
-from typing import List, Callable
+from typing import List, Callable, Dict
 
 from model_compression_toolkit.core import MixedPrecisionQuantizationConfig
-from model_compression_toolkit.core.common import Graph
+from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.hessian import HessianInfoService
@@ -100,11 +100,13 @@ def search_bit_width(graph: Graph,
                                                  fw_impl,
                                                  se,
                                                  target_resource_utilization)
-    result_bit_cfg = search_manager.search()
+    nodes_bit_cfg = search_manager.search()
 
     graph.skip_validation_check = False
 
     if mp_config.refine_mp_solution:
-        result_bit_cfg = greedy_solution_refinement_procedure(result_bit_cfg, search_manager, target_resource_utilization)
+        nodes_bit_cfg = greedy_solution_refinement_procedure(nodes_bit_cfg, search_manager, target_resource_utilization)
 
-    return result_bit_cfg
+    topo_bit_cfg = [nodes_bit_cfg[n] for n in graph.get_configurable_sorted_nodes(fw_info)]
+    assert len(topo_bit_cfg) == len(nodes_bit_cfg)
+    return topo_bit_cfg

{mct_nightly-2.3.0.20250415.557 → mct_nightly-2.3.0.20250417.547}/model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import itertools
+
 import copy
 from collections import defaultdict
 
@@ -21,7 +23,6 @@ from typing import Dict, List, Tuple
 
 import numpy as np
 
-from model_compression_toolkit.constants import EPS
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
@@ -75,34 +76,44 @@ class MixedPrecisionSearchManager:
         self.target_resource_utilization = target_resource_utilization
 
         self.mp_topo_configurable_nodes = self.mp_graph.get_configurable_sorted_nodes(fw_info)
-        self.layer_to_bitwidth_mapping = self.get_search_space()
 
         self.ru_targets = target_resource_utilization.get_restricted_targets()
         self.ru_helper = MixedPrecisionRUHelper(self.mp_graph, fw_info, fw_impl)
 
-        self.min_ru_config = self.mp_graph.get_min_candidates_config(fw_info)
-        self.max_ru_config = self.mp_graph.get_max_candidates_config(fw_info)
+        self.min_ru_config: Dict[BaseNode, int] = self.mp_graph.get_min_candidates_config(fw_info)
+        self.max_ru_config: Dict[BaseNode, int] = self.mp_graph.get_max_candidates_config(fw_info)
         self.min_ru = self.ru_helper.compute_utilization(self.ru_targets, self.min_ru_config)
 
         self.config_reconstruction_helper = ConfigReconstructionHelper(virtual_graph=self.mp_graph,
                                                                        original_graph=self.original_graph)
 
-    def search(self) -> List[int]:
+    def search(self) -> Dict[BaseNode, int]:
         """
         Run mixed precision search.
 
         Returns:
-            Indices of the selected bit-widths candidates.
+            Mapping from nodes to indices of the selected bit-widths candidate.
         """
-        candidates_sensitivity = self._build_sensitivity_mapping()
-        candidates_ru = self._compute_relative_ru_matrices()
-        rel_target_ru = self._get_relative_ru_constraint_per_mem_element()
-        solver = MixedPrecisionIntegerLPSolver(candidates_sensitivity, candidates_ru, rel_target_ru)
-        config = solver.run()
+        mp_config = self._prepare_and_run_solver()
 
         if self.using_virtual_graph:
-            config = self.config_reconstruction_helper.reconstruct_config_from_virtual_graph(config)
-        return config
+            mp_config = self.config_reconstruction_helper.reconstruct_config_from_virtual_graph(mp_config)
+
+        return mp_config
+
+    def _prepare_and_run_solver(self) -> Dict[BaseNode, int]:
+        """
+        Prepare sensitivity and ru data for LP solver and run the solver.
+
+        Returns:
+            Mapping from nodes to indices of the selected bit-widths candidate.
+        """
+        layers_candidates_sensitivity: Dict[BaseNode, List[float]] = self._build_sensitivity_mapping()
+        candidates_ru = self._compute_relative_ru_matrices()
+        rel_target_ru = self._get_relative_ru_constraint_per_mem_element()
+        solver = MixedPrecisionIntegerLPSolver(layers_candidates_sensitivity, candidates_ru, rel_target_ru)
+        mp_config = solver.run()
+        return mp_config
 
     def _get_relative_ru_constraint_per_mem_element(self) -> Dict[RUTarget, np.ndarray]:
         """
@@ -119,7 +130,7 @@ class MixedPrecisionSearchManager:
         """
         target_ru = self.target_resource_utilization.get_resource_utilization_dict(restricted_only=True)
         rel_target_ru = {
-            ru_target: ru - self.min_ru[ru_target] for ru_target, ru in target_ru.items()
+            ru_target: (ru - self.min_ru[ru_target]) for ru_target, ru in target_ru.items()
         }
         unsatisfiable_targets = {
             ru_target.value: target_ru[ru_target] for ru_target, ru in rel_target_ru.items() if any(ru < 0)
@@ -129,28 +140,31 @@ class MixedPrecisionSearchManager:
                              f"following targets: {unsatisfiable_targets}")
         return rel_target_ru
 
-    def _build_sensitivity_mapping(self, eps: float = EPS) -> Dict[int, Dict[int, float]]:
+    def _build_sensitivity_mapping(self, eps: float = 1e-6) -> Dict[BaseNode, List[float]]:
         """
         This function measures the sensitivity of a change in a bitwidth of a layer on the entire model.
-        It builds a mapping from a node's index, to its bitwidht's effect on the model sensitivity.
-        For each node and some possible node's bitwidth (according to the given search space), we use
-        the framework function compute_metric_fn in order to infer
-        a batch of images, and compute (using the inference results) the sensitivity metric of
-        the configured mixed-precision model.
 
         Args:
-            eps: Epsilon value to manually increase metric value (if necessary) for numerical stability
+            eps: if sensitivity for a non-max candidate is lower than for a max candidate, we set it to
+              sensitivity of a max candidate + epsilon.
 
         Returns:
-            Mapping from each node's index in a graph, to a dictionary from the bitwidth index (of this node) to
-            the sensitivity of the model.
-
+            Mapping from nodes to their bitwidth candidates sensitivity.
         """
 
         Logger.info('Starting to evaluate metrics')
-        layer_to_metrics_mapping = {}
 
-        compute_metric = self.sensitivity_evaluator.compute_metric
+        orig_sorted_nodes = self.original_graph.get_configurable_sorted_nodes(self.fw_info)
+
+        def topo_cfg(cfg: dict) -> list:
+            topo_cfg = [cfg[n] for n in orig_sorted_nodes]
+            assert len(topo_cfg) == len(cfg)
+            return topo_cfg
+
+        def compute_metric(cfg, node_idx=None, baseline_cfg=None):
+            return self.sensitivity_evaluator.compute_metric(topo_cfg(cfg),
+                                                             node_idx,
+                                                             topo_cfg(baseline_cfg) if baseline_cfg else None)
         if self.using_virtual_graph:
             origin_max_config = self.config_reconstruction_helper.reconstruct_config_from_virtual_graph(
                 self.max_ru_config)
@@ -158,19 +172,17 @@ class MixedPrecisionSearchManager:
         else:
             max_config_value = compute_metric(self.max_ru_config)
 
-        for node_idx, layer_possible_bitwidths_indices in tqdm(self.layer_to_bitwidth_mapping.items(),
-                                                               total=len(self.layer_to_bitwidth_mapping)):
-            layer_to_metrics_mapping[node_idx] = {}
-
-            for bitwidth_idx in layer_possible_bitwidths_indices:
-                if self.max_ru_config[node_idx] == bitwidth_idx:
+        layer_to_metrics_mapping = defaultdict(list)
+        for node_idx, node in tqdm(enumerate(self.mp_topo_configurable_nodes)):
+            for bitwidth_idx, _ in enumerate(node.candidates_quantization_cfg):
+                if self.max_ru_config[node] == bitwidth_idx:
                     # This is a computation of the metric for the max configuration, assign pre-calculated value
-                    layer_to_metrics_mapping[node_idx][bitwidth_idx] = max_config_value
+                    layer_to_metrics_mapping[node].append(max_config_value)
                     continue
 
                 # Create a configuration that differs at one layer only from the baseline model
                 mp_model_configuration = self.max_ru_config.copy()
-                mp_model_configuration[node_idx] = bitwidth_idx
+                mp_model_configuration[node] = bitwidth_idx
 
                 # Build a distance matrix using the function we got from the framework implementation.
                 if self.using_virtual_graph:
@@ -180,8 +192,8 @@ class MixedPrecisionSearchManager:
                             mp_model_configuration,
                             changed_virtual_nodes_idx=[node_idx],
                             original_base_config=origin_max_config)
-                    origin_changed_nodes_indices = [i for i, c in enumerate(origin_max_config) if
-                                                    c != origin_mp_model_configuration[i]]
+                    origin_changed_nodes_indices = [i for i, (n, c) in enumerate(origin_max_config.items()) if
+                                                    c != origin_mp_model_configuration[n]]
                     metric_value = compute_metric(
                         origin_mp_model_configuration,
                         origin_changed_nodes_indices,
@@ -191,11 +203,11 @@ class MixedPrecisionSearchManager:
                         mp_model_configuration,
                         [node_idx],
                         self.max_ru_config)
-
-                layer_to_metrics_mapping[node_idx][bitwidth_idx] = max(metric_value, max_config_value + eps)
+                metric_value = max(metric_value, max_config_value + eps)
+                layer_to_metrics_mapping[node].append(metric_value)
 
         # Finalize distance metric mapping
-        self.finalize_distance_metric(layer_to_metrics_mapping)
+        self._finalize_distance_metric(layer_to_metrics_mapping)
 
         return layer_to_metrics_mapping
 
@@ -221,22 +233,6 @@ class MixedPrecisionSearchManager:
 
         return graph, False
 
-    def get_search_space(self) -> Dict[int, List[int]]:
-        """
-        The search space is a mapping from a node's index to a list of integers (possible bitwidths candidates indeces
-        for the node).
-
-        Returns:
-            The entire search space of the graph.
-        """
-
-        indices_mapping = {}
-        for idx, n in enumerate(self.mp_topo_configurable_nodes):
-            # For each node, get all possible bitwidth indices for it
-            # (which is a list from 0 to the length of the candidates mp_config list of the node).
-            indices_mapping[idx] = list(range(len(n.candidates_quantization_cfg)))  # all search_methods space
-        return indices_mapping
-
     def _compute_relative_ru_matrices(self) -> Dict[RUTarget, np.ndarray]:
         """
         Computes and builds a resource utilization matrix for all restricted targets, to be used for the
@@ -248,55 +244,41 @@ class MixedPrecisionSearchManager:
             per ru target. Num memory elements depends on the target, e.g. num cuts or 1 for cumulative metrics.
         """
         rus_per_candidate = defaultdict(list)
-        for c, c_n in enumerate(self.mp_topo_configurable_nodes):
-            for candidate_idx in range(len(c_n.candidates_quantization_cfg)):
-                if candidate_idx == self.min_ru_config[c]:
+        for node in self.mp_topo_configurable_nodes:
+            for candidate_idx, _ in enumerate(node.candidates_quantization_cfg):
+                if candidate_idx == self.min_ru_config[node]:
                     candidate_rus = self.min_ru
                 else:
-                    candidate_rus = self.compute_ru_for_candidate(c, candidate_idx)
+                    cfg = self.min_ru_config.copy()
+                    cfg[node] = candidate_idx
+                    candidate_rus = self.ru_helper.compute_utilization(self.ru_targets, cfg)
 
                 for target, ru in candidate_rus.items():
                     rus_per_candidate[target].append(ru)
 
         # Each target contains a matrix of num configurations X num elements
-        relative_rus = {target: np.array(ru) - self.min_ru[target] for target, ru in rus_per_candidate.items()}
+        relative_rus = {target: (np.array(ru) - self.min_ru[target]) for target, ru in rus_per_candidate.items()}
         return relative_rus
 
-    def compute_ru_for_candidate(self, conf_node_idx: int, candidate_idx: int) -> Dict[RUTarget, np.ndarray]:
-        """
-        Computes a resource utilization vector after replacing the given node's configuration candidate in the minimal
-        target configuration with the given candidate index.
-
-        Args:
-            conf_node_idx: The index of a node in a sorted configurable nodes list.
-            candidate_idx: Quantization config candidate to be used for the node's resource utilization computation.
-
-        Returns:
-            Node's resource utilization vector.
-
-        """
-        cfg = self.replace_config_in_index(self.min_ru_config, conf_node_idx, candidate_idx)
-        return self.ru_helper.compute_utilization(self.ru_targets, cfg)
-
     @staticmethod
-    def replace_config_in_index(mp_cfg: List[int], idx: int, value: int) -> List[int]:
+    def copy_config_with_replacement(mp_cfg: Dict[BaseNode, int], node: BaseNode, candidate_idx: int) -> Dict[BaseNode, int]:
         """
-        Replacing the quantization configuration candidate in a given mixed-precision configuration at the given
-        index (node's index) with the given value (candidate index).
+        Create a copy of the given mixed-precision configuration and update the candidate index for a specific node.
 
         Args:
-            mp_cfg: Mixed-precision configuration (list of candidates' indices)
-            idx: A configurable node's index.
-            value: A new candidate index to configure.
+            mp_cfg: Mixed-precision configuration.
+            node: Node to update the config for.
+            candidate_idx: A new candidate index to configure.
 
-        Returns: A new mixed-precision configuration.
+        Returns:
+            A new mixed-precision configuration.
 
         """
         updated_cfg = mp_cfg.copy()
-        updated_cfg[idx] = value
+        updated_cfg[node] = candidate_idx
         return updated_cfg
 
-    def compute_resource_utilization_for_config(self, config: List[int]) -> ResourceUtilization:
+    def compute_resource_utilization_for_config(self, config: Dict[BaseNode, int]) -> ResourceUtilization:
         """
         Computes the resource utilization values for a given mixed-precision configuration.
 
@@ -313,7 +295,7 @@ class MixedPrecisionSearchManager:
             w_qcs=w_qcs, ru_targets=self.ru_targets, allow_unused_qcs=True)
         return ru
 
-    def finalize_distance_metric(self, layer_to_metrics_mapping: Dict[int, Dict[int, float]]):
+    def _finalize_distance_metric(self, layer_to_metrics_mapping: Dict[BaseNode, List[float]]):
         """
         Finalizing the distance metric building.
         The method checks to see if the maximal distance value is larger than a given threshold, and if so,
@@ -321,21 +303,20 @@ class MixedPrecisionSearchManager:
         Modification to the dictionary is done inplace.
 
         Args:
-            layer_to_metrics_mapping: A mapping between a node index to a mapping between
-            a bitwidth index to a distance value.
+            layer_to_metrics_mapping: A mapping between a node to a list of distance values per bitwidth candidate.
 
         """
         # normalize metric for numerical stability
+        max_dist = max(itertools.chain.from_iterable(layer_to_metrics_mapping.values()))
 
-        max_dist = max([max([d for b, d in dists.items()]) for layer, dists in layer_to_metrics_mapping.items()])
         if max_dist >= self.sensitivity_evaluator.quant_config.metric_normalization_threshold:
             Logger.warning(f"The mixed precision distance metric values indicate a large error in the quantized model."
                            f"this can cause numerical issues."
                            f"The program will proceed with mixed precision search after scaling the metric values,"
                            f"which can lead to unstable results.")
             for layer, dists in layer_to_metrics_mapping.items():
-                for b, d in dists.items():
-                    layer_to_metrics_mapping[layer][b] /= max_dist
+                for i, _ in enumerate(dists):
+                    layer_to_metrics_mapping[layer][i] /= max_dist
 
 
 class ConfigReconstructionHelper:
@@ -363,7 +344,8 @@ class ConfigReconstructionHelper:
         self.fw_info = original_graph.fw_info
 
         self.virtual_sorted_nodes_names = self.virtual_graph.get_configurable_sorted_nodes_names(self.fw_info)
-        self.origin_sorted_conf_nodes_names = self.original_graph.get_configurable_sorted_nodes_names(self.fw_info)
+        self.origin_sorted_conf_nodes = self.original_graph.get_configurable_sorted_nodes(self.fw_info)
+        self.origin_sorted_conf_nodes_names = [n.name for n in self.origin_sorted_conf_nodes]
 
         self.origin_node_idx_to_cfg = {}
 
@@ -375,9 +357,9 @@ class ConfigReconstructionHelper:
         self.origin_node_idx_to_cfg = {}
 
     def reconstruct_config_from_virtual_graph(self,
-                                              virtual_mp_cfg: List[int],
+                                              virtual_mp_cfg: Dict[BaseNode, int],
                                               changed_virtual_nodes_idx: List[int] = None,
-                                              original_base_config: List[int] = None) -> List[int]:
+                                              original_base_config: Dict[BaseNode, int] = None) -> Dict[BaseNode, int]:
         """
         Reconstructs the original config for a given virtual graph mixed-precision config.
         It iterates over all virtual configurable node (that has some chosen bit-width virtual candidate)
@@ -405,21 +387,21 @@ class ConfigReconstructionHelper:
                 [(idx, self.virtual_graph.get_configurable_sorted_nodes(self.fw_info)[idx]) for idx in changed_virtual_nodes_idx]
             # Iterating only over the virtual nodes that have updated config
             for virtual_node_idx, n in updated_virtual_nodes:
-                self.reconstruct_node_config(n, virtual_mp_cfg, virtual_node_idx)
+                self.reconstruct_node_config(n, list(virtual_mp_cfg.values()), virtual_node_idx)
             # Updating reconstructed config for all other nodes based on provided base_config
             original_sorted_conf_nodes = self.original_graph.get_configurable_sorted_nodes(self.fw_info)
-            for i in range(len(original_base_config)):
+            for i, (n, qc_ind) in enumerate(original_base_config.items()):
                 if i not in list(self.origin_node_idx_to_cfg.keys()):
-                    self.update_config_at_original_idx(n=original_sorted_conf_nodes[i],
-                                                       origin_cfg_idx=original_base_config[i])
+                    self.update_config_at_original_idx(n=n, origin_cfg_idx=qc_ind)
         else:
             # Reconstruct entire config
             for virtual_node_idx, n in enumerate(self.virtual_graph.get_configurable_sorted_nodes(self.fw_info)):
-                self.reconstruct_node_config(n, virtual_mp_cfg, virtual_node_idx)
+                self.reconstruct_node_config(n, list(virtual_mp_cfg.values()), virtual_node_idx)
 
         res_config = [self.origin_node_idx_to_cfg[key] for key in sorted(self.origin_node_idx_to_cfg.keys())]
         self._clear_reconstruction_dict()
-        return res_config
+        assert len(res_config) == len(self.origin_sorted_conf_nodes)
+        return {n: candidate_idx for n, candidate_idx in zip(self.origin_sorted_conf_nodes, res_config)}
 
     def reconstruct_node_config(self,
                                 n: BaseNode,