tico 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,78 @@
+# 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 Any, Dict, Optional
+
+import torch
+
+from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
+
+from tico.experimental.quantization.algorithm.pt2e.annotation.annotator import (
+    get_asymmetric_quantization_config,
+    PT2EAnnotator,
+)
+from tico.experimental.quantization.quantizer import BaseQuantizer
+
+
+class PT2EQuantizer(BaseQuantizer):
+    """
+    Quantizer for applying pytorch 2.0 export quantization (typically for activation quantization).
+    """
+
+    def prepare(
+        self,
+        model: torch.nn.Module,
+        args: Optional[Any] = None,
+        kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Prepare the model for pt2e quantization.
+
+        Registers activation observers using the provided example inputs.
+
+        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 pt2e quantization.
+        """
+        # Program capture
+        assert isinstance(args, tuple)
+        model = torch.export.export_for_training(
+            model, args=args, kwargs=kwargs
+        ).module()
+        quantizer = PT2EAnnotator()
+        quantizer = quantizer.set_global(get_asymmetric_quantization_config())
+
+        # Register observers in each nodes
+        assert isinstance(model, torch.fx.GraphModule)
+        model = prepare_pt2e(model, quantizer)
+
+        return model
+
+    def convert(self, model: torch.fx.GraphModule):
+        """
+        Convert the prepared model to its pt2e quantized version.
+
+        Applies the pt2e quantization on activations based on the collected statistics.
+
+        Parameters:
+            model: The prepared PyTorch model.
+
+        Returns:
+            The quantized model.
+        """
+        return convert_pt2e(model)

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

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

tico/experimental/quantization/algorithm/pt2e/transformation/convert_scalars_to_attrs.py

@@ -0,0 +1,58 @@
+# 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
+from torch.ao.quantization.fx.utils import get_new_attr_name_with_prefix
+
+
+def convert_scalars_to_attrs(model: torch.fx.GraphModule) -> torch.fx.GraphModule:
+    """
+    Convert scalar values in the graph to `get_attr` nodes.
+
+    This function identifies scalar constants in the graph and transforms them
+     into `get_attr` nodes to ensure compatibility with quantization workflows.
+    """
+    for n in model.graph.nodes:
+        if n.op != "call_function" or n.target not in [
+            # The operators that have scalar parameters.
+            torch.ops.aten.add.Tensor,
+        ]:
+            continue
+        args = list(n.args)
+        new_args = []
+        for arg in args:
+            if isinstance(arg, torch.fx.Node):
+                new_args.append(arg)
+                continue
+
+            assert isinstance(arg, float)
+            prefix = "_tensor_constant_"
+            get_new_attr_name = get_new_attr_name_with_prefix(prefix)
+            tensor_constant_name = get_new_attr_name(model)
+            float_tensor = torch.tensor(float(arg))
+            model.register_buffer(tensor_constant_name, float_tensor)
+
+            fake_mode = n.meta["val"].fake_mode
+            with model.graph.inserting_before(n):
+                get_attr_node = model.graph.create_node(
+                    "get_attr", tensor_constant_name, (), {}
+                )
+                get_attr_node.meta["val"] = fake_mode.from_tensor(
+                    float_tensor, static_shapes=True
+                )
+                new_args.append(get_attr_node)
+        n.args = tuple(new_args)
+    model.recompile()
+
+    return model

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

@@ -0,0 +1,138 @@
+# 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, List, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from torch.ao.quantization.quantizer import QuantizationSpec
+from torch.ao.quantization.quantizer.utils import _get_module_name_filter
+from torch.utils import _pytree as pytree
+
+from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
+    QuantizationConfig,
+)
+
+
+def get_module_type_filter(tp: Callable):
+    """
+    Get the module_type_filter function for a given module type.
+
+    The filter accepts a node and checks if the node comes from a module
+     that has certain module type.
+
+    For example:
+        node: linear_op = call_function[...](...)  # comes from a module with type Block -> Sub -> Linear
+
+
+    >> module_type_filter = get_module_type_filter(Sub)  # submodule with type `Sub`, under the `Block` submodule
+    >> print(module_type_filter(node))
+    True  # the node is from the submodule `Sub`
+    """
+
+    tp_str = tp.__module__ + "." + tp.__qualname__
+
+    def module_type_filter(n: torch.fx.Node) -> bool:
+        # example: {
+        #     'L__self___sub': ("L['self'].sub", <class '....Sub'>),
+        #     'L__self___sub_linear': ("L['self'].sub.linear", <class 'torch.nn.modules.linear.Linear'>)
+        # }
+        nn_module_stack = n.meta.get("nn_module_stack", {})
+        types = []
+        for _, t in nn_module_stack.values():
+            # export() returns str, but older APIs (e.g. capture_pre_autograd_graph)
+            # return type. Handle both cases.
+            if isinstance(t, type):
+                t = t.__module__ + "." + t.__qualname__
+            types.append(t)
+        return tp_str in types
+
+    return module_type_filter
+
+
+def get_not_module_type_or_name_filter(
+    tp_list: List[Callable], module_name_list: List[str]
+) -> Callable[[torch.fx.Node], bool]:
+    module_type_filters = [get_module_type_filter(tp) for tp in tp_list]
+    module_name_list_filters = [_get_module_name_filter(m) for m in module_name_list]
+
+    def not_module_type_or_name_filter(n: torch.fx.Node) -> bool:
+        return not any(f(n) for f in module_type_filters + module_name_list_filters)
+
+    return not_module_type_or_name_filter
+
+
+def get_input_act_qspec(quantization_config: Optional[QuantizationConfig]):
+    if quantization_config is None:
+        return None
+    if quantization_config.input_activation is None:
+        return None
+    quantization_spec: QuantizationSpec = quantization_config.input_activation
+    assert quantization_spec.qscheme in [
+        torch.per_tensor_affine,
+    ]
+    return quantization_spec
+
+
+def get_output_act_qspec(quantization_config: Optional[QuantizationConfig]):
+    if quantization_config is None:
+        return None
+    if quantization_config.output_activation is None:
+        return None
+    quantization_spec: QuantizationSpec = quantization_config.output_activation
+    assert quantization_spec.qscheme in [
+        torch.per_tensor_affine,
+    ]
+    return quantization_spec
+
+
+def get_weight_qspec(quantization_config: Optional[QuantizationConfig]):
+    if quantization_config is None:
+        return None
+    if quantization_config.weight is None:
+        return None
+    quantization_spec: QuantizationSpec = quantization_config.weight
+    if quantization_spec.qscheme not in [
+        torch.per_tensor_affine,
+        torch.per_channel_affine,
+    ]:
+        raise ValueError(
+            f"Unsupported quantization_spec {quantization_spec} for weight"
+        )
+    return quantization_spec
+
+
+def get_bias_qspec(quantization_config: Optional[QuantizationConfig]):
+    if quantization_config is None:
+        return None
+    if quantization_config.bias is None:
+        return None
+    quantization_spec: QuantizationSpec = quantization_config.bias
+    return quantization_spec
+
+
+def is_annotated(nodes: List[torch.fx.Node] | torch.fx.Node):
+    """
+    Check if any of the node in the given list is annotated.
+    """
+    annotated = False
+    if isinstance(nodes, torch.fx.Node):
+        nodes = [nodes]
+    for node in nodes:
+        annotated = annotated or (
+            "quantization_annotation" in node.meta
+            and node.meta["quantization_annotation"]._annotated
+        )
+    return annotated

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

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

tico/experimental/quantization/algorithm/smoothquant/observer.py

@@ -0,0 +1,78 @@
+# 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 functools
+from typing import Any, Dict, List
+
+import torch
+
+
+class ChannelwiseMaxActsObserver:
+    """
+    Observer to calcuate channelwise maximum activation
+    """
+
+    def __init__(self, model):
+        """
+        model
+            A torch module whose activations are to be analyzed.
+        hooks
+            A list to store the hooks which are registered to collect activation statistics.
+        max_acts
+            A dictionary to store the maximum activation values
+        """
+        self.model = model
+        self.hooks: List[Any] = []
+        self.max_acts: Dict[str, torch.Tensor] = {}
+
+    def attach(self):
+        """
+        Attach hooks to compute the maximum activation values per channel by running the given model
+        on a dataset.
+
+        WHAT IT DOES:
+            Set hooks to collect activation values at the per-channel level.
+            For each channel, it will calculate the maximum observed activation across
+                 all processed samples.
+        """
+        self.model.eval()
+
+        def stat_tensor(name, tensor: torch.Tensor):
+            hidden_dim = tensor.shape[-1]
+            tensor = tensor.view(-1, hidden_dim).abs().detach()
+            coming_max = torch.max(tensor, dim=0)[0]
+            if name in self.max_acts:
+                self.max_acts[name] = torch.max(self.max_acts[name], coming_max)
+            else:
+                self.max_acts[name] = coming_max
+
+        def stat_input_hook(m, input, name):
+            if isinstance(input, tuple):
+                input = input[0]
+            stat_tensor(name, input)
+
+        for name, m in self.model.named_modules():
+            if isinstance(m, torch.nn.Linear):
+                self.hooks.append(
+                    m.register_forward_pre_hook(
+                        functools.partial(stat_input_hook, name=name)
+                    )
+                )
+
+    def remove(self):
+        for hook in self.hooks:
+            hook.remove()
+
+    def get_max_acts(self):
+        return self.max_acts

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

@@ -0,0 +1,81 @@
+# 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 Any, Dict, Optional
+
+import torch
+
+from tico.experimental.quantization.algorithm.smoothquant.observer import (
+    ChannelwiseMaxActsObserver,
+)
+
+from tico.experimental.quantization.algorithm.smoothquant.smooth_quant import (
+    apply_smoothing,
+)
+from tico.experimental.quantization.config import SmoothQuantConfig
+from tico.experimental.quantization.quantizer import BaseQuantizer
+
+
+class SmoothQuantQuantizer(BaseQuantizer):
+    """
+    Quantizer for applying the SmoothQuant algorithm
+    """
+
+    def __init__(self, config: SmoothQuantConfig):
+        super().__init__(config)
+
+        self.alpha = config.alpha
+        self.custom_alpha_map = config.custom_alpha_map
+        self.observer: Optional[ChannelwiseMaxActsObserver] = None
+
+    @torch.no_grad()
+    def prepare(
+        self,
+        model: torch.nn.Module,
+        args: Optional[Any] = None,
+        kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        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 SmoothQuant quantization.
+        """
+        self.observer = ChannelwiseMaxActsObserver(model)
+        self.observer.attach()
+
+        return model
+
+    @torch.no_grad()
+    def convert(self, model):
+        """
+        Convert the prepared model to its SmoothQuant quantized version.
+        Applies the SmoothQuant quantization on weights based on the collected statistics.
+
+        Parameters:
+            model: The prepared PyTorch model.
+
+        Returns:
+            The quantized model.
+        """
+        if self.observer is not None:
+            self.observer.remove()
+            apply_smoothing(
+                model, self.observer.get_max_acts(), self.alpha, self.custom_alpha_map
+            )
+
+        return model

tico/experimental/quantization/algorithm/smoothquant/smooth_quant.py

@@ -0,0 +1,164 @@
+# 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 Any, Dict, List, Optional
+
+import torch
+
+
+@torch.no_grad()
+def smooth_weights(
+    front_module: torch.nn.Module,
+    back_modules: torch.nn.Module | List[torch.nn.Module],
+    activation_max: torch.Tensor,
+    alpha: float,
+):
+    """
+    Applies SmoothQuant-style smoothing to the weights and biases of two
+     connected modules using activation maximum values.
+
+    NOTE All modules **MUST** have `weight` and optionally `bias` attributes.
+
+    Parameters
+    -----------
+        front_module
+            The front module whose weights and biases will be adjusted.
+        back_modules
+            A list of back modules whose weights and biases will be adjusted.
+        activation_max
+            A tensor of channel-wise maximum activation values for the front module.
+        alpha
+            The smoothing factor that determines the scaling for weight adjustments.
+
+    Raises
+    -------
+    AttributeError
+        If `front_module` or any module in `back_modules` does not have `weight` attributes.
+    ValueError
+        If the shape of tensors in `activation_max` does not match the number of channels
+         in `front_module`'s weight.
+    NoteImplementedError
+        If `front_module` or any module in `back_modules` is of an unsupported type.
+    """
+    from transformers.models.llama.modeling_llama import LlamaRMSNorm
+
+    if not isinstance(back_modules, list):
+        back_modules = [back_modules]
+
+    # Check attributes
+    if not hasattr(front_module, "weight"):
+        raise AttributeError(
+            f"The front module '{type(front_module).__name__}' does not have a 'weight' attribute."
+        )
+    for back_m in back_modules:
+        if not hasattr(back_m, "weight"):
+            raise AttributeError(
+                f"The front module '{type(back_m).__name__}' does not have a 'weight' attribute."
+            )
+    # Check shapes
+    if isinstance(front_module, LlamaRMSNorm):
+        front_numel = front_module.weight.numel()
+    else:
+        raise NotImplementedError(
+            f"Unsupported module type: {type(front_module).__name__}"
+        )
+    for back_m in back_modules:
+        if isinstance(back_m, torch.nn.Linear):
+            back_numel = back_m.in_features
+        else:
+            raise NotImplementedError(
+                f"Unsupported module type: {type(front_module).__name__}"
+            )
+
+        if front_numel != back_numel or back_numel != activation_max.numel():
+            raise ValueError(
+                f"Shape mismatch: front_numel({front_numel}), back_numel({back_numel}), activation_max_numel({activation_max.numel()})"
+            )
+
+    # Compute scales
+    device, dtype = back_modules[0].weight.device, back_modules[0].weight.dtype
+    activation_max = activation_max.to(device=device, dtype=dtype)  # type: ignore[arg-type]
+    weight_scales = torch.cat(
+        [back_m.weight.abs().max(dim=0, keepdim=True)[0] for back_m in back_modules],  # type: ignore[operator]
+        dim=0,
+    )
+    weight_scales = weight_scales.max(dim=0)[0].clamp(min=1e-5)
+    scales = (
+        (activation_max.pow(alpha) / weight_scales.pow(1 - alpha))
+        .clamp(min=1e-5)
+        .to(device)  # type: ignore[arg-type]
+        .to(dtype)  # type: ignore[arg-type]
+    )
+
+    # Smooth
+    front_module.weight.div_(scales)
+    if hasattr(front_module, "bias"):
+        front_module.bias.div_(scales)
+
+    for back_m in back_modules:
+        back_m.weight.mul_(scales.view(1, -1))  # type: ignore[operator]
+
+
+@torch.no_grad()
+def apply_smoothing(
+    model: torch.nn.Module,
+    activation_max: Dict[str, torch.Tensor],
+    alpha: float = 0.5,
+    custom_alpha_map: Optional[Dict[str, float]] = None,
+):
+    """
+    Applies SmoothQuant-style smoothing to the model's weights using activation maximum values.
+
+    Parameters
+    -----------
+        model
+            A torch module whose weights will be smoothed.
+        activation_max
+            The channel-wise maximum activation values for the model.
+        alpha
+            The default smoothing factor to apply across all modules.
+        custom_alpha_map
+            A dictionary mapping layer/module names to custom alpha values.
+            Layers specified in this dictionary will use the corresponding alpha
+             value instead of the default.
+    """
+    from transformers.models.llama.modeling_llama import LlamaDecoderLayer
+
+    for name, module in model.named_modules():
+        alpha_to_apply = alpha
+        if custom_alpha_map and name in custom_alpha_map:
+            alpha_to_apply = custom_alpha_map[name]
+        if alpha_to_apply > 1.0:
+            raise RuntimeError(
+                f"Alpha value cannot exceed 1.0. Given alpha: {alpha_to_apply}"
+            )
+        # SmoothQuant is applied before capturing the graph. Therefore, it needs to know
+        #  specific module information.
+        # TODO Suport more modules.
+        if isinstance(module, LlamaDecoderLayer):
+            attn_ln = module.input_layernorm
+            qkv = [
+                module.self_attn.q_proj,
+                module.self_attn.k_proj,
+                module.self_attn.v_proj,
+            ]
+
+            qkv_input_scales = activation_max[name + ".self_attn.q_proj"]
+            smooth_weights(attn_ln, qkv, qkv_input_scales, alpha_to_apply)
+
+            ffn_ln = module.post_attention_layernorm
+            fcs = [module.mlp.gate_proj, module.mlp.up_proj]
+            fcs_input_scales = activation_max[name + ".mlp.gate_proj"]
+
+            smooth_weights(ffn_ln, fcs, fcs_input_scales, alpha_to_apply)

tico/experimental/quantization/config.py

@@ -0,0 +1,68 @@
+# 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 abc import ABC, abstractmethod
+from typing import Dict, Optional
+
+
+class BaseConfig(ABC):
+    """
+    Base configuration class for quantization.
+    """
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        pass
+
+
+class PT2EConfig(BaseConfig):
+    """
+    Configuration for pytorch 2.0 export quantization.
+    """
+
+    @property
+    def name(self) -> str:
+        return "pt2e"
+
+
+class GPTQConfig(BaseConfig):
+    """
+    Configuration for GPTQ.
+    """
+
+    def __init__(self, verbose: bool = False):
+        self.verbose = verbose
+
+    @property
+    def name(self) -> str:
+        return "gptq"
+
+
+class SmoothQuantConfig(BaseConfig):
+    """
+    Configuration for smooth quant.
+    """
+
+    def __init__(
+        self,
+        alpha: float = 0.5,
+        custom_alpha_map: Optional[Dict[str, float]] = None,
+    ):
+        self.alpha = alpha
+        self.custom_alpha_map = custom_alpha_map
+
+    @property
+    def name(self) -> str:
+        return "smooth_quant"

tico/experimental/quantization/evaluation/__init__.py

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