tico 0.1.0__py3-none-any.whl

tico/experimental/quantization/passes/propagate_qparam_backward.py

@@ -0,0 +1,91 @@
+# 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 TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import copy
+
+import torch
+from torch.export import ExportedProgram
+
+from tico.serialize.quant_param import QPARAM_KEY
+from tico.utils import logging
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.validate_args_kwargs import CatArgs, PermuteArgs, ReshapeArgs
+
+
+@trace_graph_diff_on_pass
+class PropagateQParamBackward(PassBase):
+    """
+    This pass propagates quantization parameters backward.
+
+    BEFORE)
+
+    node -> reshape (with meta[QPARAM_KEY])
+
+    AFTER)
+
+    node (with meta[QPARAM_KEY]) -> reshape (with meta[QPARAM_KEY])
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        graph_module = exported_program.graph_module
+        graph: torch.fx.Graph = graph_module.graph
+
+        def _propagate_qparam_if_possible(src: torch.fx.Node, dst: torch.fx.Node):
+            if QPARAM_KEY not in src.meta:
+                return
+
+            if (
+                QPARAM_KEY in dst.meta
+                and src.meta[QPARAM_KEY].dtype != dst.meta[QPARAM_KEY].dtype
+            ):
+                return
+
+            dst.meta[QPARAM_KEY] = copy.deepcopy(src.meta[QPARAM_KEY])
+
+            logger.debug(f"{src.name}'s quantparam is propagated to {dst.name}.")
+
+        # Do reverse-order traversal for backward propagation
+        for node in reversed(graph.nodes):
+            if node.op != "call_function":
+                continue
+            if node.target == torch.ops.aten.cat.default:
+                concat_args = CatArgs(*node.args, **node.kwargs)
+                concat_inputs = concat_args.tensors
+
+                for concat_input in concat_inputs:
+                    _propagate_qparam_if_possible(node, concat_input)
+            elif node.target == torch.ops.aten.reshape.default:
+                args = ReshapeArgs(*node.args, **node.kwargs)
+                _propagate_qparam_if_possible(node, args.input)
+            elif node.target == torch.ops.aten.permute.default:
+                permute_args = PermuteArgs(*node.args, **node.kwargs)
+                _propagate_qparam_if_possible(node, permute_args.input)
+            # TODO Support more ops.
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        # Run only once.
+        return PassResult(False)

tico/experimental/quantization/passes/propagate_qparam_forward.py

@@ -0,0 +1,141 @@
+# 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 TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import copy
+
+import torch
+from torch.export import ExportedProgram
+
+from tico.serialize.quant_param import QPARAM_KEY
+from tico.utils import logging
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.validate_args_kwargs import (
+    CatArgs,
+    NegArgs,
+    PermuteArgs,
+    ReshapeArgs,
+    SliceArgs,
+)
+
+
+@trace_graph_diff_on_pass
+class PropagateQParamForward(PassBase):
+    """
+    A pass propagates quantization parameters through operations that do not alter them.
+
+    This pass identifies and propagates quantization parameters through operations that
+     do not change their values, such as `permute`, `reshape`, `transpose`, `view` and
+    similar tensor transformations.
+
+    By ensuring that quantization parameters remain consistent across such operations,
+    this pass helps maintain correctness in quantization-aware representations.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        def _propagate_qparam_if_possible(src: torch.fx.Node, dst: torch.fx.Node):
+            if QPARAM_KEY not in src.meta:
+                return
+
+            if (
+                QPARAM_KEY in dst.meta
+                and src.meta[QPARAM_KEY].dtype != dst.meta[QPARAM_KEY].dtype
+            ):
+                return
+
+            dst.meta[QPARAM_KEY] = copy.deepcopy(src.meta[QPARAM_KEY])
+
+            logger.debug(f"{src.name}'s quantparam is propagated to {dst.name}.")
+
+        graph_module = exported_program.graph_module
+        graph: torch.fx.Graph = graph_module.graph
+        for node in graph.nodes:
+            if node.op != "call_function":
+                continue
+            if node.target == torch.ops.aten.permute.default:
+                permute_args = PermuteArgs(*node.args, **node.kwargs)
+                _propagate_qparam_if_possible(permute_args.input, node)
+            elif node.target == torch.ops.aten.reshape.default:
+                reshape_args = ReshapeArgs(*node.args, **node.kwargs)
+                _propagate_qparam_if_possible(reshape_args.input, node)
+            elif node.target == torch.ops.aten.slice.Tensor:
+                slice_args = SliceArgs(*node.args, **node.kwargs)
+                _propagate_qparam_if_possible(slice_args.input, node)
+            elif node.target == torch.ops.aten.neg.default:
+                neg_args = NegArgs(*node.args, **node.kwargs)
+
+                if QPARAM_KEY not in neg_args.input.meta:
+                    continue
+                # Only support int16 for now
+                if neg_args.input.meta[QPARAM_KEY].dtype != "int16":
+                    continue
+
+                _propagate_qparam_if_possible(neg_args.input, node)
+
+            elif node.target == torch.ops.aten.cat.default:
+                concat_args = CatArgs(*node.args, **node.kwargs)
+                concat_inputs = concat_args.tensors
+
+                cond = True
+                for concat_input in concat_inputs:
+                    # Check all inputs have qparam
+                    if QPARAM_KEY not in concat_input.meta:
+                        cond = False
+                        break
+
+                    # Only support int16 for now
+                    if concat_input.meta[QPARAM_KEY].dtype != "int16":
+                        cond = False
+                        break
+
+                    if concat_input.meta[QPARAM_KEY].scale is None:
+                        cond = False
+                        break
+
+                    if len(concat_input.meta[QPARAM_KEY].scale) != 1:
+                        cond = False
+                        break
+
+                if not cond:
+                    continue
+
+                # Find max scale node
+                max_scale = 0.0
+                max_scale_node = None
+                for concat_input in concat_inputs:
+                    scale = concat_input.meta[QPARAM_KEY].scale[0]
+                    if max_scale < scale:
+                        max_scale = scale
+                        max_scale_node = concat_input
+
+                assert max_scale_node is not None
+                _propagate_qparam_if_possible(max_scale_node, node)
+
+            # TODO Support more ops.
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        # Run only once.
+        return PassResult(False)

tico/experimental/quantization/passes/quantize_bias.py

@@ -0,0 +1,123 @@
+# 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 TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import copy
+
+import torch
+from torch.export import ExportedProgram
+
+from tico.serialize.quant_param import QPARAM_KEY, QuantParam, to_qparam_dtype
+from tico.utils import logging
+from tico.utils.graph import add_placeholder, get_torch_param_value, is_torch_param
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.validate_args_kwargs import LinearArgs
+
+
+@trace_graph_diff_on_pass
+class QuantizeBias(PassBase):
+    """
+    Quantize bias.
+
+    This pass identifies fp32 biases, quantizes them using scales of input and weights.
+
+    This pass assumes that if bias is fp32, input and weights must have been quantized.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        graph_module = exported_program.graph_module
+        graph: torch.fx.Graph = graph_module.graph
+        for node in graph.nodes:
+            if node.op != "call_function":
+                continue
+            if node.target == torch.ops.aten.linear.default:
+                lin_args = LinearArgs(*node.args, **node.kwargs)
+                inp = lin_args.input
+                weights = lin_args.weight
+                bias = lin_args.bias
+
+                if bias is None:
+                    continue
+
+                # Only support bias is Parameter
+                # TODO Is it possible that bias is not Parameter?
+                if not is_torch_param(bias, exported_program):
+                    continue
+
+                bias_val: torch.Tensor = get_torch_param_value(bias, exported_program)
+                if bias_val.dtype != torch.float32:
+                    continue
+
+                if QPARAM_KEY not in inp.meta:
+                    continue
+
+                if QPARAM_KEY not in weights.meta:
+                    continue
+
+                quant_dtype = None
+                if inp.meta[QPARAM_KEY].dtype == "int16":
+                    quant_dtype = torch.int64
+                elif inp.meta[QPARAM_KEY].dtype == "uint8":
+                    quant_dtype = torch.int32
+                else:
+                    continue
+
+                type_info = torch.iinfo(quant_dtype)
+
+                assert quant_dtype is not None
+
+                i_scale = inp.meta[QPARAM_KEY].scale
+                w_scale = weights.meta[QPARAM_KEY].scale
+
+                assert i_scale is not None
+                assert w_scale is not None
+                assert len(i_scale) == 1
+                assert len(w_scale) == bias_val.shape[0]
+
+                bias_scale = torch.tensor(i_scale) * torch.tensor(w_scale)
+                q_bias = torch.round(bias_val / bias_scale)
+                q_bias = torch.clamp(q_bias, min=type_info.min, max=type_info.max)
+                q_bias = q_bias.to(quant_dtype)
+
+                q_bias_node = add_placeholder(exported_program, q_bias, bias.name)
+
+                qparam = QuantParam()
+                qparam.scale = bias_scale.tolist()
+                assert qparam.scale is not None
+                qparam.zero_point = [0] * len(qparam.scale)
+                qparam.dtype = to_qparam_dtype(quant_dtype)
+                qparam.quantized_dimension = 0
+                q_bias_node.meta[QPARAM_KEY] = qparam
+
+                node.update_arg(2, q_bias_node)
+
+                logger.debug(f"Bias ({bias.name}) is quantized to {q_bias_node.name}.")
+
+            # TODO Support more ops.
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        # Run only once.
+        return PassResult(False)

tico/experimental/quantization/passes/remove_weight_dequant_op.py

@@ -0,0 +1,177 @@
+# 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 List, Optional, TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from torch._export.utils import (
+    get_buffer,
+    get_lifted_tensor_constant,
+    is_buffer,
+    is_lifted_tensor_constant,
+)
+from torch._subclasses.fake_tensor import FakeTensor
+from torch.export import ExportedProgram
+
+from tico.serialize.quant_param import QPARAM_KEY, QuantParam, to_qparam_dtype
+from tico.utils import logging
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.validate_args_kwargs import (
+    DequantizePerChannelArgs,
+    DequantizePerTensorArgs,
+)
+
+
+def get_constant(exported_program: ExportedProgram, node: torch.fx.Node):
+    assert isinstance(node, torch.fx.Node)
+    if node.name in exported_program.constants:
+        return exported_program.constants[node.name]
+    elif is_buffer(exported_program, node):
+        return get_buffer(exported_program, node)
+    elif is_lifted_tensor_constant(exported_program, node):
+        return get_lifted_tensor_constant(exported_program, node)
+    else:
+        raise RuntimeError("NYI constant")
+
+
+class ValRange:
+    def __init__(self, val: Union[torch.Tensor, List[int]]):
+        if isinstance(val, torch.Tensor):
+            self.max = torch.max(val).item()
+            self.min = torch.min(val).item()
+        elif type(val) == list:
+            self.max = max(val)
+            self.min = min(val)
+        else:
+            raise RuntimeError("Wrong dtype (val)")
+
+    def within(self, min_val, max_val):
+        return self.min >= min_val and self.max <= max_val
+
+
+# Infer dtype using weight, zero point, and dtype
+def infer_dtype(weight: torch.Tensor, zerop: List[int], dtype: torch.dtype) -> str:
+    weight_val = ValRange(weight)
+    zp_val = ValRange(zerop)
+
+    if weight_val.within(0, 15) and zp_val.within(0, 15) and dtype == torch.uint8:
+        return "uint4"
+    else:
+        return to_qparam_dtype(dtype)
+
+
+@trace_graph_diff_on_pass
+class RemoveWeightDequantOp(PassBase):
+    """
+    This pass identifies and removes any remaining Dequantize ops associated with
+     quantized weights.
+
+    Since weights already quantized earlier (and possibly kept in float by
+     attaching a DQ), the final stage of the quantization pipeline typically
+    does not require those DQ ops anymore.
+
+    NOTE Removing 'DQ' causes a sementic change: f32 -> quantized
+
+    [BEFORE]
+      W (quantized) - Dequantize (float)
+
+    [AFTER]
+      W (quantized)
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        graph_module = exported_program.graph_module
+        graph: torch.fx.Graph = graph_module.graph
+        for dq in graph.nodes:
+            if not dq.op == "call_function":
+                continue
+
+            if dq.target not in [
+                torch.ops.quantized_decomposed.dequantize_per_channel.default,
+                torch.ops.quantized_decomposed.dequantize_per_tensor.default,
+            ]:
+                continue
+            dq_args: Optional[DequantizePerChannelArgs | DequantizePerTensorArgs] = None
+
+            if (
+                dq.target
+                == torch.ops.quantized_decomposed.dequantize_per_channel.default
+            ):
+                dq_args = DequantizePerChannelArgs(*dq.args, **dq.kwargs)
+            elif (
+                dq.target
+                == torch.ops.quantized_decomposed.dequantize_per_tensor.default
+            ):
+                dq_args = DequantizePerTensorArgs(*dq.args, **dq.kwargs)
+            else:
+                raise RuntimeError(f"Invalid DQ target: {dq.target}")
+
+            q_weight = dq_args.input
+            # All weights are placehoders.
+            if q_weight.op != "placeholder":
+                continue
+            # Check if DQ already has quant param because DQ can be shared.
+            if QPARAM_KEY in q_weight.meta:
+                continue
+
+            q_weight_meta = q_weight.meta["val"]
+            assert isinstance(q_weight_meta, FakeTensor)
+            # Weight should have quantized values.
+            assert q_weight_meta.dtype != torch.float
+
+            q_weight_val = get_constant(exported_program, q_weight)
+            assert isinstance(q_weight_val, torch.Tensor)
+
+            quant_param = QuantParam()
+            if isinstance(dq_args, DequantizePerChannelArgs):
+                scales = get_constant(exported_program, dq_args.scales)
+                zero_ps = get_constant(exported_program, dq_args.zero_points)
+
+                # Sometimes users can give fp32 zero point. Let's update dtype here.
+                zero_ps = zero_ps.to(torch.int64)
+                quant_param.scale = scales.tolist()
+                quant_param.zero_point = zero_ps.tolist()
+                assert quant_param.zero_point is not None  # To avoid mypy error
+                quant_param.quantized_dimension = dq_args.axis
+                quant_param.dtype = infer_dtype(
+                    q_weight_val, quant_param.zero_point, q_weight_meta.dtype
+                )
+            elif isinstance(dq_args, DequantizePerTensorArgs):
+                quant_param.scale = [dq_args.scale]
+                quant_param.zero_point = [dq_args.zero_point]
+                assert quant_param.zero_point is not None  # To avoid mypy error
+                quant_param.dtype = infer_dtype(
+                    q_weight_val, quant_param.zero_point, q_weight_meta.dtype
+                )
+            else:
+                raise RuntimeError(f"Invalid DQ target: {dq.target}")
+
+            q_weight.meta[QPARAM_KEY] = quant_param
+            dq.replace_all_uses_with(q_weight, propagate_meta=False)
+            logger.debug(f"{dq.name} is removed.")
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        # Run only once.
+        return PassResult(False)

tico/experimental/quantization/public_interface.py

@@ -0,0 +1,108 @@
+# 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 copy
+from typing import Any, Dict, Optional, Type
+
+import torch
+
+from tico.experimental.quantization.algorithm.gptq.quantizer import GPTQQuantizer
+from tico.experimental.quantization.algorithm.pt2e.quantizer import PT2EQuantizer
+from tico.experimental.quantization.algorithm.smoothquant.quantizer import (
+    SmoothQuantQuantizer,
+)
+from tico.experimental.quantization.config import BaseConfig
+from tico.experimental.quantization.quantizer import BaseQuantizer
+
+
+config_to_quantizer: Dict[str, Type[BaseQuantizer]] = {
+    "pt2e": PT2EQuantizer,
+    "gptq": GPTQQuantizer,
+    "smooth_quant": SmoothQuantQuantizer,
+}
+
+QUANTIZER_ATTRIBUTE_NAME = "tico_quantizer"
+
+
+def prepare(
+    model: torch.nn.Module,
+    quant_config: BaseConfig,
+    args: Optional[Any] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    inplace: Optional[bool] = False,
+):
+    """
+    Prepare the model for quantization using the provided configuration.
+
+    Determines the appropriate quantizer based on the type of `quant_config` and
+    prepares the model accordingly.
+
+    Parameters:
+        model: The PyTorch model to be quantized.
+        quant_config (BaseConfig): The quantization configuration.
+        args (Any, optional): Positional example inputs required for activation quantization.
+        kwargs (Dict[str, Any], optional): Keyword example inputs required for activation quantization.
+        inplace (bool, optional): If true, the model will be modified in place;
+                                   otherwise, a new prepared model is returned.
+
+    Returns:
+        The model prepared for quantization.
+    """
+    if quant_config.name == "pt2e" and inplace:
+        raise RuntimeError(
+            "In-place is not supported for PT2E quantization due to limitation in the underlying Torch APIs. Please set 'inplace=False' to proceed."
+        )
+
+    model = model if inplace else copy.deepcopy(model)
+
+    quantizer = config_to_quantizer[quant_config.name](quant_config)
+    model = quantizer.prepare(model, args, kwargs)
+    setattr(model, QUANTIZER_ATTRIBUTE_NAME, quantizer)
+
+    return model
+
+
+def convert(model, inplace: Optional[bool] = False):
+    """
+    Convert the prepared model to a quantized model using the provided configuration.
+
+    Determines the appropriate quantizer based on the type of quant_config and
+    converts the model accordingly.
+
+    Parameters:
+        model: The prepared PyTorch model.
+        inplace (bool, optional): If true, the model will be modified in place;
+                                   otherwise, a new prepared model is returned.
+
+    Returns:
+        The quantized model.
+    """
+    # Get quantizer first before calling deepcopy that does not copy attributes properly.
+    if hasattr(model, QUANTIZER_ATTRIBUTE_NAME):
+        quantizer = getattr(model, QUANTIZER_ATTRIBUTE_NAME)
+        delattr(model, QUANTIZER_ATTRIBUTE_NAME)
+    else:
+        raise RuntimeError("Call prepare() function first.")
+
+    if isinstance(quantizer, PT2EQuantizer) and inplace:
+        raise RuntimeError(
+            "In-place is not supported for PT2E quantization due to limitation in the underlying Torch APIs. Please set 'inplace=False' to proceed."
+        )
+
+    model = model if inplace else copy.deepcopy(model)
+
+    assert isinstance(quantizer, BaseQuantizer)
+    model = quantizer.convert(model)
+
+    return model