tico 0.1.0.dev250803__py3-none-any.whl → 0.1.0.dev251106__py3-none-any.whl

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

tico/quantization/wrapq/utils/introspection.py

@@ -0,0 +1,167 @@
+# 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, Dict, List, Optional, Tuple
+
+import torch
+
+from tico.quantization.evaluation.metric import MetricCalculator
+from tico.quantization.wrapq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+
+
+def build_fqn_map(root: torch.nn.Module) -> dict[torch.nn.Module, str]:
+    """
+    Return {module_object: full_qualified_name} without touching the modules.
+    """
+    return {m: n for n, m in root.named_modules()}
+
+
+def save_fp_outputs(
+    model: torch.nn.Module,
+) -> Tuple[List[torch.utils.hooks.RemovableHandle], Dict[str, torch.Tensor]]:
+    """
+    Register forward-hooks on every `QuantModuleBase` wrapper itself (not the
+    wrapped `module`) and cache its output while the wrapper runs in CALIB mode.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+        The model whose wrappers are already switched to CALIB mode
+        (`enable_calibration()` has been called).
+
+    Returns
+    -------
+    handles : list[RemovableHandle]
+        Hook handles; call `.remove()` on each one to detach the hooks.
+    cache : dict[str, torch.Tensor]
+        Mapping "wrapper-name → cached FP32 activation" captured from the first
+        forward pass. Keys default to `wrapper.fp_name`; if that attribute is
+        `None`, the `id(wrapper)` string is used instead.
+    """
+    cache: Dict[str, torch.Tensor] = {}
+    handles: List[torch.utils.hooks.RemovableHandle] = []
+
+    def _save(name: str):
+        def hook(_, __, out: torch.Tensor | Tuple):
+            if isinstance(out, tuple):
+                out = out[0]
+            assert isinstance(out, torch.Tensor)
+            cache[name] = out.detach()
+
+        return hook
+
+    for m in model.modules():
+        if isinstance(m, QuantModuleBase):
+            name = m.fp_name or str(id(m))
+            handles.append(m.register_forward_hook(_save(name)))
+
+    return handles, cache
+
+
+def compare_layer_outputs(
+    model: torch.nn.Module,
+    cache: Dict[str, torch.Tensor],
+    *,
+    metrics: Optional[List[str]] = None,
+    custom_metrics: Optional[Dict[str, Callable]] = None,
+    rtol: float = 1e-3,
+    atol: float = 1e-3,
+    collect: bool = False,
+):
+    """
+    Register forward-hooks on every `QuantModuleBase` wrapper to compare its
+    QUANT-mode output to the FP32 reference saved by `save_fp_outputs()`.
+
+    Each hook prints a per-layer diff report:
+
+        ✓  layer_name  max=1.23e-02  mean=8.45e-04     (within tolerance)
+        ⚠️ layer_name  max=3.07e+00  mean=5.12e-01     (exceeds tolerance)
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+        The model whose wrappers are now in QUANT mode
+        (`freeze_qparams()` has been called).
+    cache : dict[str, torch.Tensor]
+        The reference activations captured during CALIB mode.
+    metrics
+        Metrics to compute. Defaults to `["diff"]`. Add `peir` to print PEIR.
+    custom_metrics
+        Optional user metric functions. Same signature as built-ins.
+    rtol, atol : float, optional
+        Relative / absolute tolerances used to flag large deviations
+        (similar to `torch.allclose` semantics).
+    collect : bool, optional
+        • False (default) → print one-line report per layer, return `None`
+        • True            → suppress printing, return a nested dict
+                                {layer_name -> {metric -> value}}
+
+    Returns
+    -------
+    handles
+        Hook handles; call `.remove()` once diffing is complete.
+    results
+        Only if *collect* is True.
+    """
+    metrics = metrics or ["diff"]
+    calc = MetricCalculator(custom_metrics)
+    handles: List[torch.utils.hooks.RemovableHandle] = []
+    results: Dict[
+        str, Dict[str, float]
+    ] = {}  # Dict[layer_name, Dict[metric_name, value]]
+
+    def _cmp(name: str):
+        ref = cache.get(name)
+
+        def hook(_, __, out):
+            if ref is None:
+                if not collect:
+                    print(f"[{name}]  no cached reference")
+                return
+            if isinstance(out, tuple):
+                out = out[0]
+            assert isinstance(out, torch.Tensor)
+
+            # Compute all requested metrics
+            res = calc.compute([ref], [out], metrics)  # lists with length-1 tensors
+            res = {k: v[0] for k, v in res.items()}  # flatten
+
+            if collect:
+                results[name] = res  # type: ignore[assignment]
+                return
+
+            # Pretty print ------------------------------------------------ #
+            diff_val = res.get("diff") or res.get("max_abs_diff")
+            thresh = atol + rtol * ref.abs().max().item()
+            flag = "⚠️" if (diff_val is not None and diff_val > thresh) else "✓"  # type: ignore[operator]
+
+            pieces = [f"{flag} {name:45s}"]
+            for key, val in res.items():
+                pieces.append(f"{key}={val:<7.4}")
+            print("  ".join(pieces))
+
+        return hook
+
+    for m in model.modules():
+        if isinstance(m, PTQWrapper):
+            # skip the internal fp module inside the wrapper
+            continue
+        if isinstance(m, QuantModuleBase):
+            lname = m.fp_name or str(id(m))
+            handles.append(m.register_forward_hook(_cmp(lname)))
+
+    if collect:
+        return handles, results
+    return handles

tico/quantization/wrapq/utils/metrics.py

@@ -0,0 +1,124 @@
+# 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 Optional
+
+import torch
+import tqdm
+
+
+def perplexity(
+    model: torch.nn.Module,
+    encodings: torch.Tensor,
+    device: torch.device | str,
+    *,
+    max_length: Optional[int] = None,
+    stride: int = 512,
+    ignore_index: int | None = -100,
+    show_progress: bool = True,
+) -> float:
+    """
+    Compute perplexity (PPL) using a "strided sliding-window"
+     evaluation strategy.
+
+    The function:
+    1. Splits the token sequence into overlapping windows of length
+       `max_length` (model context size).
+    2. Masks tokens that were already scored in previous windows
+       (`labels == -100`), so each token's negative log-likelihood (NLL)
+       is counted EXACTLY once.
+    3. Aggregates token-wise NLL to return corpus-level PPL.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+        Causal LM loaded in evaluation mode (`model.eval()`).
+    encodings : torch.Tensor | transformers.BatchEncoding
+        Tokenised corpus.  If a `BatchEncoding` is passed, its
+        `.input_ids` field is used.  Shape must be `(1, seq_len)`.
+    device : torch.device | str
+        CUDA or CPU device on which to run evaluation.
+    max_length : int, optional
+        Context window size.  Defaults to `model.config.max_position_embeddings`.
+    stride : int, default = 512
+        Step size by which the sliding window advances.  Must satisfy
+        `1 ≤ stride ≤ max_length`.
+    ignore_index : int, default = -100
+        Label value to ignore in loss computation. This should match
+        the `ignore_index` used by the model's internal
+        `CrossEntropyLoss`. For Hugging Face causal LMs, the
+        convention is `-100`.
+    show_progress : bool, default = True
+        If True, displays a tqdm progess bar while evaluating.
+
+    Returns
+    -------
+    float
+        Corpus-level perplexity.
+    """
+    # -------- input preparation -------- #
+    try:
+        # transformers.BatchEncoding has `input_ids`
+        input_ids_full = encodings.input_ids  # type: ignore[attr-defined]
+    except AttributeError:  # already a tensor
+        input_ids_full = encodings
+    assert isinstance(input_ids_full, torch.Tensor)
+    input_ids_full = input_ids_full.to(device)
+
+    if max_length is None:
+        assert hasattr(model, "config")
+        assert hasattr(model.config, "max_position_embeddings")
+        assert isinstance(model.config.max_position_embeddings, int)
+        max_length = model.config.max_position_embeddings
+    assert max_length is not None
+    assert (
+        1 <= stride <= max_length
+    ), f"stride ({stride}) must be in [1, max_length ({max_length})]"
+
+    seq_len = input_ids_full.size(1)
+    nll_sum = 0.0
+    n_tokens = 0
+    prev_end = 0
+
+    # -------- main loop -------- #
+    for begin in tqdm.trange(0, seq_len, stride, desc="PPL", disable=not show_progress):
+        end = min(begin + max_length, seq_len)
+        trg_len = end - prev_end  # fresh tokens in this window
+
+        input_ids = input_ids_full[:, begin:end]
+        target_ids = input_ids.clone()
+        # mask previously-scored tokens
+        target_ids[:, :-trg_len] = ignore_index  # type: ignore[assignment]
+
+        with torch.no_grad():
+            outputs = model(input_ids, labels=target_ids)
+            # loss is already averaged over non-masked labels
+            neg_log_likelihood = outputs.loss
+
+        # exact number of labels that contributed to loss
+        loss_tokens = (target_ids[:, 1:] != ignore_index).sum().item()  # type: ignore[attr-defined]
+        nll_sum += neg_log_likelihood * loss_tokens
+        n_tokens += int(loss_tokens)
+
+        prev_end = end
+        if end == seq_len:
+            break
+
+    avg_nll: float | torch.Tensor = nll_sum / n_tokens
+    if not isinstance(avg_nll, torch.Tensor):
+        avg_nll = torch.tensor(avg_nll)
+    assert isinstance(avg_nll, torch.Tensor)
+    ppl = torch.exp(avg_nll)
+
+    return ppl.item()

tico/quantization/wrapq/utils/reduce_utils.py

@@ -0,0 +1,25 @@
+# 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 channelwise_minmax(x: torch.Tensor, channel_axis: int):
+    """
+    Compute per-channel (min, max) by reducing all axes except `channel_axis`.
+    """
+    channel_axis = channel_axis % x.ndim  # handle negative indices safely
+    dims = tuple(d for d in range(x.ndim) if d != channel_axis)
+
+    return x.amin(dim=dims), x.amax(dim=dims)

tico/quantization/wrapq/wrappers/__init__.py

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

tico/quantization/wrapq/wrappers/fairseq/__init__.py

@@ -0,0 +1,5 @@
+from tico.quantization.wrapq.wrappers.fairseq.quant_mha import (
+    QuantFairseqMultiheadAttention,
+)
+
+__all__ = ["QuantFairseqMultiheadAttention"]