mct-nightly 2.3.0.20250506.618__tar.gz → 2.3.0.20250507.555__tar.gz

{mct_nightly-2.3.0.20250506.618 → mct_nightly-2.3.0.20250507.555}/PKG-INFO

@@ -1,7 +1,8 @@
 Metadata-Version: 2.4
 Name: mct-nightly
-Version: 2.3.0.20250506.618
+Version: 2.3.0.20250507.555
 Summary: A Model Compression Toolkit for neural networks
+Author-email: ssi-dnn-dev@sony.com
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
@@ -23,6 +24,7 @@ Requires-Dist: protobuf
 Requires-Dist: mct-quantizers-nightly
 Requires-Dist: pydantic>=2.0
 Requires-Dist: sony-custom-layers-dev==0.4.0.dev6
+Dynamic: author-email
 Dynamic: classifier
 Dynamic: description
 Dynamic: description-content-type
@@ -51,7 +53,7 @@ ______________________________________________________________________
 </p>
 <p align="center">
   <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/pytorch-2.2%20%7C%202.3%20%7C%202.4%20%7C%202.5-blue" /></a>
-  <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/tensorflow-02.14%20%7C%202.15-blue" /></a>
+  <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/tensorflow-2.14%20%7C%202.15-blue" /></a>
   <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/python-3.9%20%7C%203.10%20%7C%203.11%20%7C%203.12-blue" /></a>
   <a href="https://github.com/sony/model_optimization/releases"><img src="https://img.shields.io/github/v/release/sony/model_optimization" /></a>
   <a href="https://github.com/sony/model_optimization/blob/main/LICENSE.md"><img src="https://img.shields.io/badge/license-Apache%202.0-blue" /></a>

{mct_nightly-2.3.0.20250506.618 → mct_nightly-2.3.0.20250507.555}/README.md

@@ -18,7 +18,7 @@ ______________________________________________________________________
 </p>
 <p align="center">
   <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/pytorch-2.2%20%7C%202.3%20%7C%202.4%20%7C%202.5-blue" /></a>
-  <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/tensorflow-02.14%20%7C%202.15-blue" /></a>
+  <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/tensorflow-2.14%20%7C%202.15-blue" /></a>
   <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/python-3.9%20%7C%203.10%20%7C%203.11%20%7C%203.12-blue" /></a>
   <a href="https://github.com/sony/model_optimization/releases"><img src="https://img.shields.io/github/v/release/sony/model_optimization" /></a>
   <a href="https://github.com/sony/model_optimization/blob/main/LICENSE.md"><img src="https://img.shields.io/badge/license-Apache%202.0-blue" /></a>

{mct_nightly-2.3.0.20250506.618 → mct_nightly-2.3.0.20250507.555}/mct_nightly.egg-info/PKG-INFO

@@ -1,7 +1,8 @@
 Metadata-Version: 2.4
 Name: mct-nightly
-Version: 2.3.0.20250506.618
+Version: 2.3.0.20250507.555
 Summary: A Model Compression Toolkit for neural networks
+Author-email: ssi-dnn-dev@sony.com
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
@@ -23,6 +24,7 @@ Requires-Dist: protobuf
 Requires-Dist: mct-quantizers-nightly
 Requires-Dist: pydantic>=2.0
 Requires-Dist: sony-custom-layers-dev==0.4.0.dev6
+Dynamic: author-email
 Dynamic: classifier
 Dynamic: description
 Dynamic: description-content-type
@@ -51,7 +53,7 @@ ______________________________________________________________________
 </p>
 <p align="center">
   <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/pytorch-2.2%20%7C%202.3%20%7C%202.4%20%7C%202.5-blue" /></a>
-  <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/tensorflow-02.14%20%7C%202.15-blue" /></a>
+  <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/tensorflow-2.14%20%7C%202.15-blue" /></a>
   <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/python-3.9%20%7C%203.10%20%7C%203.11%20%7C%203.12-blue" /></a>
   <a href="https://github.com/sony/model_optimization/releases"><img src="https://img.shields.io/github/v/release/sony/model_optimization" /></a>
   <a href="https://github.com/sony/model_optimization/blob/main/LICENSE.md"><img src="https://img.shields.io/badge/license-Apache%202.0-blue" /></a>

{mct_nightly-2.3.0.20250506.618 → mct_nightly-2.3.0.20250507.555}/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.20250506.000618"
+__version__ = "2.3.0.20250507.000555"

{mct_nightly-2.3.0.20250506.618 → mct_nightly-2.3.0.20250507.555}/model_compression_toolkit/core/common/graph/virtual_activation_weights_node.py

@@ -12,22 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import abc
 import uuid
 
-from typing import Dict, Any, Tuple
-
 from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.constants import VIRTUAL_ACTIVATION_WEIGHTS_NODE_PREFIX, \
     VIRTUAL_WEIGHTS_SUFFIX, VIRTUAL_ACTIVATION_SUFFIX, FLOAT_BITWIDTH
-
+from model_compression_toolkit.core.common.framework_info import DEFAULT_KERNEL_ATTRIBUTES
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
-import numpy as np
-
 from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import \
     CandidateNodeQuantizationConfig
+from model_compression_toolkit.core.common.quantization.node_quantization_config import ActivationQuantizationMode
+
 
+class VirtualNode(BaseNode, abc.ABC):
+    """ Base class for all virtual nodes. """
+    pass
 
-class VirtualSplitNode(BaseNode):
+
+class VirtualSplitNode(VirtualNode, abc.ABC):
     """
     A class that represents a node that was split from a kernel node (node with weights).
     """
@@ -73,14 +76,11 @@ class VirtualSplitWeightsNode(VirtualSplitNode):
         super().__init__(origin_node)
 
         self.name = origin_node.name + VIRTUAL_WEIGHTS_SUFFIX
-        # Virtual weights node is created only to be absorbed into virtual composed node right away.
-        # However, in some cases composition is impossible and virtual weights node can remain in the graph.
-        # In such case it messes up resource utilization computation, specifically activation cuts. In order to minimize
-        # the impact, we preserve the behavior of the original node wrt activation (shape and quantization),
-        # so that prev - virtualW cut is identical to prev-origin_node. Only the cut virtualW-virtualA will be different
-        # from the original graph, so in the worst case the utilization will be higher in virtual graph.
-        # This should guarantee that the utilization of the original graph does not exceed the requested target.
-        self.candidates_quantization_cfg = origin_node.candidates_quantization_cfg
+
+        self.candidates_quantization_cfg = origin_node.get_unique_weights_candidates(kernel_attr)
+        for c in self.candidates_quantization_cfg:
+            c.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.NO_QUANT
+            c.activation_quantization_cfg.activation_n_bits = FLOAT_BITWIDTH
 
 
 class VirtualSplitActivationNode(VirtualSplitNode):
@@ -113,7 +113,7 @@ class VirtualSplitActivationNode(VirtualSplitNode):
             c.weights_quantization_cfg.weights_n_bits = FLOAT_BITWIDTH
 
 
-class VirtualActivationWeightsNode(BaseNode):
+class VirtualActivationWeightsNode(VirtualNode):
     """
     A node that represents a composition of pair of sequential activation node and weights (kernel) node.
     This structure is used for mixed-precision search with bit-operation constraint.
@@ -149,7 +149,7 @@ class VirtualActivationWeightsNode(BaseNode):
         weights = weights_node.weights.copy()
         act_node_w_rename = {}
         if act_node.weights:
-            if not fw_info.get_kernel_op_attributes(act_node)[0] is None:
+            if fw_info.get_kernel_op_attributes(act_node) != DEFAULT_KERNEL_ATTRIBUTES:
                 raise NotImplementedError(f'Node {act_node} with kernel cannot be used as activation for '
                                           f'VirtualActivationWeightsNode.')
             if act_node.has_any_configurable_weight():

mct_nightly-2.3.0.20250507.555/model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py

@@ -0,0 +1,505 @@
+# Copyright 2021 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import itertools
+
+import copy
+from collections import defaultdict
+
+from tqdm import tqdm
+
+from typing import Dict, List, Tuple, Optional
+
+import numpy as np
+
+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
+from model_compression_toolkit.core.common.graph.base_graph import Graph
+from model_compression_toolkit.core.common.graph.virtual_activation_weights_node import VirtualActivationWeightsNode, \
+    VirtualSplitWeightsNode, VirtualSplitActivationNode, VirtualNode
+from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import \
+    RUTarget, ResourceUtilization
+from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization_calculator import \
+    TargetInclusionCriterion, BitwidthMode
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_ru_helper import \
+    MixedPrecisionRUHelper
+from model_compression_toolkit.core.common.mixed_precision.search_methods.linear_programming import \
+    MixedPrecisionIntegerLPSolver
+from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import SensitivityEvaluation
+from model_compression_toolkit.core.common.substitutions.apply_substitutions import substitute
+from model_compression_toolkit.logger import Logger
+
+
+class MixedPrecisionSearchManager:
+    """
+    Class to wrap and manage the search process of a mixed-precision configuration.
+    """
+
+    def __init__(self,
+                 graph: Graph,
+                 fw_info: FrameworkInfo,
+                 fw_impl: FrameworkImplementation,
+                 sensitivity_evaluator: SensitivityEvaluation,
+                 target_resource_utilization: ResourceUtilization):
+        """
+
+        Args:
+            graph: Graph to search for its MP configuration.
+            fw_info: FrameworkInfo object about the specific framework (e.g., attributes of different layers' weights to quantize).
+            fw_impl: FrameworkImplementation object with specific framework methods implementation.
+            sensitivity_evaluator: A SensitivityEvaluation which provides a function that evaluates the sensitivity of
+                a bit-width configuration for the MP model.
+            target_resource_utilization: Target Resource Utilization to bound our feasible solution space s.t the configuration does not violate it.
+        """
+
+        self.fw_info = fw_info
+        self.fw_impl = fw_impl
+
+        self.original_graph = graph
+        # graph for mp search
+        self.mp_graph, self.using_virtual_graph = self._get_mp_graph(graph, target_resource_utilization)
+        del graph  # so that it's not used by mistake
+
+        self.sensitivity_evaluator = sensitivity_evaluator
+        self.target_resource_utilization = target_resource_utilization
+
+        self.mp_topo_configurable_nodes = self.mp_graph.get_configurable_sorted_nodes(fw_info)
+
+        self.ru_targets = target_resource_utilization.get_restricted_targets()
+        self.ru_helper = MixedPrecisionRUHelper(self.original_graph, fw_info, fw_impl)
+
+        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.config_reconstruction_helper = ConfigReconstructionHelper(self.original_graph)
+        if self.using_virtual_graph:
+            real_min_ru_config = self.config_reconstruction_helper.reconstruct_full_configuration(self.min_ru_config)
+            self.min_ru = self.ru_helper.compute_utilization(self.ru_targets, real_min_ru_config)
+        else:
+            self.min_ru = self.ru_helper.compute_utilization(self.ru_targets, self.min_ru_config)
+
+    def search(self) -> Dict[BaseNode, int]:
+        """
+        Run mixed precision search.
+
+        Returns:
+            Mapping from nodes to indices of the selected bit-widths candidate.
+        """
+        mp_config = self._prepare_and_run_solver()
+
+        if self.using_virtual_graph:
+            mp_config = self.config_reconstruction_helper.reconstruct_full_configuration(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.
+        """
+        candidates_ru = self._compute_relative_ru_matrices()
+        rel_target_ru = self._get_relative_ru_constraint_per_mem_element()
+        layers_candidates_sensitivity: Dict[BaseNode, List[float]] = self._build_sensitivity_mapping()
+        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]:
+        """
+        Computes resource utilization constraint with respect to the minimal bit configuration, i.e. corresponding
+        constraint for each memory element is the relative utilization between the target utilization and
+        element's utilization for min-bit configuration.
+
+        Returns:
+            A dictionary of relative resource utilization constraints per ru target.
+
+        Raises:
+            ValueError: if target resource utilization cannot be satisfied (utilization for the minimal bit
+              configuration exceeds the requested target utilization for any target).
+        """
+        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()
+        }
+        unsatisfiable_targets = {
+            ru_target.value: target_ru[ru_target] for ru_target, ru in rel_target_ru.items() if any(ru < 0)
+        }
+        if unsatisfiable_targets:
+            raise ValueError(f"The model cannot be quantized to meet the specified resource utilization for the "
+                             f"following targets: {unsatisfiable_targets}")
+        return rel_target_ru
+
+    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.
+
+        Args:
+            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 nodes to their bitwidth candidates sensitivity.
+        """
+
+        Logger.info('Starting to evaluate metrics')
+
+        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_full_configuration(self.max_ru_config)
+            max_config_value = compute_metric(origin_max_config)
+        else:
+            max_config_value = compute_metric(self.max_ru_config)
+
+        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].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] = bitwidth_idx
+
+                # Build a distance matrix using the function we got from the framework implementation.
+                if self.using_virtual_graph:
+                    # Reconstructing original graph's configuration from virtual graph's configuration
+                    orig_mp_config = self.config_reconstruction_helper.reconstruct_full_configuration(mp_model_configuration)
+                    changed_nodes = [orig_sorted_nodes.index(n) for n, ind in orig_mp_config.items()
+                                     if origin_max_config[n] != ind]
+                    metric_value = compute_metric(orig_mp_config, changed_nodes, origin_max_config)
+                else:
+                    metric_value = compute_metric(mp_model_configuration, [node_idx], self.max_ru_config)
+                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)
+
+        return layer_to_metrics_mapping
+
+    def _get_mp_graph(self, graph: Graph, target_resource_utilization: ResourceUtilization) -> Tuple[Graph, bool]:
+        """
+        Get graph for mixed precision search. Virtual graph is built if bops is restricted and both activation and
+        weights are configurable.
+
+        Args:
+            graph: input graph.
+            target_resource_utilization: target resource utilization.
+
+        Returns:
+            Graph for mixed precision search (virtual or original), and a boolean flag whether a virtual graph has been
+            constructed.
+        """
+        if (target_resource_utilization.bops_restricted() and
+                graph.has_any_configurable_activation() and
+                graph.has_any_configurable_weights()):
+            mp_graph = substitute(copy.deepcopy(graph),
+                                  self.fw_impl.get_substitutions_virtual_weights_activation_coupling())
+            return mp_graph, True
+
+        return graph, False
+
+    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
+        mixed-precision search problem formalization.
+        Utilization is computed relative to the minimal configuration, i.e. utilization for it will be 0.
+
+        Returns:
+            A dictionary containing resource utilization matrix of shape (num configurations, num memory elements)
+            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 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:
+                    cfg = self.min_ru_config.copy()
+                    cfg[node] = candidate_idx
+                    real_cfg = self.config_reconstruction_helper.reconstruct_full_configuration(cfg)
+                    candidate_rus = self.ru_helper.compute_utilization(self.ru_targets, real_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()}
+        return relative_rus
+
+    @staticmethod
+    def copy_config_with_replacement(mp_cfg: Dict[BaseNode, int], node: BaseNode, candidate_idx: int) -> Dict[BaseNode, int]:
+        """
+        Create a copy of the given mixed-precision configuration and update the candidate index for a specific node.
+
+        Args:
+            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.
+
+        """
+        updated_cfg = mp_cfg.copy()
+        updated_cfg[node] = candidate_idx
+        return updated_cfg
+
+    def compute_resource_utilization_for_config(self, config: Dict[BaseNode, int]) -> ResourceUtilization:
+        """
+        Computes the resource utilization values for a given mixed-precision configuration.
+
+        Args:
+            config: A mixed-precision configuration (list of candidates indices)
+
+        Returns: A ResourceUtilization object with the model's resource utilization values when quantized 
+        with the given config.
+
+        """
+        act_qcs, w_qcs = self.ru_helper.get_quantization_candidates(config)
+        ru = self.ru_helper.ru_calculator.compute_resource_utilization(
+            target_criterion=TargetInclusionCriterion.AnyQuantized, bitwidth_mode=BitwidthMode.QCustom, act_qcs=act_qcs,
+            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[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,
+        it scales all metric values to prevent possible numerical issues.
+        Modification to the dictionary is done inplace.
+
+        Args:
+            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()))
+
+        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 i, _ in enumerate(dists):
+                    layer_to_metrics_mapping[layer][i] /= max_dist
+
+
+class ConfigReconstructionHelper:
+    def __init__(self, original_graph):
+        # mapping in order to return the actual node objects from the original graph
+        self.orig_nodes = {n.name: n for n in original_graph.nodes}
+
+    def reconstruct_full_configuration(self,
+                                       virtual_cfg: Dict[BaseNode, int],
+                                       include_non_configurable: bool = False) -> Dict[BaseNode, int]:
+        """
+        Convert a configuration of a virtual graph into the corresponding configuration of the original graph.
+        Note that a configurable VirtualActivationWeightsNode might comprise one configurable and one non-configurable
+        original nodes.
+
+        Args:
+            virtual_cfg: a mapping from nodes in the virtual graph to selected candidate index. Should contain all
+                configurable nodes of the virtual graph, and only configurable nodes.
+            include_non_configurable: whether to return configs for non-configurable original nodes.
+
+        Returns:
+            A mapping from configurable nodes in the original graph to their candidate indices.
+        """
+        # Original candidate of a node that has been split might be determined by two different virtual nodes, one
+        # determines activation and one - weights. First, for each virtual node we collect the original
+        # activation / weights nodes, with all original candidates that match the virtual candidate
+        # activation / weights config. If both activation and weights of the original node are determined by virtual
+        # candidates, we look for a common candidate.
+        orig_nodes_a_candidates = {}
+        orig_nodes_w_candidates = {}
+        for virtual_node, virtual_qc_ind in virtual_cfg.items():
+            assert virtual_node.has_configurable_activation() or virtual_node.has_any_configurable_weight()
+            orig_a_node, orig_a_candidates = self._retrieve_matching_orig_a_candidates(virtual_node, virtual_qc_ind)
+            if orig_a_node and (include_non_configurable or orig_a_node.has_configurable_activation()):
+                assert orig_a_node not in orig_nodes_a_candidates
+                orig_nodes_a_candidates[orig_a_node] = orig_a_candidates
+            orig_w_node, orig_w_candidates = self._retrieve_matching_orig_w_candidates(virtual_node, virtual_qc_ind)
+            if orig_w_node and (include_non_configurable or orig_w_node.has_any_configurable_weight()):
+                assert orig_w_node not in orig_nodes_w_candidates
+                orig_nodes_w_candidates[orig_w_node] = orig_w_candidates
+
+        orig_cfg = {}
+        common_orig_nodes = set(orig_nodes_a_candidates.keys()).intersection(set(orig_nodes_w_candidates))
+        for orig_node in common_orig_nodes:
+            a_candidates = orig_nodes_a_candidates[orig_node]
+            w_candidates = orig_nodes_w_candidates[orig_node]
+            # find the common candidate
+            common_candidates = set(a_candidates).intersection(set(w_candidates))
+            if len(common_candidates) != 1:
+                raise ValueError(f'Expected to find exactly one candidate with the required activation and weights '
+                                 f'quantization configuration for node {orig_node}. Found {len(common_candidates)}')
+            # in theory it's possible that original non-configurable node gets split and each part is combined
+            # with a configurable part of another node and we end up here
+            if orig_node.has_configurable_activation() or orig_node.has_any_configurable_weight():
+                orig_cfg[orig_node] = common_candidates.pop()
+            del orig_nodes_a_candidates[orig_node]
+            del orig_nodes_w_candidates[orig_node]
+
+        # remaining a nodes
+        for orig_node, a_candidates in orig_nodes_a_candidates.items():
+            assert not orig_node.has_any_configurable_weight()  # if it had we should have caught it above
+            assert len(a_candidates) == 1
+            assert orig_node not in orig_cfg
+            if include_non_configurable or orig_node.has_configurable_activation():
+                orig_cfg[orig_node] = a_candidates[0]
+
+        # remaining w nodes
+        for orig_node, w_candidates in orig_nodes_w_candidates.items():
+            assert not orig_node.has_configurable_activation()  # if it had we should have caught it above
+            assert len(w_candidates) == 1
+            assert orig_node not in orig_cfg
+            if include_non_configurable or orig_node.has_any_configurable_weight():
+                orig_cfg[orig_node] = w_candidates[0]
+
+        return orig_cfg
+
+    def reconstruct_separate_aw_configs(self, virtual_cfg: Dict[BaseNode, int], include_non_configurable: bool) \
+            -> Tuple[Dict[BaseNode, int], Dict[BaseNode, int]]:
+        """
+        Retrieves original activation and weights nodes and corresponding candidates for a given configuration of the
+        virtual graph. Only returns configuration specified by the virtual config, per configurable target (activation
+        or weights). For example, if 'virtual_cfg' contains a single VirtualActivationWeightsNode, the returned
+        configuration will contain only activation config for the original activation node, and only weights config
+        for the original weights node).
+        In practice, we return candidate index in both cases, instead of actual activation or weights config, since
+        sensitivity evaluator heavily depends on it, so we must ignore activation config in weights candidate and vice
+        versa. This is bad!!! TODO
+
+        Args:
+            virtual_cfg: a mapping from nodes in the virtual graph to selected candidate index.
+            include_non_configurable: whether to return configs for non-configurable target (i.e. activation config
+              for non-configurable activation, and weights config for non-configurable weight).
+
+        Returns:
+            Configuration for original activation nodes and a separate configuration for original weights nodes.
+        """
+        a_cfg = {}
+        w_cfg = {}
+        for virtual_node, virtual_qc_ind in virtual_cfg.items():
+            orig_a_node, orig_a_candidates = self._retrieve_matching_orig_a_candidates(virtual_node, virtual_qc_ind)
+            if orig_a_node and (include_non_configurable or orig_a_node.has_configurable_activation()):
+                # we may have retrieved multiple candidates with different weights configs and identical activation
+                # configs, so we just take the first
+                a_cfg[orig_a_node] = orig_a_candidates[0]
+
+            orig_w_node, orig_w_candidates = self._retrieve_matching_orig_w_candidates(virtual_node, virtual_qc_ind)
+            if orig_w_node and (include_non_configurable or orig_w_node.has_any_configurable_weight()):
+                # we may have retrieved multiple candidates with different activation configs and identical weights
+                # configs, so we just take the first
+                w_cfg[orig_w_node] = orig_w_candidates[0]
+
+        return a_cfg, w_cfg
+
+    def _retrieve_matching_orig_a_candidates(self,
+                                             virtual_node: BaseNode,
+                                             virtual_qc_ind: int) -> Tuple[Optional[BaseNode], Optional[List[int]]]:
+        """
+        Retrieve the original activation node and all its candidates matching activation quantization config of the
+        given virtual candidate (candidate of a node in the virtual graph).
+        Note that we do simple matching, without any filtering, so disabled activation quantization will be also matched.
+
+        Args:
+            virtual_node: node in the virtual graph (can be virtual or regular).
+            virtual_qc_ind: candidate index of the virtual node.
+
+        Returns:
+            The original activation node (actual object from the original graph) and a list of its matching candidates.
+        """
+        if not isinstance(virtual_node, VirtualNode):
+            return self.orig_nodes[virtual_node.name], [virtual_qc_ind]
+        if isinstance(virtual_node, VirtualSplitWeightsNode):
+            return None, None
+        if isinstance(virtual_node, VirtualActivationWeightsNode):
+            orig_a_node = virtual_node.original_activation_node
+            if isinstance(orig_a_node, VirtualSplitActivationNode):
+                orig_a_node = orig_a_node.origin_node
+        else:
+            assert isinstance(virtual_node, VirtualSplitActivationNode)
+            orig_a_node = virtual_node.origin_node
+
+        virtual_qc = virtual_node.candidates_quantization_cfg[virtual_qc_ind]
+        matching_orig_a_cfgs = [i for i, orig_qc in enumerate(orig_a_node.candidates_quantization_cfg)
+                                if orig_qc.activation_quantization_cfg == virtual_qc.activation_quantization_cfg]
+        if not matching_orig_a_cfgs:    # pragma: no cover
+            raise ValueError(f'Could not find matching activation quantization config in the original node '
+                             f'{orig_a_node} for candidate {virtual_qc_ind} of the virtual node {virtual_node}')
+        return self.orig_nodes[orig_a_node.name], matching_orig_a_cfgs
+
+    def _retrieve_matching_orig_w_candidates(self,
+                                             virtual_node: BaseNode,
+                                             virtual_qc_ind: int) -> Tuple[Optional[BaseNode], Optional[List[int]]]:
+        """
+        Retrieve the original weights node and all its candidates matching weights quantization config of the
+        given virtual candidate (candidate of a node in the virtual graph).
+
+        Args:
+            virtual_node: node in the virtual graph (can be virtual or regular).
+            virtual_qc_ind: candidate index of the virtual node.
+
+        Returns:
+            The original weights node (actual object from the original graph) and a list of all its matching candidates.
+        """
+        if not isinstance(virtual_node, VirtualNode):
+            if virtual_node.weights:
+                return self.orig_nodes[virtual_node.name], [virtual_qc_ind]
+            return None, None
+        if isinstance(virtual_node, VirtualSplitActivationNode):
+            return None, None
+
+        if isinstance(virtual_node, VirtualActivationWeightsNode):
+            assert isinstance(virtual_node.original_weights_node, VirtualSplitWeightsNode)
+            orig_w_node = virtual_node.original_weights_node.origin_node
+        else:
+            assert isinstance(virtual_node, VirtualSplitWeightsNode)
+            orig_w_node = virtual_node.origin_node
+
+        virtual_qc = virtual_node.candidates_quantization_cfg[virtual_qc_ind]
+
+        # Matching candidate is a candidate with matching configs for configurable weights. We cannot compare the entire
+        # weights config since the virtual node may contain additional non-configurable weights from the activation node
+        orig_configurable_attrs = [attr for attr in orig_w_node.weights if virtual_node.is_configurable_weight(attr)]
+        assert all(virtual_node.is_configurable_weight(attr) for attr in orig_configurable_attrs)
+
+        def get_configurable_attrs_cfgs(qc):
+            return {attr: qc.weights_quantization_cfg.get_attr_config(attr) for attr in orig_configurable_attrs}
+        virtual_cfg = get_configurable_attrs_cfgs(virtual_qc)
+        matching_orig_w_cfgs = [i for i, orig_qc in enumerate(orig_w_node.candidates_quantization_cfg)
+                                if get_configurable_attrs_cfgs(orig_qc) == virtual_cfg]
+        if not matching_orig_w_cfgs:    # pragma: no cover
+            raise ValueError(f'Could not find matching weights quantization config in the original node '
+                             f'{orig_w_node} for candidate {virtual_qc_ind} of the virtual node {virtual_node}')
+        return self.orig_nodes[orig_w_node.name], matching_orig_w_cfgs