tico 0.1.0.dev250924__py3-none-any.whl → 0.1.0.dev251109__py3-none-any.whl

tico/{experimental/quantization/ptq → quantization/wrapq}/examples/quantize_with_gptq.py

@@ -33,16 +33,14 @@ import tqdm
 from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from tico.experimental.quantization import convert, prepare
-from tico.experimental.quantization.config.gptq import GPTQConfig
-from tico.experimental.quantization.ptq.observers.affine_base import AffineObserverBase
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.utils.introspection import build_fqn_map
-from tico.experimental.quantization.ptq.utils.metrics import perplexity
-from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
-from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
-    QuantModuleBase,
-)
+from tico.quantization import convert, prepare
+from tico.quantization.config.gptq import GPTQConfig
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.wrapq.observers.affine_base import AffineObserverBase
+from tico.quantization.wrapq.utils.introspection import build_fqn_map
+from tico.quantization.wrapq.utils.metrics import perplexity
+from tico.quantization.wrapq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
 
 
 # Token-budget presets for activation calibration
@@ -215,22 +213,8 @@ def main():
     # 4. Wrap every layer with PTQWrapper (activation UINT-8)
     # -------------------------------------------------------------------------
     print("Wrapping layers with PTQWrapper …")
-    layers = q_m.model.layers
-    if not isinstance(layers, (list, torch.nn.ModuleList)):
-        raise TypeError(f"'model.layers' must be list/ModuleList, got {type(layers)}")
-
-    qcfg = QuantConfig()  # default: per-tensor UINT8
-    wrapped = torch.nn.ModuleList()
-    for idx, fp_layer in enumerate(layers):
-        layer_cfg = qcfg.child(f"layer{idx}")
-        wrapped.append(
-            PTQWrapper(
-                fp_layer,
-                qcfg=layer_cfg,
-                fp_name=m_to_fqn.get(fp_layer),
-            )
-        )
-    q_m.model.layers = wrapped
+    qcfg = PTQConfig()  # default: per-tensor UINT8
+    prepare(q_m, qcfg)
 
     # -------------------------------------------------------------------------
     # 5. Single-pass activation calibration
@@ -242,11 +226,7 @@ def main():
     calib_txt = " ".join(dataset_train["text"])[:CALIB_TOKENS]
     train_ids = tokenizer(calib_txt, return_tensors="pt").input_ids.to(device)
 
-    # (a) Enable CALIB mode on every QuantModuleBase
-    for l in q_m.model.layers:
-        l.enable_calibration()
-
-    # (b) Overwrite weight observers with GPTQ statistics
+    # Overwrite weight observers with GPTQ statistics
     if hasattr(q_m, "quantizers") and isinstance(q_m.quantizers, dict):
         inject_gptq_qparams(q_m, q_m.quantizers)
     else:
@@ -254,7 +234,7 @@ def main():
             "[Warn] q_m.quantizers not found or not a dict; skipping GPTQ qparam injection."
         )
 
-    # (c) Forward passes to collect activation ranges
+    # Forward passes to collect activation ranges
     iterator = range(0, train_ids.size(1) - 1, args.stride)
     if not args.no_tqdm:
         iterator = tqdm.tqdm(iterator, desc="Act-calibration")
@@ -262,9 +242,8 @@ def main():
         for i in iterator:
             q_m(train_ids[:, i : i + args.stride])
 
-    # (d) Freeze all Q-params (scale, zero-point)
-    for l in q_m.model.layers:
-        l.freeze_qparams()
+    # Freeze all Q-params (scale, zero-point)
+    convert(q_m)
 
     # -------------------------------------------------------------------------
     # 6. Evaluate perplexity on Wikitext-2

tico/{experimental/quantization/ptq → quantization/wrapq}/observers/affine_base.py

@@ -17,9 +17,9 @@ from typing import Optional, Tuple
 
 import torch
 
-from tico.experimental.quantization.ptq.dtypes import DType, UINT8
-from tico.experimental.quantization.ptq.observers.base import ObserverBase
-from tico.experimental.quantization.ptq.qscheme import QScheme
+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):

tico/{experimental/quantization/ptq → quantization/wrapq}/observers/base.py

@@ -17,8 +17,8 @@ from typing import Optional, Tuple
 
 import torch
 
-from tico.experimental.quantization.ptq.dtypes import DType, UINT8
-from tico.experimental.quantization.ptq.qscheme import QScheme
+from tico.quantization.wrapq.dtypes import DType, UINT8
+from tico.quantization.wrapq.qscheme import QScheme
 
 
 class ObserverBase(ABC):

tico/{experimental/quantization/ptq → quantization/wrapq}/observers/ema.py

@@ -14,8 +14,8 @@
 
 import torch
 
-from tico.experimental.quantization.ptq.observers.affine_base import AffineObserverBase
-from tico.experimental.quantization.ptq.utils.reduce_utils import channelwise_minmax
+from tico.quantization.wrapq.observers.affine_base import AffineObserverBase
+from tico.quantization.wrapq.utils.reduce_utils import channelwise_minmax
 
 
 class EMAObserver(AffineObserverBase):

tico/{experimental/quantization/ptq → quantization/wrapq}/observers/identity.py

@@ -24,7 +24,7 @@ performing any statistics gathering or fake-quantization.
 """
 import torch
 
-from tico.experimental.quantization.ptq.observers.affine_base import AffineObserverBase
+from tico.quantization.wrapq.observers.affine_base import AffineObserverBase
 
 
 class IdentityObserver(AffineObserverBase):

tico/{experimental/quantization/ptq → quantization/wrapq}/observers/minmax.py

@@ -14,8 +14,8 @@
 
 import torch
 
-from tico.experimental.quantization.ptq.observers.affine_base import AffineObserverBase
-from tico.experimental.quantization.ptq.utils.reduce_utils import channelwise_minmax
+from tico.quantization.wrapq.observers.affine_base import AffineObserverBase
+from tico.quantization.wrapq.utils.reduce_utils import channelwise_minmax
 
 
 class MinMaxObserver(AffineObserverBase):

tico/{experimental/quantization/ptq → quantization/wrapq}/observers/mx.py

@@ -14,7 +14,7 @@
 
 import torch
 
-from tico.experimental.quantization.ptq.observers.base import ObserverBase
+from tico.quantization.wrapq.observers.base import ObserverBase
 from tico.utils.mx.mx_ops import quantize_mx
 
 

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/{experimental/quantization/ptq → quantization/wrapq}/utils/introspection.py

@@ -16,11 +16,9 @@ from typing import Callable, Dict, List, Optional, Tuple
 
 import torch
 
-from tico.experimental.quantization.evaluation.metric import MetricCalculator
-from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
-from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
-    QuantModuleBase,
-)
+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]:

tico/{experimental/quantization/ptq → quantization/wrapq}/utils/metrics.py

@@ -98,7 +98,8 @@ def perplexity(
 
         input_ids = input_ids_full[:, begin:end]
         target_ids = input_ids.clone()
-        target_ids[:, :-trg_len] = ignore_index  # mask previously-scored tokens
+        # mask previously-scored tokens
+        target_ids[:, :-trg_len] = ignore_index  # type: ignore[assignment]
 
         with torch.no_grad():
             outputs = model(input_ids, labels=target_ids)
@@ -106,7 +107,7 @@ def perplexity(
             neg_log_likelihood = outputs.loss
 
         # exact number of labels that contributed to loss
-        loss_tokens = (target_ids[:, 1:] != ignore_index).sum().item()
+        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)
 

tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/fairseq/__init__.py

@@ -1,4 +1,4 @@
-from tico.experimental.quantization.ptq.wrappers.fairseq.quant_mha import (
+from tico.quantization.wrapq.wrappers.fairseq.quant_mha import (
     QuantFairseqMultiheadAttention,
 )
 

tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/fairseq/quant_decoder.py

@@ -25,12 +25,10 @@ import torch
 import torch.nn.functional as F
 from torch import nn, Tensor
 
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
-from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
-    QuantModuleBase,
-)
-from tico.experimental.quantization.ptq.wrappers.registry import try_register
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.wrapq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+from tico.quantization.wrapq.wrappers.registry import try_register
 
 
 @try_register("fairseq.models.transformer.TransformerDecoderBase")
@@ -53,7 +51,7 @@ class QuantFairseqDecoder(QuantModuleBase):
         self,
         fp_decoder: nn.Module,
         *,
-        qcfg: Optional[QuantConfig] = None,
+        qcfg: Optional[PTQConfig] = None,
         fp_name: Optional[str] = None,
     ):
         super().__init__(qcfg, fp_name=fp_name)
@@ -116,7 +114,7 @@ class QuantFairseqDecoder(QuantModuleBase):
 
         prefix = _safe_prefix(fp_name)
 
-        # Prepare child QuantConfig namespaces: layers/<idx>
+        # Prepare child PTQConfig namespaces: layers/<idx>
         layers_qcfg = qcfg.child("layers") if qcfg else None
         for i, layer in enumerate(fp_layers):
             child_cfg = layers_qcfg.child(str(i)) if layers_qcfg else None

tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/fairseq/quant_decoder_layer.py

@@ -23,15 +23,13 @@ from typing import Dict, Iterable, List, Optional, Tuple
 import torch
 from torch import nn, Tensor
 
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.wrappers.fairseq.quant_mha import (
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.wrapq.wrappers.fairseq.quant_mha import (
     QuantFairseqMultiheadAttention,
 )
-from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
-from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
-    QuantModuleBase,
-)
-from tico.experimental.quantization.ptq.wrappers.registry import try_register
+from tico.quantization.wrapq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+from tico.quantization.wrapq.wrappers.registry import try_register
 
 
 @try_register("fairseq.modules.transformer_layer.TransformerDecoderLayerBase")
@@ -55,7 +53,7 @@ class QuantFairseqDecoderLayer(QuantModuleBase):
         self,
         fp_layer: nn.Module,
         *,
-        qcfg: Optional[QuantConfig] = None,
+        qcfg: Optional[PTQConfig] = None,
         fp_name: Optional[str] = None,
     ):
         super().__init__(qcfg, fp_name=fp_name)

tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/fairseq/quant_encoder.py

@@ -25,12 +25,10 @@ import torch
 import torch.nn as nn
 from torch import Tensor
 
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
-from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
-    QuantModuleBase,
-)
-from tico.experimental.quantization.ptq.wrappers.registry import try_register
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.wrapq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+from tico.quantization.wrapq.wrappers.registry import try_register
 
 
 @try_register("fairseq.models.transformer.TransformerEncoderBase")
@@ -56,7 +54,7 @@ class QuantFairseqEncoder(QuantModuleBase):
         self,
         fp_encoder: nn.Module,
         *,
-        qcfg: Optional[QuantConfig] = None,
+        qcfg: Optional[PTQConfig] = None,
         fp_name: Optional[str] = None,
         use_external_inputs: bool = False,  # export-mode flag
         return_type: Literal["tensor", "dict"] = "dict",
@@ -100,7 +98,7 @@ class QuantFairseqEncoder(QuantModuleBase):
         fp_layers = list(fp_encoder.layers)  # type: ignore[arg-type]
         self.layers = nn.ModuleList()
 
-        # Prepare child QuantConfig namespaces: layers/<idx>
+        # Prepare child PTQConfig namespaces: layers/<idx>
         layers_qcfg = qcfg.child("layers") if qcfg else None
         for i, layer in enumerate(fp_layers):
             child_cfg = layers_qcfg.child(str(i)) if layers_qcfg else None

tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/fairseq/quant_encoder_layer.py

@@ -23,15 +23,13 @@ from typing import Optional
 import torch.nn as nn
 from torch import Tensor
 
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.wrappers.fairseq.quant_mha import (
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.wrapq.wrappers.fairseq.quant_mha import (
     QuantFairseqMultiheadAttention,
 )
-from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
-from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
-    QuantModuleBase,
-)
-from tico.experimental.quantization.ptq.wrappers.registry import try_register
+from tico.quantization.wrapq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+from tico.quantization.wrapq.wrappers.registry import try_register
 
 
 @try_register("fairseq.modules.transformer_layer.TransformerEncoderLayerBase")
@@ -49,7 +47,7 @@ class QuantFairseqEncoderLayer(QuantModuleBase):
         self,
         fp_layer: nn.Module,
         *,
-        qcfg: Optional[QuantConfig] = None,
+        qcfg: Optional[PTQConfig] = None,
         fp_name: Optional[str] = None,
     ):
         super().__init__(qcfg, fp_name=fp_name)

tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/fairseq/quant_mha.py

@@ -24,12 +24,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
-from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
-    QuantModuleBase,
-)
-from tico.experimental.quantization.ptq.wrappers.registry import try_register
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.wrapq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+from tico.quantization.wrapq.wrappers.registry import try_register
 
 
 @try_register("fairseq.modules.multihead_attention.MultiheadAttention")
@@ -59,7 +57,7 @@ class QuantFairseqMultiheadAttention(QuantModuleBase):
         self,
         fp_attn: nn.Module,
         *,
-        qcfg: Optional[QuantConfig] = None,
+        qcfg: Optional[PTQConfig] = None,
         fp_name: Optional[str] = None,
         max_seq: int = 4096,
         use_static_causal: bool = False,

tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/llama/quant_attn.py

@@ -17,12 +17,10 @@ from typing import Optional, Tuple
 import torch
 import torch.nn as nn
 
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
-from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
-    QuantModuleBase,
-)
-from tico.experimental.quantization.ptq.wrappers.registry import try_register
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.wrapq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+from tico.quantization.wrapq.wrappers.registry import try_register
 
 
 @try_register(
@@ -34,7 +32,7 @@ class QuantLlamaAttention(QuantModuleBase):
         self,
         fp_attn: nn.Module,
         *,
-        qcfg: Optional[QuantConfig] = None,
+        qcfg: Optional[PTQConfig] = None,
         fp_name: Optional[str] = None,
     ):
         super().__init__(qcfg, fp_name=fp_name)

tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/llama/quant_decoder_layer.py

@@ -17,16 +17,12 @@ from typing import Optional, Tuple
 import torch
 import torch.nn as nn
 
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.wrappers.llama.quant_attn import (
-    QuantLlamaAttention,
-)
-from tico.experimental.quantization.ptq.wrappers.llama.quant_mlp import QuantLlamaMLP
-from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
-from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
-    QuantModuleBase,
-)
-from tico.experimental.quantization.ptq.wrappers.registry import try_register
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.wrapq.wrappers.llama.quant_attn import QuantLlamaAttention
+from tico.quantization.wrapq.wrappers.llama.quant_mlp import QuantLlamaMLP
+from tico.quantization.wrapq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+from tico.quantization.wrapq.wrappers.registry import try_register
 
 
 @try_register("transformers.models.llama.modeling_llama.LlamaDecoderLayer")
@@ -56,7 +52,7 @@ class QuantLlamaDecoderLayer(QuantModuleBase):
         self,
         fp_layer: nn.Module,
         *,
-        qcfg: Optional[QuantConfig] = None,
+        qcfg: Optional[PTQConfig] = None,
         fp_name: Optional[str] = None,
         return_type: Optional[str] = None,
     ):
@@ -165,7 +161,7 @@ class QuantLlamaDecoderLayer(QuantModuleBase):
         # - If use_cache: always return (hidden_states, present_key_value)
         # - Else: return as configured (tuple/tensor) for HF compatibility
         if use_cache:
-            return hidden_states, present_key_value
+            return hidden_states, present_key_value  # type: ignore[return-value]
 
         if self.return_type == "tuple":
             return (hidden_states,)

tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/llama/quant_mlp.py

@@ -17,12 +17,10 @@ from typing import Optional
 import torch
 import torch.nn as nn
 
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
-from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
-    QuantModuleBase,
-)
-from tico.experimental.quantization.ptq.wrappers.registry import try_register
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.wrapq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+from tico.quantization.wrapq.wrappers.registry import try_register
 
 
 @try_register("transformers.models.llama.modeling_llama.LlamaMLP")
@@ -31,7 +29,7 @@ class QuantLlamaMLP(QuantModuleBase):
         self,
         mlp_fp: nn.Module,
         *,
-        qcfg: Optional[QuantConfig] = None,
+        qcfg: Optional[PTQConfig] = None,
         fp_name: Optional[str] = None,
     ):
         super().__init__(qcfg, fp_name=fp_name)

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

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

tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/nn/quant_layernorm.py

@@ -17,12 +17,11 @@ from typing import Iterable, Optional, Tuple
 import torch
 import torch.nn as nn
 
-from tico.experimental.quantization.ptq.mode import Mode
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
-    QuantModuleBase,
-)
-from tico.experimental.quantization.ptq.wrappers.registry import register
+from tico.quantization.config.ptq import PTQConfig
+
+from tico.quantization.wrapq.mode import Mode
+from tico.quantization.wrapq.wrappers.quant_module_base import QuantModuleBase
+from tico.quantization.wrapq.wrappers.registry import register
 
 
 @register(nn.LayerNorm)
@@ -46,7 +45,7 @@ class QuantLayerNorm(QuantModuleBase):
         self,
         fp: nn.LayerNorm,
         *,
-        qcfg: Optional[QuantConfig] = None,
+        qcfg: Optional[PTQConfig] = None,
         fp_name: Optional[str] = None
     ):
         super().__init__(qcfg, fp_name=fp_name)