tico 0.1.0.dev250411__py3-none-any.whl

tico/experimental/quantization/algorithm/gptq/quantizer.py

@@ -0,0 +1,225 @@
+# Copyright (c) 2024 Intel Corporation
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. 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 types
+from typing import Any, Dict, List, Optional
+
+import torch
+
+from tico.experimental.quantization.algorithm.gptq.gptq import GPTQ
+from tico.experimental.quantization.algorithm.gptq.utils import (
+    find_layers,
+    gather_single_batch_from_dict,
+    gather_single_batch_from_list,
+)
+from tico.experimental.quantization.config import BaseConfig, GPTQConfig
+from tico.experimental.quantization.quantizer import BaseQuantizer
+
+
+class GPTQQuantizer(BaseQuantizer):
+    """
+    Quantizer for applying the GPTQ algorithm (typically for weight quantization)
+    """
+
+    def __init__(self, config: BaseConfig):
+        super().__init__(config)
+
+        self.cache_args: List[Any] = []
+        self.cache_kwargs: Dict[str, Any] = {"batch_num": 0}
+
+    @torch.no_grad()
+    def prepare(
+        self,
+        model: torch.nn.Module,
+        args: Optional[Any] = None,
+        kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Overrides the forward method of the first LLaMA layer (layer 0) to capture the
+         input required for calibration.
+
+        This method modifies the original forward pass of LLaMA layer 0 so that the
+         inputs used during inference are intercepted and recorded. These captured inputs
+        are then utilized to calibrate the quantization parameters for the GPTQ.
+
+        Parameters:
+            model: The target PyTorch model.
+            args: Positional example inputs required for capturing graph.
+            kwargs: Keyword example inputs required for capturing graph.
+
+        Returns:
+            The model prepared for GPTQ quantization.
+        """
+        if args is None and kwargs is None:
+            raise RuntimeError(
+                "Either args or kwargs must be provided for captruing graph."
+            )
+        # Define a function to capture input activations and associated parameters.
+        def forward(layer, *args, **kwargs):
+            self.cache_kwargs["batch_num"] += 1
+            for idx, item in enumerate(args):
+                if (idx + 1) > len(self.cache_args):
+                    self.cache_args.append([])
+                self.cache_args[idx].append(item)
+            for arg in kwargs:
+                if self.cache_kwargs.get(arg, None) is None:
+                    self.cache_kwargs[arg] = []
+                self.cache_kwargs[arg].append(kwargs[arg])
+            # Raise an error to interrupt the forward pass after capturing data.
+            raise ValueError
+
+        # Replace the first layer with defined function to capture calibration data.
+        if hasattr(model, "model"):
+            assert hasattr(model.model, "layers")
+            assert isinstance(model.model.layers, torch.nn.ModuleList)
+            layer_forward_cache = model.model.layers[0].forward
+            model.model.layers[0].forward = types.MethodType(
+                forward, model.model.layers[0]
+            )
+        else:
+            assert hasattr(model, "forward")
+            layer_forward_cache = model.forward
+            model.forward = types.MethodType(forward, model.forward)
+
+        model_forward_cache = model.forward
+        # Replace model's forward to avoid ValueError
+        def model_forward(model, *args, **kwargs):
+            nonlocal model_forward_cache
+            try:
+                model_forward_cache(*args, **kwargs)
+            except ValueError:
+                pass
+
+        model.forward = types.MethodType(model_forward, model)
+        kwargs = kwargs or {}
+        model(*args, **kwargs)  # type: ignore[misc]
+
+        # Recover original forward
+        model.forward = model_forward_cache
+        if hasattr(model, "model"):
+            assert hasattr(model.model, "layers")
+            assert isinstance(model.model.layers, torch.nn.ModuleList)
+            model.model.layers[0].forward = layer_forward_cache
+        else:
+            model.forward = layer_forward_cache
+
+        return model
+
+    @torch.no_grad()
+    def convert(self, model):
+        """
+        Convert the prepared model to its GPTQ quantized version.
+
+        Applies the GPTQ quantization on weights based on the collected statistics.
+
+        Parameters:
+            model: The prepared PyTorch model.
+
+        Returns:
+            The quantized model.
+        """
+        gptq_conf = self.config
+        assert isinstance(gptq_conf, GPTQConfig)
+
+        # Save the original cache setting and disable caching during calibration/inference.
+        if hasattr(model, "config"):
+            use_cache = model.config.use_cache
+            model.config.use_cache = False
+
+        quantizers = {}
+        if hasattr(model, "model"):
+            target_layers = model.model.layers
+        else:
+            target_layers = [model]
+        for l_idx, layer in enumerate(target_layers):
+            # Identify quantizable submodules within the layer.
+            full = find_layers(layer)
+
+            sequential = [list(full.keys())]
+            for names in sequential:
+                subset = {n: full[n] for n in names}
+
+                gptq: Dict[str, GPTQ] = {}
+                for name in subset:
+                    gptq[name] = GPTQ(subset[name])
+                    gptq[name].quantizer.configure(
+                        8, perchannel=True, sym=False, mse=False
+                    )
+                # Define a hook to collect input/output batches for quantizer calibration.
+                def add_batch(name):
+                    def tmp(_, inp, out):
+                        gptq[name].add_batch(inp[0].data, out.data)
+
+                    return tmp
+
+                handles = []
+                for name in subset:
+                    handles.append(subset[name].register_forward_hook(add_batch(name)))
+                # Run the current layer on the stored calibration inputs to capture activation stats.
+                batch_num = self.cache_kwargs.pop("batch_num")
+                for batch_idx in range(batch_num):
+                    cache_args_batch = gather_single_batch_from_list(
+                        self.cache_args, batch_idx
+                    )
+                    cache_kwargs_batch = gather_single_batch_from_dict(
+                        self.cache_kwargs, batch_idx
+                    )
+                    layer(*cache_args_batch, **cache_kwargs_batch)[0]
+                self.cache_kwargs["batch_num"] = batch_num
+                for h in handles:
+                    h.remove()
+                # Quantize each submodule using the collected calibration data.
+                for name in subset:
+                    if gptq_conf.verbose:
+                        print(l_idx, name)
+                        print("Quantizing ...")
+                    gptq[name].fasterquant(
+                        percdamp=0.01,
+                        groupsize=-1,
+                        actorder=True,
+                        static_groups=False,
+                        verbose=gptq_conf.verbose,
+                    )
+                    quantizers["model.layers.%d.%s" % (l_idx, name)] = gptq[
+                        name
+                    ].quantizer
+                    gptq[name].free()
+            """
+            Execute the quantized layer with the calibration inputs to obtain ouptuts
+             that will serve as inputs for the next layer.
+
+            This ensures that the quantization effects are correctly propagated to subsequent
+             layers.
+            """
+            batch_num = self.cache_kwargs.pop("batch_num")
+            for batch_idx in range(batch_num):
+                cache_args_batch = gather_single_batch_from_list(
+                    self.cache_args, batch_idx
+                )
+                cache_kwargs_batch = gather_single_batch_from_dict(
+                    self.cache_kwargs, batch_idx
+                )
+                outs = layer(*cache_args_batch, **cache_kwargs_batch)[0]
+                # Update inputs for next iteration.
+                self.cache_args[0][batch_idx] = outs
+            self.cache_kwargs["batch_num"] = batch_num
+
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        # Restore the original cache configuration.
+        if hasattr(model, "config"):
+            model.config.use_cache = use_cache
+
+        return model

tico/experimental/quantization/algorithm/gptq/utils.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2024 Intel Corporation
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. 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 torch
+
+
+def find_layers(module, layers=[torch.nn.Linear], name=""):
+    if type(module) in layers:
+        return {name: module}
+    res = {}
+    for name1, child in module.named_children():
+        res.update(
+            find_layers(
+                child, layers=layers, name=name + "." + name1 if name != "" else name1
+            )
+        )
+    return res
+
+
+def gather_single_batch_from_dict(data_dict, idx):
+    """
+    Gather single batch from a dict.
+
+    Args:
+        data_dict (dict): data dict.
+        idx (int): index
+
+    Returns:
+        dict: single batch.
+    """
+    # obtain a set of keyword input from cache
+    single_batch = {}
+    for k, v in data_dict.items():
+        single_batch[k] = data_dict[k][idx]
+    return single_batch
+
+
+def gather_single_batch_from_list(data_list, idx):
+    """
+    Gather single batch from a list.
+
+    Args:
+        data_dict (dict): data list.
+        idx (int): index
+
+    Returns:
+        list: single batch.
+    """
+    # obtain a set of keyword input from cache
+    single_batch = []
+    for data_item in data_list:
+        single_batch.append(data_item[idx])
+    return single_batch

tico/experimental/quantization/algorithm/pt2e/__init__.py

@@ -0,0 +1 @@
+# DO NOT REMOVE THIS FILE

tico/experimental/quantization/algorithm/pt2e/annotation/__init__.py

@@ -0,0 +1 @@
+# DO NOT REMOVE THIS FILE

tico/experimental/quantization/algorithm/pt2e/annotation/annotator.py

@@ -0,0 +1,215 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. 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.
+
+from __future__ import annotations
+
+import functools
+from typing import Any, Callable, Dict, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+    from torch.ao.quantization.qconfig import _ObserverOrFakeQuantizeConstructor
+import torch
+from torch.ao.quantization.observer import (
+    MinMaxObserver,
+    MovingAverageMinMaxObserver,
+    MovingAveragePerChannelMinMaxObserver,
+    PerChannelMinMaxObserver,
+)
+from torch.ao.quantization.quantizer import QuantizationSpec, Quantizer
+from torch.ao.quantization.quantizer.utils import _get_module_name_filter
+
+from tico.experimental.quantization.algorithm.pt2e.annotation.op import *
+import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+from tico.experimental.quantization.algorithm.pt2e.transformation.convert_scalars_to_attrs import (
+    convert_scalars_to_attrs,
+)
+
+
+class PT2EAnnotator(Quantizer):
+    """
+    The class annotates quantization configurations on each nodes.
+
+    Observers would be attached according to those configurations in
+     'torch.prepare_pt2e'.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.global_config: Optional[QuantizationConfig] = None
+        self.operator_type_config: Dict[
+            torch._ops.OpOverloadPacket, QuantizationConfig
+        ] = {}
+        self.module_type_config: Dict[Callable, QuantizationConfig] = {}
+        self.module_name_config: Dict[str, QuantizationConfig] = {}
+
+    def set_global(self, quantization_config: QuantizationConfig) -> PT2EAnnotator:
+        """
+        Set quantization config globally.
+        """
+        assert quantization_config is not None
+        self.global_config = quantization_config
+        return self
+
+    def set_operator_type(
+        self,
+        operator_type: torch._ops.OpOverloadPacket,
+        quantization_config: QuantizationConfig,
+    ) -> PT2EAnnotator:
+        """
+        Set quantization config for given operator type.
+        """
+        assert quantization_config is not None
+        self.operator_type_config[operator_type] = quantization_config
+        return self
+
+    def set_module_type(
+        self, module_type: Callable, quantization_config: QuantizationConfig
+    ):
+        """
+        Set quantization config for given module type.
+
+        For example, let's say quantizer.set_module_type(nn.Linear).
+        It will quantize all 'nn.Linear' modules with the `quantization_config`.
+        """
+        assert quantization_config is not None
+        self.module_type_config[module_type] = quantization_config
+        return self
+
+    def set_module_name(
+        self, module_name: str, quantization_config: QuantizationConfig
+    ):
+        """
+        Set quantization config for given module name.
+
+        For example, let's say quantizer.set_module_name("blocks.sub").
+        It will quantize all nodes that come from a module whose name is "blocks.sub"
+         with the `quantization_config`.
+        """
+        assert quantization_config is not None
+        self.module_name_config[module_name] = quantization_config
+        return self
+
+    def transform_for_annotation(
+        self, model: torch.fx.GraphModule
+    ) -> torch.fx.GraphModule:
+        """Allows for user defined transforms to run before annotating the graph."""
+        model = convert_scalars_to_attrs(model)
+        return model
+
+    def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
+        model = self._annotate_for_quantization(model)
+        annot_utils.propagate_annotation_forward(model)
+        return model
+
+    def validate(self, model: torch.fx.GraphModule) -> None:
+        # TODO Consider this method.
+        pass
+
+    def _annotate_by_config_and_filter(
+        self,
+        model: torch.fx.GraphModule,
+        quantization_config: Optional[QuantizationConfig],
+        filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+    ) -> torch.fx.GraphModule:
+        assert quantization_config is not None
+
+        for node in model.graph.nodes:
+            if node.target not in annot_spec.OP_TO_ANNOTATOR:
+                continue
+            annot_spec.OP_TO_ANNOTATOR[node.target](
+                model, node, quantization_config, filter_fn
+            )
+        return model
+
+    def _annotate_for_quantization(
+        self, model: torch.fx.GraphModule
+    ) -> torch.fx.GraphModule:
+        # Annotate according to the given module names.
+        module_name_list = list(self.module_name_config.keys())
+        for module_name, config in self.module_name_config.items():
+            self._annotate_by_config_and_filter(
+                model, config, _get_module_name_filter(module_name)
+            )
+
+        # Annotate according to the given module types.
+        tp_list = list(self.module_type_config.keys())
+        for module_type, config in self.module_type_config.items():
+            self._annotate_by_config_and_filter(
+                model, config, quant_utils.get_module_type_filter(module_type)
+            )
+
+        # TODO Annotate according to the given operator types.
+
+        self._annotate_by_config_and_filter(
+            model,
+            self.global_config,
+            quant_utils.get_not_module_type_or_name_filter(tp_list, module_name_list),
+        )
+        return model
+
+
+@functools.lru_cache
+def get_asymmetric_quantization_config(
+    weight_is_per_channel: bool = True,
+    act_qmin: int = 0,
+    act_qmax: int = 255,
+    weight_qmin: int = 0,
+    weight_qmax: int = 255,
+) -> QuantizationConfig:
+    # activation
+    act_extra_args: Dict[str, Any] = {"eps": 2**-12}
+    act_observer = MinMaxObserver
+    act_qspec = QuantizationSpec(
+        dtype=torch.uint8,
+        quant_min=act_qmin,
+        quant_max=act_qmax,
+        qscheme=torch.per_tensor_affine,
+        is_dynamic=False,
+        observer_or_fake_quant_ctr=act_observer.with_args(
+            **act_extra_args,
+        ),
+    )
+    # weight
+    weight_extra_args: Dict[str, Any] = {"eps": 2**-12}
+    weight_qscheme = (
+        torch.per_channel_affine if weight_is_per_channel else torch.per_tensor_affine
+    )
+    weight_observer: _ObserverOrFakeQuantizeConstructor = (
+        PerChannelMinMaxObserver if weight_is_per_channel else MinMaxObserver
+    )
+    weight_qspec = QuantizationSpec(
+        dtype=torch.uint8,
+        quant_min=weight_qmin,
+        quant_max=weight_qmax,
+        qscheme=weight_qscheme,
+        ch_axis=0,
+        is_dynamic=False,
+        observer_or_fake_quant_ctr=weight_observer.with_args(**weight_extra_args),
+    )
+
+    # Set bias qspec in each annotation functions.
+    bias_qspec = None
+    quantization_config = QuantizationConfig(
+        act_qspec,
+        act_qspec,
+        weight_qspec,
+        bias_qspec,
+    )
+    return quantization_config

tico/experimental/quantization/algorithm/pt2e/annotation/config.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. 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.
+
+from dataclasses import dataclass
+from typing import Optional
+
+from torch.ao.quantization.quantizer import QuantizationSpec
+
+
+@dataclass(eq=True, frozen=True)
+class QuantizationConfig:
+    input_activation: Optional[QuantizationSpec]
+    output_activation: Optional[QuantizationSpec]
+    weight: Optional[QuantizationSpec]
+    bias: Optional[QuantizationSpec]

tico/experimental/quantization/algorithm/pt2e/annotation/op/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. 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 glob
+from os.path import basename, dirname, isfile, join
+
+modules = glob.glob(join(dirname(__file__), "*.py"))
+__all__ = [
+    basename(f)[:-3] for f in modules if isfile(f) and not f.endswith("__init__.py")
+]

tico/experimental/quantization/algorithm/pt2e/annotation/op/adaptive_avg_pool2d.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. 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.
+
+from typing import Callable, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from torch.ao.quantization.quantizer import SharedQuantizationSpec
+
+import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+from tico.utils.validate_args_kwargs import AdaptiveAvgPool2dArgs
+
+
+@annot_spec.register_annotator([torch.ops.aten.adaptive_avg_pool2d.default])
+def _annotate_adaptive_avg_pool2d(
+    gm: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+):
+    if (
+        node.op != "call_function"
+        or node.target != torch.ops.aten.adaptive_avg_pool2d.default
+    ):
+        return
+    if filter_fn and not filter_fn(node):
+        return
+    if quant_utils.is_annotated(node):
+        return
+
+    args = AdaptiveAvgPool2dArgs(*node.args)  # type: ignore[arg-type]
+    input = args.input
+
+    assert isinstance(input, torch.fx.Node)
+    if (
+        "quantization_annotation" not in input.meta
+        or not input.meta["quantization_annotation"]._annotated
+        or input.meta["quantization_annotation"].output_qspec is None
+    ):
+        input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
+    else:
+        input_act_qspec = SharedQuantizationSpec(input)
+    annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
+
+    output_act_qspec = SharedQuantizationSpec((input, node))
+    annot_utils.annotate_output_qspec(node, output_act_qspec)
+
+    annot_utils.mark_nodes_as_annotated(node)

tico/experimental/quantization/algorithm/pt2e/annotation/op/add.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. 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.
+
+from typing import Callable, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+
+import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+from tico.utils.validate_args_kwargs import AddTensorArgs
+
+
+@annot_spec.register_annotator([torch.ops.aten.add.Tensor])
+def _annotate_add(
+    gm: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
+):
+    if node.op != "call_function" or node.target != torch.ops.aten.add.Tensor:
+        return
+    if filter_fn and not filter_fn(node):
+        return
+    if quant_utils.is_annotated(node):
+        return
+
+    args = AddTensorArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+    input_ = args.input
+    other = args.other
+
+    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
+    if isinstance(input_, torch.fx.Node):
+        annot_utils.annotate_input_qspec_map(node, input_, input_act_qspec)
+    if isinstance(other, torch.fx.Node):
+        annot_utils.annotate_input_qspec_map(node, other, input_act_qspec)
+
+    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
+    annot_utils.annotate_output_qspec(node, output_act_qspec)
+
+    annot_utils.mark_nodes_as_annotated(node)