tico 0.1.0.dev250714__py3-none-any.whl → 0.1.0.dev251102__py3-none-any.whl

tico/quantization/wrapq/observers/affine_base.py

@@ -0,0 +1,128 @@
+# 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 math
+from typing import Optional, Tuple
+
+import torch
+
+from tico.quantization.wrapq.dtypes import DType, UINT8
+from tico.quantization.wrapq.observers.base import ObserverBase
+from tico.quantization.wrapq.qscheme import QScheme
+
+
+class AffineObserverBase(ObserverBase):
+    """Base for affine observers (min/max → scale/zp)."""
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        dtype: DType = UINT8,
+        qscheme: QScheme = QScheme.PER_TENSOR_ASYMM,
+        channel_axis: Optional[int] = None,
+    ):
+        super().__init__(
+            name=name, dtype=dtype, qscheme=qscheme, channel_axis=channel_axis
+        )
+
+    def reset(self) -> None:
+        """
+        Reset running min/max and drop cached qparams.
+        """
+        self.min_val: torch.Tensor = torch.tensor(math.inf)
+        self.max_val: torch.Tensor = torch.tensor(-math.inf)
+        if hasattr(self, "_cached_scale"):
+            del self._cached_scale
+        if hasattr(self, "_cached_zp"):
+            del self._cached_zp
+
+    def load_qparams(self, scale: torch.Tensor, zp: torch.Tensor, *, lock: bool = True):
+        """
+        Inject externally computed qparams and optionally lock the observer.
+
+        When locked, subsequent `collect()` calls are ignored.
+        """
+        self._cached_scale = scale.detach()
+        self._cached_zp = zp.to(torch.int)
+        if lock:
+            self.enabled = False
+
+    @property
+    def has_qparams(self) -> bool:
+        return hasattr(self, "_cached_scale")
+
+    def compute_qparams(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        qmin, qmax = self.dtype.qmin, self.dtype.qmax
+        rng = self.max_val - self.min_val
+        eps = 1e-12
+
+        if self.qscheme.is_symmetric():
+            max_abs = torch.maximum(self.max_val.abs(), self.min_val.abs())
+            scale = torch.clamp(max_abs, min=eps) / qmax
+            zp = torch.zeros_like(scale, dtype=torch.int)
+            self._cached_scale, self._cached_zp = scale, zp
+            return scale, zp
+
+        if self.channel_axis is None:
+            if torch.all(rng.abs() < 1e-8):
+                C = self.min_val
+                if torch.allclose(C, torch.zeros_like(C)):
+                    scale = torch.ones_like(C)
+                    zp = torch.zeros_like(C, dtype=torch.int)
+                elif (C > 0).all():
+                    scale = torch.clamp(C, min=eps)
+                    zp = torch.zeros_like(C, dtype=torch.int)
+                else:
+                    scale = torch.clamp(C.abs(), min=eps)
+                    zp = torch.full_like(C, qmax, dtype=torch.int)
+            else:
+                scale = torch.clamp(rng, min=eps) / (qmax - qmin)
+                zp = (
+                    torch.round(qmin - self.min_val / scale)
+                    .clamp(qmin, qmax)
+                    .to(torch.int)
+                )
+        else:
+            scale = torch.clamp(rng, min=eps) / (qmax - qmin)
+            zp = (
+                torch.round(qmin - self.min_val / scale).clamp(qmin, qmax).to(torch.int)
+            )
+
+        self._cached_scale, self._cached_zp = scale, zp
+        return scale, zp
+
+    def fake_quant(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.has_qparams:
+            raise RuntimeError(
+                "Call compute_qparams()/freeze_qparams() or load_qparams() first."
+            )
+        scale, zp = self._cached_scale, self._cached_zp
+        if self.channel_axis is None:
+            return torch.fake_quantize_per_tensor_affine(
+                x,
+                scale=scale,
+                zero_point=zp,
+                quant_min=self.dtype.qmin,
+                quant_max=self.dtype.qmax,
+            )
+        else:
+            return torch.fake_quantize_per_channel_affine(
+                x,
+                scale=scale,
+                zero_point=zp,
+                axis=self.channel_axis,
+                quant_min=self.dtype.qmin,
+                quant_max=self.dtype.qmax,
+            )

tico/quantization/wrapq/observers/base.py

@@ -0,0 +1,98 @@
+# 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 Optional, Tuple
+
+import torch
+
+from tico.quantization.wrapq.dtypes import DType, UINT8
+from tico.quantization.wrapq.qscheme import QScheme
+
+
+class ObserverBase(ABC):
+    """
+    Minimal abstract base for all observers/quantizers.
+
+    Subclasses must implement:
+      - reset()
+      - collect(x)
+      - fake_quant(x)
+      - compute_qparams(): optional in practice for some observers (e.g., MX),
+        but still part of the interface; those can return None.
+    """
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        dtype: DType = UINT8,
+        qscheme: QScheme = QScheme.PER_TENSOR_ASYMM,
+        channel_axis: Optional[int] = None,  # None → per-tensor
+    ):
+        self.name = name
+        self.dtype = dtype
+        self.qscheme = qscheme
+        self.channel_axis = channel_axis if qscheme.is_per_channel() else None
+        self.enabled = True
+        self.reset()
+
+    @abstractmethod
+    def reset(self) -> None:
+        """Clear any running statistics or cached params."""
+        raise NotImplementedError
+
+    def collect(self, x: torch.Tensor) -> None:
+        """
+        Update running statistics with a new batch of data.
+
+        This base implementation guards on `enabled` and then calls `_update_stats(x)`.
+        Subclasses should implement `_update_stats(x)` instead of overriding `collect`.
+        """
+        if not self.enabled:
+            return
+        self._update_stats(x)
+
+    @abstractmethod
+    def _update_stats(self, x: torch.Tensor) -> None:
+        """
+        Update running statistics (min/max, hist, mse buffers, ...).
+
+        Must be implemented by subclasses (e.g., MinMax, EMA, Histogram, MSE).
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def fake_quant(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Apply the observer's quantization.
+        Implementations may or may not rely on qparams.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def compute_qparams(self) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Compute and (if applicable) cache quantization params.
+        Affine observers typically return (scale, zero_point).
+        Observers that do not use qparams (e.g., MX) may return None.
+        """
+        raise NotImplementedError
+
+    # String repr helps debugging
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(name={self.name}, dtype={str(self.dtype)}, "
+            f"qscheme={str(self.qscheme)}, channel_axis={self.channel_axis}, enabled={self.enabled})"
+        )

tico/quantization/wrapq/observers/ema.py

@@ -0,0 +1,62 @@
+# 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 tico.quantization.wrapq.observers.affine_base import AffineObserverBase
+from tico.quantization.wrapq.utils.reduce_utils import channelwise_minmax
+
+
+class EMAObserver(AffineObserverBase):
+    """
+    Exponential-Moving-Average min/max tracker.
+
+    Why?
+    -----
+    • Smoother than raw MinMax (reduces outlier shock).
+    • Much cheaper than histogram/MSE observers.
+
+    The update rule follows the common "momentum" form:
+
+        ema = momentum * ema + (1 - momentum) * new_value
+
+    With momentum → 0: FAST adaptation, momentum → 1: SLOW adaptation.
+    """
+
+    def __init__(
+        self,
+        *,
+        momentum: float = 0.9,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        assert 0.0 < momentum < 1.0, "momentum must be in (0, 1)"
+        self.momentum = momentum
+
+    @torch.no_grad()
+    def _update_stats(self, x: torch.Tensor):
+        if self.channel_axis is None:
+            curr_min, curr_max = x.min(), x.max()
+        else:
+            curr_min, curr_max = channelwise_minmax(x, self.channel_axis)
+
+        if (
+            torch.isinf(self.min_val).any() and torch.isinf(self.max_val).any()
+        ):  # first batch → hard init
+            self.min_val, self.max_val = curr_min, curr_max
+            return
+
+        m = self.momentum
+        self.min_val = m * self.min_val + (1 - m) * curr_min
+        self.max_val = m * self.max_val + (1 - m) * curr_max

tico/quantization/wrapq/observers/identity.py

@@ -0,0 +1,74 @@
+# 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.
+
+"""
+IdentityObserver: a "no-op" observer for FP-only modules.
+
+Motivation
+----------
+Some layers should stay in full precision even when the rest of the model
+is quantized.  Attaching an `IdentityObserver` satisfies the wrapper API
+(`_update_stats()`, `compute_qparams()`, `fake_quant()`) without actually
+performing any statistics gathering or fake-quantization.
+"""
+import torch
+
+from tico.quantization.wrapq.observers.affine_base import AffineObserverBase
+
+
+class IdentityObserver(AffineObserverBase):
+    """
+    Passthrough observer that NEVER alters the tensor.
+
+    • `_update_stats()`   → does nothing
+    • `compute_qparams()` → returns (1.0, 0) "dummy" q-params
+    • `fake_quant()`      → returns `x` unchanged
+    """
+
+    def __init__(self, **kwargs):
+        # Call parent so the usual fields (`dtype`, `qscheme`, …) exist,
+        # but immediately disable any stateful behaviour.
+        super().__init__(**kwargs)
+
+        # Deactivate statistics collection permanently.
+        self.enabled = False
+
+        # Pre-cache sentinel q-params so wrapper code that blindly
+        # accesses them won't crash.
+        self._cached_scale = torch.tensor(1.0)
+        self._cached_zp = torch.tensor(0, dtype=torch.int)
+
+    def reset(self) -> None:  # (simple override – nothing to do)
+        """No internal state to reset."""
+        pass
+
+    def _update_stats(self, x: torch.Tensor) -> None:
+        """Skip statistic collection entirely."""
+        return
+
+    def compute_qparams(self):
+        """
+        Return the pre-cached (scale, zero_point) tuple.
+
+        Keeping the signature identical to other observers allows uniform
+        lifecycle management in wrapper code.
+        """
+        return self._cached_scale, self._cached_zp
+
+    def fake_quant(self, x: torch.Tensor):
+        """Identity mapping — leaves `x` in FP."""
+        return x
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"

tico/quantization/wrapq/observers/minmax.py

@@ -0,0 +1,39 @@
+# 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 tico.quantization.wrapq.observers.affine_base import AffineObserverBase
+from tico.quantization.wrapq.utils.reduce_utils import channelwise_minmax
+
+
+class MinMaxObserver(AffineObserverBase):
+    """Plain min/max range tracker."""
+
+    @torch.no_grad()
+    def _update_stats(self, x: torch.Tensor) -> None:
+        """
+        Update running min/max with the incoming batch.
+
+        Per-tensor: use global min/max.
+        Per-channel: reduce all axes except the channel axis.
+        """
+        if self.channel_axis is None:
+            curr_min, curr_max = x.min(), x.max()
+        else:
+            curr_min, curr_max = channelwise_minmax(x, self.channel_axis)
+
+        # Broadcasting handles scalar-vs-vector cases
+        self.min_val = torch.minimum(self.min_val, curr_min)
+        self.max_val = torch.maximum(self.max_val, curr_max)

tico/quantization/wrapq/observers/mx.py

@@ -0,0 +1,60 @@
+# 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 tico.quantization.wrapq.observers.base import ObserverBase
+from tico.utils.mx.mx_ops import quantize_mx
+
+
+class MXObserver(ObserverBase):
+    """MX (micro-scaling) observer: no min/max, no affine qparams."""
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        elem_format: str = "int8",
+        axis: int = 0,
+        shared_exp_method: str = "max",
+        round: str = "nearest",
+        **base_kwargs,
+    ):
+        super().__init__(name=name, **base_kwargs)
+        self.elem_format = elem_format
+        self.axis = axis
+        self.shared_exp_method = shared_exp_method
+        self.round = round
+
+    def reset(self) -> None:
+        # No state to reset
+        return
+
+    @torch.no_grad()
+    def _update_stats(self, x: torch.Tensor) -> None:
+        # No stats required
+        return None
+
+    def compute_qparams(self):
+        # MX path does not produce affine qparams; keep interface contract.
+        return None
+
+    def fake_quant(self, x: torch.Tensor) -> torch.Tensor:
+        return quantize_mx(
+            x,
+            elem_format=self.elem_format,
+            axis=self.axis,
+            shared_exp_method=self.shared_exp_method,
+            round=self.round,
+        )

tico/quantization/wrapq/qscheme.py

@@ -0,0 +1,40 @@
+# 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 enum import auto, Enum
+
+
+class QScheme(Enum):
+    # ───── Per-tensor ────────────
+    PER_TENSOR_ASYMM = auto()
+    PER_TENSOR_SYMM = auto()
+    # ───── Per-channel ───────────
+    PER_CHANNEL_ASYMM = auto()
+    PER_CHANNEL_SYMM = auto()
+
+    # helper
+    def is_per_channel(self) -> bool:
+        return self in {
+            QScheme.PER_CHANNEL_ASYMM,
+            QScheme.PER_CHANNEL_SYMM,
+        }
+
+    def is_symmetric(self) -> bool:
+        return self in {
+            QScheme.PER_TENSOR_SYMM,
+            QScheme.PER_CHANNEL_SYMM,
+        }
+
+    def __str__(self) -> str:
+        return self.name.lower()

tico/quantization/wrapq/quantizer.py

@@ -0,0 +1,179 @@
+# 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
+import torch.nn as nn
+
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.quantizer import BaseQuantizer
+from tico.quantization.quantizer_registry import register_quantizer
+
+from tico.quantization.wrapq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+
+
+@register_quantizer(PTQConfig)
+class PTQQuantizer(BaseQuantizer):
+    """
+    Post-Training Quantization (PTQ) quantizer integrated with the public interface.
+
+    Features
+    --------
+    • Automatically wraps quantizable modules using PTQWrapper.
+    • Supports leaf-level (single-module) quantization (e.g., prepare(model.fc, PTQConfig())).
+    • Enforces strict wrapping if `strict_wrap=True`: raises NotImplementedError if
+      no quantizable module was found at any boundary.
+    • If `strict_wrap=False`, unquantizable modules are silently skipped.
+    """
+
+    def __init__(self, config: PTQConfig):
+        super().__init__(config)
+        self.qcfg: PTQConfig = config
+        self.strict_wrap: bool = bool(getattr(config, "strict_wrap", True))
+
+    @torch.no_grad()
+    def prepare(
+        self,
+        model: torch.nn.Module,
+        args: Optional[Any] = None,
+        kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        # Wrap the tree (or single module) according to strictness policy
+        model = self._wrap_supported(model, self.qcfg)
+
+        # Switch all quant modules into calibration mode
+        if isinstance(model, QuantModuleBase):
+            model.enable_calibration()
+        for m in model.modules():
+            if isinstance(m, QuantModuleBase):
+                m.enable_calibration()
+        return model
+
+    @torch.no_grad()
+    def convert(self, model):
+        # Freeze qparams across the tree (QUANT mode)
+        if isinstance(model, QuantModuleBase):
+            model.freeze_qparams()
+        for m in model.modules():
+            if isinstance(m, QuantModuleBase):
+                m.freeze_qparams()
+        return model
+
+    def _wrap_supported(
+        self,
+        root: nn.Module,
+        qcfg: PTQConfig,
+    ) -> nn.Module:
+        """
+        Recursively attempt to wrap boundaries. Strictness is applied at every boundary.
+        """
+        assert not isinstance(root, QuantModuleBase), "The module is already wrapped."
+
+        # Case A: HuggingFace-style transformers: model.model.layers
+        lm = getattr(root, "model", None)
+        layers = getattr(lm, "layers", None) if isinstance(lm, nn.Module) else None
+        if isinstance(layers, nn.ModuleList):
+            new_list = nn.ModuleList()
+            for idx, layer in enumerate(layers):
+                child_scope = f"layer{idx}"
+                child_cfg = qcfg.child(child_scope)
+
+                # Enforce strictness at the child boundary
+                wrapped = self._try_wrap(
+                    layer,
+                    child_cfg,
+                    fp_name=child_scope,
+                    raise_on_fail=self.strict_wrap,
+                )
+                new_list.append(wrapped)
+            lm.layers = new_list  # type: ignore[union-attr]
+            return root
+
+        # Case B: Containers
+        if isinstance(root, (nn.Sequential, nn.ModuleList)):
+            for i, child in enumerate(list(root)):
+                name = str(i)
+                child_cfg = qcfg.child(name)
+
+                wrapped = self._try_wrap(
+                    child, child_cfg, fp_name=name, raise_on_fail=self.strict_wrap
+                )
+                if wrapped is child:
+                    assert not self.strict_wrap
+                    wrapped = self._wrap_supported(wrapped, child_cfg)
+                root[i] = wrapped  # type: ignore[index]
+
+        if isinstance(root, nn.ModuleDict):
+            for k, child in list(root.items()):
+                name = k
+                child_cfg = qcfg.child(name)
+
+                wrapped = self._try_wrap(
+                    child, child_cfg, fp_name=name, raise_on_fail=self.strict_wrap
+                )
+                if wrapped is child:
+                    assert not self.strict_wrap
+                    wrapped = self._wrap_supported(wrapped, child_cfg)
+                root[k] = wrapped  # type: ignore[index]
+
+        # Case C: Leaf node
+        root_name = getattr(root, "_get_name", lambda: None)()
+        wrapped = self._try_wrap(
+            root, qcfg, fp_name=root_name, raise_on_fail=self.strict_wrap
+        )
+        if wrapped is not root:
+            return wrapped
+
+        assert not self.strict_wrap
+        # Case D: Named children
+        for name, child in list(root.named_children()):
+            child_cfg = qcfg.child(name)
+
+            wrapped = self._try_wrap(
+                child, child_cfg, fp_name=name, raise_on_fail=self.strict_wrap
+            )
+            if wrapped is child:
+                assert not self.strict_wrap
+                wrapped = self._wrap_supported(wrapped, child_cfg)
+            setattr(root, name, wrapped)
+
+        return root
+
+    def _try_wrap(
+        self,
+        module: nn.Module,
+        qcfg_for_child: PTQConfig,
+        *,
+        fp_name: Optional[str],
+        raise_on_fail: bool,
+    ) -> nn.Module:
+        """
+        Attempt to wrap a boundary with PTQWrapper.
+
+        Behavior:
+          • If PTQWrapper succeeds: return wrapped module.
+          • If PTQWrapper raises NotImplementedError:
+                - raise_on_fail=True  -> re-raise (strict)
+                - raise_on_fail=False -> return original module (permissive)
+        """
+        try:
+            return PTQWrapper(module, qcfg=qcfg_for_child, fp_name=fp_name)
+        except NotImplementedError as e:
+            if raise_on_fail:
+                raise NotImplementedError(
+                    f"PTQQuantizer: no quantization wrapper for {type(module).__name__}"
+                ) from e
+            return module

tico/quantization/wrapq/utils/__init__.py

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