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

tico/{experimental/quantization → quantization}/algorithm/pt2e/quantizer.py

@@ -18,13 +18,13 @@ import torch
 
 from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
 
-from tico.experimental.quantization.algorithm.pt2e.annotation.annotator import (
+from tico.quantization.algorithm.pt2e.annotation.annotator import (
     get_asymmetric_quantization_config,
     PT2EAnnotator,
 )
-from tico.experimental.quantization.config.pt2e import PT2EConfig
-from tico.experimental.quantization.quantizer import BaseQuantizer
-from tico.experimental.quantization.quantizer_registry import register_quantizer
+from tico.quantization.config.pt2e import PT2EConfig
+from tico.quantization.quantizer import BaseQuantizer
+from tico.quantization.quantizer_registry import register_quantizer
 
 
 @register_quantizer(PT2EConfig)

tico/{experimental/quantization → quantization}/algorithm/pt2e/utils.py

@@ -20,9 +20,7 @@ import torch
 from torch.ao.quantization.quantizer import QuantizationSpec
 from torch.ao.quantization.quantizer.utils import _get_module_name_filter
 
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 
 
 def get_module_type_filter(tp: Callable):

tico/{experimental/quantization → quantization}/algorithm/smoothquant/quantizer.py

@@ -16,16 +16,12 @@ 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.smoothquant import SmoothQuantConfig
-from tico.experimental.quantization.quantizer import BaseQuantizer
-from tico.experimental.quantization.quantizer_registry import register_quantizer
+from tico.quantization.algorithm.smoothquant.observer import ChannelwiseMaxActsObserver
+
+from tico.quantization.algorithm.smoothquant.smooth_quant import apply_smoothing
+from tico.quantization.config.smoothquant import SmoothQuantConfig
+from tico.quantization.quantizer import BaseQuantizer
+from tico.quantization.quantizer_registry import register_quantizer
 
 
 @register_quantizer(SmoothQuantConfig)

tico/quantization/config/fpi_gptq.py

@@ -0,0 +1,29 @@
+# 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 tico.quantization.config.gptq import GPTQConfig
+
+
+class FPIGPTQConfig(GPTQConfig):
+    """
+    Configuration for FPIGPTQ (Fixed Point Iteration).
+    """
+
+    def __init__(self, verbose: bool = False, show_progress: bool = True):
+        self.verbose = verbose
+        self.show_progress = show_progress
+
+    @property
+    def name(self) -> str:
+        return "fpi_gptq"

tico/{experimental/quantization → quantization}/config/gptq.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from tico.experimental.quantization.config.base import BaseConfig
+from tico.quantization.config.base import BaseConfig
 
 
 class GPTQConfig(BaseConfig):

tico/{experimental/quantization → quantization}/config/pt2e.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from tico.experimental.quantization.config.base import BaseConfig
+from tico.quantization.config.base import BaseConfig
 
 
 class PT2EConfig(BaseConfig):

tico/{experimental/quantization/ptq/quant_config.py → quantization/config/ptq.py}

@@ -15,14 +15,15 @@
 from dataclasses import dataclass, field
 from typing import Any, Dict, Mapping, Type
 
-from tico.experimental.quantization.ptq.dtypes import DType
-from tico.experimental.quantization.ptq.observers.base import ObserverBase
-from tico.experimental.quantization.ptq.observers.minmax import MinMaxObserver
-from tico.experimental.quantization.ptq.qscheme import QScheme
+from tico.quantization.config.base import BaseConfig
+from tico.quantization.wrapq.dtypes import DType
+from tico.quantization.wrapq.observers.base import ObserverBase
+from tico.quantization.wrapq.observers.minmax import MinMaxObserver
+from tico.quantization.wrapq.qscheme import QScheme
 
 
 @dataclass
-class QuantConfig:
+class PTQConfig(BaseConfig):
     """
     One object describes the quantization preferences for a single wrapper
     and its descendants.
@@ -54,9 +55,9 @@ class QuantConfig:
     Example
     -------
     ```python
-    from ptq.observers import PercentileObserver
+    from wrapq.observers import PercentileObserver
 
-    cfg = QuantConfig(
+    cfg = PTQConfig(
         default_dtype   = DType.uint(8),
         default_qscheme  = QScheme.PER_TENSOR_SYMM,        # <- global scheme
         default_observer = PercentileObserver,             # <- global algorithm
@@ -74,6 +75,12 @@ class QuantConfig:
     default_observer: Type[ObserverBase] = MinMaxObserver
     default_qscheme: QScheme = QScheme.PER_TENSOR_ASYMM
     overrides: Mapping[str, Mapping[str, Any]] = field(default_factory=dict)
+    # If True, any module that cannot be wrapped will raise.
+    strict_wrap: bool = True
+
+    @property
+    def name(self) -> str:
+        return "ptq"
 
     def get_kwargs(self, obs_name: str) -> Dict[str, Any]:
         """
@@ -87,7 +94,7 @@ class QuantConfig:
         """
         return dict(self.overrides.get(obs_name, {}))
 
-    def child(self, scope: str) -> "QuantConfig":
+    def child(self, scope: str) -> "PTQConfig":
         """
         Produce a *view* for a child wrapper.
 
@@ -100,12 +107,13 @@ class QuantConfig:
         Other scopes remain invisible to the child.
         """
         sub_overrides = self.overrides.get(scope, {})
-        return QuantConfig(
+        return PTQConfig(
             self.default_dtype,
             self.default_observer,
             default_qscheme=self.default_qscheme,
             overrides=sub_overrides,
+            strict_wrap=self.strict_wrap,
         )
 
     def __repr__(self):
-        return f"QuantConfig(default_dtype={self.default_dtype}, default_observer={self.default_observer}, default_qscheme={self.default_qscheme}, overrides={dict(self.overrides)})"
+        return f"PTQConfig(default_dtype={self.default_dtype}, default_observer={self.default_observer}, default_qscheme={self.default_qscheme}, overrides={dict(self.overrides)}, strict_wrap={self.strict_wrap})"

tico/{experimental/quantization → quantization}/config/smoothquant.py

@@ -14,7 +14,7 @@
 
 from typing import Dict, Literal, Optional
 
-from tico.experimental.quantization.config.base import BaseConfig
+from tico.quantization.config.base import BaseConfig
 
 
 class SmoothQuantConfig(BaseConfig):

tico/{experimental/quantization → quantization}/evaluation/evaluate.py

@@ -20,18 +20,12 @@ import torch
 from circle_schema import circle
 from torch.utils import _pytree as pytree
 
-from tico.experimental.quantization.evaluation.backend import BACKEND
-from tico.experimental.quantization.evaluation.executor.backend_executor import (
-    BackendExecutor,
-)
-from tico.experimental.quantization.evaluation.executor.circle_executor import (
-    CircleExecutor,
-)
-from tico.experimental.quantization.evaluation.executor.triv24_executor import (
-    Triv24Executor,
-)
-from tico.experimental.quantization.evaluation.metric import MetricCalculator
-from tico.experimental.quantization.evaluation.utils import (
+from tico.quantization.evaluation.backend import BACKEND
+from tico.quantization.evaluation.executor.backend_executor import BackendExecutor
+from tico.quantization.evaluation.executor.circle_executor import CircleExecutor
+from tico.quantization.evaluation.executor.triv24_executor import Triv24Executor
+from tico.quantization.evaluation.metric import MetricCalculator
+from tico.quantization.evaluation.utils import (
     ensure_list,
     find_invalid_types,
     get_graph_input_output,

tico/{experimental/quantization → quantization}/evaluation/executor/circle_executor.py

@@ -19,9 +19,7 @@ from typing import List
 import numpy as np
 import torch
 
-from tico.experimental.quantization.evaluation.executor.backend_executor import (
-    BackendExecutor,
-)
+from tico.quantization.evaluation.executor.backend_executor import BackendExecutor
 from tico.utils.model import CircleModel
 from tico.utils.utils import run_bash_cmd
 

tico/{experimental/quantization → quantization}/evaluation/executor/triv24_executor.py

@@ -20,10 +20,8 @@ import numpy as np
 import torch
 from circle_schema import circle
 
-from tico.experimental.quantization.evaluation.executor.backend_executor import (
-    BackendExecutor,
-)
-from tico.experimental.quantization.evaluation.utils import (
+from tico.quantization.evaluation.executor.backend_executor import BackendExecutor
+from tico.quantization.evaluation.utils import (
     dequantize,
     get_graph_input_output,
     quantize,

tico/{experimental/quantization → quantization}/evaluation/utils.py

@@ -44,7 +44,7 @@ def quantize(
         data = np.array(data)
     # Perfrom quantization
     if not scale:
-        logger.warn("WARNING: scale value is 0. 1e-7 will be used instead.")
+        logger.warning("WARNING: scale value is 0. 1e-7 will be used instead.")
         scale = 1e-7
     rescaled = np.round(data / scale) + zero_point
     # Clamp the values

tico/{experimental/quantization → quantization}/public_interface.py

@@ -17,11 +17,11 @@ from typing import Any, Dict, Optional
 
 import torch
 
-from tico.experimental.quantization.algorithm.gptq.quantizer import GPTQQuantizer
-from tico.experimental.quantization.algorithm.pt2e.quantizer import PT2EQuantizer
-from tico.experimental.quantization.config.base import BaseConfig
-from tico.experimental.quantization.quantizer import BaseQuantizer
-from tico.experimental.quantization.quantizer_registry import get_quantizer
+from tico.quantization.algorithm.gptq.quantizer import GPTQQuantizer
+from tico.quantization.algorithm.pt2e.quantizer import PT2EQuantizer
+from tico.quantization.config.base import BaseConfig
+from tico.quantization.quantizer import BaseQuantizer
+from tico.quantization.quantizer_registry import get_quantizer
 
 
 QUANTIZER_ATTRIBUTE_NAME = "tico_quantizer"
@@ -32,7 +32,7 @@ def prepare(
     quant_config: BaseConfig,
     args: Optional[Any] = None,
     kwargs: Optional[Dict[str, Any]] = None,
-    inplace: Optional[bool] = False,
+    inplace: Optional[bool] = True,
 ):
     """
     Prepare the model for quantization using the provided configuration.
@@ -68,7 +68,7 @@ def prepare(
     return model
 
 
-def convert(model, inplace: Optional[bool] = False):
+def convert(model, inplace: Optional[bool] = True):
     """
     Convert the prepared model to a quantized model using the provided configuration.
 

tico/{experimental/quantization → quantization}/quantizer.py

@@ -17,7 +17,7 @@ from typing import Any, Dict, Optional
 
 import torch
 
-from tico.experimental.quantization.config.base import BaseConfig
+from tico.quantization.config.base import BaseConfig
 
 
 class BaseQuantizer(ABC):

tico/{experimental/quantization → quantization}/quantizer_registry.py

@@ -15,8 +15,8 @@
 import importlib
 from typing import Dict, Optional, Type, TypeVar
 
-from tico.experimental.quantization.config.base import BaseConfig
-from tico.experimental.quantization.quantizer import BaseQuantizer
+from tico.quantization.config.base import BaseConfig
+from tico.quantization.quantizer import BaseQuantizer
 
 TQ = TypeVar("TQ", bound=BaseQuantizer)
 
@@ -53,14 +53,15 @@ def get_quantizer(cfg: BaseConfig) -> BaseQuantizer:
     # Lazy import by naming convention
     name = getattr(cfg, "name", None)
     if name:
-        try:
-            importlib.import_module(
-                f"tico.experimental.quantization.algorithm.{name}.quantizer"
-            )
-        except Exception as e:
-            raise RuntimeError(
-                f"Failed to import quantizer module for config name='{name}': {e}"
-            )
+        if name == "ptq":
+            importlib.import_module(f"tico.quantization.wrapq.quantizer")
+        else:
+            try:
+                importlib.import_module(f"tico.quantization.algorithm.{name}.quantizer")
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to import quantizer module for config name='{name}': {e}"
+                )
 
     qcls = _lookup(cfg)
     if qcls is not None:

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

@@ -22,16 +22,15 @@
 
 import argparse
 import sys
-from typing import Optional
 
 import torch
 import tqdm
 from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.utils.metrics import perplexity
-from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization import convert, prepare
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.wrapq.utils.metrics import perplexity
 
 # Token-budget presets for activation calibration
 TOKENS: dict[str, int] = {
@@ -165,13 +164,8 @@ def main():
         # ---------------------------------------------------------------------
         # 2. Wrap every Transformer layer with PTQWrapper
         # ---------------------------------------------------------------------
-        qcfg = QuantConfig()  # all-uint8 defaults
-
-        wrapped_layers = torch.nn.ModuleList()
-        for idx, layer in enumerate(uint8_model.model.layers):
-            layer_cfg = qcfg.child(f"layer{idx}")
-            wrapped_layers.append(PTQWrapper(layer, qcfg=layer_cfg))
-        uint8_model.model.layers = wrapped_layers
+        qcfg = PTQConfig()  # all-uint8 defaults
+        prepare(uint8_model, qcfg)
 
         # ---------------------------------------------------------------------
         # 3. Single-pass activation calibration
@@ -182,11 +176,7 @@ def main():
         )[:CALIB_TOKENS]
         ids = tokenizer(calib_txt, return_tensors="pt").input_ids.to(device)
 
-        # (a) switch every QuantModuleBase to CALIB mode
-        for l in uint8_model.model.layers:
-            l.enable_calibration()
-
-        # (b) run inference to collect ranges
+        # Run inference to collect ranges
         iterator = range(0, ids.size(1) - 1, args.stride)
         if not args.no_tqdm:
             iterator = tqdm.tqdm(iterator, desc="Calibration")
@@ -194,9 +184,8 @@ def main():
             for i in iterator:
                 uint8_model(ids[:, i : i + args.stride])
 
-        # (c) freeze (scale, zero-point)
-        for l in uint8_model.model.layers:
-            l.freeze_qparams()
+        # Freeze (scale, zero-point)
+        convert(uint8_model)
 
     # -------------------------------------------------------------------------
     # 4. Evaluate perplexity

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

@@ -38,13 +38,14 @@ import tqdm
 from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.utils.introspection import (
+from tico.quantization import convert, prepare
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.wrapq.utils.introspection import (
     build_fqn_map,
     compare_layer_outputs,
     save_fp_outputs,
 )
-from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
+from tico.quantization.wrapq.wrappers.ptq_wrapper import PTQWrapper
 
 # Token-budget presets for activation calibration
 TOKENS: dict[str, int] = {
@@ -176,19 +177,8 @@ def main():
     # 2. Wrap every layer with PTQWrapper (UINT-8 activations)
     # -------------------------------------------------------------------------
     print("Wrapping layers with PTQWrapper …")
-    qcfg = QuantConfig()  # default: per-tensor UINT8
-
-    new_layers = torch.nn.ModuleList()
-    for idx, fp_layer in enumerate(model.model.layers):
-        layer_cfg = qcfg.child(f"layer{idx}")
-        q_layer = PTQWrapper(
-            fp_layer,
-            qcfg=layer_cfg,
-            fp_name=m_to_fqn.get(fp_layer),
-        )
-        new_layers.append(q_layer)
-
-    model.model.layers = new_layers  # swap in quant wrappers
+    qcfg = PTQConfig()  # default: per-tensor UINT8
+    prepare(model, qcfg)
 
     # -------------------------------------------------------------------------
     # 3. Activation calibration plus FP-vs-UINT8 diffing
@@ -197,10 +187,6 @@ def main():
     calib_txt = " ".join(dataset["text"])[:CALIB_TOKENS]
     ids = tokenizer(calib_txt, return_tensors="pt").input_ids.to(device)
 
-    # (a) Enable CALIB mode on every QuantModuleBase
-    for l in model.model.layers:
-        l.enable_calibration()
-
     # Save reference FP activations before observers clamp/quantize
     save_handles, act_cache = save_fp_outputs(model)
 
@@ -216,11 +202,10 @@ def main():
     for h in save_handles:
         h.remove()
 
-    # (b) Freeze (scale, zero-point) after calibration
-    for l in model.model.layers:
-        l.freeze_qparams()
+    # Freeze (scale, zero-point) after calibration
+    convert(model)
 
-    # (c) Register diff hooks and measure per-layer deltas
+    # Register diff hooks and measure per-layer deltas
     cmp_handles = compare_layer_outputs(model, act_cache, metrics=["diff", "peir"])
     # Use same inputs for comparison.
     with torch.no_grad():

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

@@ -29,13 +29,15 @@ import pathlib
 import torch
 import torch.nn as nn
 
-from tico.experimental.quantization.evaluation.metric import compute_peir
-from tico.experimental.quantization.evaluation.utils import plot_two_outputs
-
-from tico.experimental.quantization.ptq.mode import Mode
-from tico.experimental.quantization.ptq.wrappers.nn.quant_linear import QuantLinear
+from tico.quantization import convert, prepare
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.evaluation.metric import compute_peir
+from tico.quantization.evaluation.utils import plot_two_outputs
+from tico.quantization.wrapq.mode import Mode
+from tico.quantization.wrapq.wrappers.nn.quant_linear import QuantLinear
 from tico.utils.utils import SuppressWarning
 
+
 # -------------------------------------------------------------------------
 # 0. Define a toy model (1 Linear layer only)
 # -------------------------------------------------------------------------
@@ -60,20 +62,19 @@ fp32_layer = model.fc
 # -------------------------------------------------------------------------
 # 1. Replace the Linear with QuantLinear wrapper
 # -------------------------------------------------------------------------
-model.fc = QuantLinear(fp32_layer)  # type: ignore[assignment]
-# model.fc = PTQWrapper(fp32_layer) (Wrapping helper class)
+model.fc = prepare(fp32_layer, PTQConfig())  # type: ignore[assignment]
 qlayer = model.fc  # alias for brevity
 
 # -------------------------------------------------------------------------
 # 2. Single-pass calibration (collect activation ranges)
 # -------------------------------------------------------------------------
-assert isinstance(qlayer, QuantLinear)
+assert isinstance(qlayer.wrapped, QuantLinear)
 with torch.no_grad():
-    qlayer.enable_calibration()
     for _ in range(16):  # small toy batch
         x = torch.randn(4, 16)  # (batch=4, features=16)
         _ = model(x)
-    qlayer.freeze_qparams()  # lock scales & zero-points
+
+convert(qlayer)
 
 assert qlayer._mode is Mode.QUANT, "Quantization mode should be active now."
 

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

@@ -17,13 +17,12 @@ import pathlib
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from tico.experimental.quantization.evaluation.metric import compute_peir
-from tico.experimental.quantization.evaluation.utils import plot_two_outputs
-
-from tico.experimental.quantization.ptq.mode import Mode
-from tico.experimental.quantization.ptq.wrappers.llama.quant_attn import (
-    QuantLlamaAttention,
-)
+from tico.quantization import convert, prepare
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.evaluation.metric import compute_peir
+from tico.quantization.evaluation.utils import plot_two_outputs
+from tico.quantization.wrapq.mode import Mode
+from tico.quantization.wrapq.wrappers.llama.quant_attn import QuantLlamaAttention
 from tico.utils.utils import SuppressWarning
 
 name = "Maykeye/TinyLLama-v0"
@@ -34,12 +33,11 @@ tokenizer = AutoTokenizer.from_pretrained(name)
 # 1. Replace layer-0’s MLP with QuantLlamaMLP
 # -------------------------------------------------------------------------
 orig_attn = model.model.layers[0].self_attn
-model.model.layers[0].self_attn = QuantLlamaAttention(
-    orig_attn
-)  # PTQWrapper(orig_attn) is also fine
+model.model.layers[0].self_attn = prepare(orig_attn, PTQConfig())
 model.eval()
 
 attn_q = model.model.layers[0].self_attn  # quant wrapper
+assert isinstance(attn_q.wrapped, QuantLlamaAttention)
 rotary = model.model.rotary_emb
 
 # -------------------------------------------------------------------------
@@ -55,7 +53,6 @@ PROMPTS = [
 ]
 
 with torch.no_grad():
-    attn_q.enable_calibration()
     for prompt in PROMPTS:
         ids = tokenizer(prompt, return_tensors="pt")
         embeds = model.model.embed_tokens(ids["input_ids"])
@@ -63,7 +60,8 @@ with torch.no_grad():
         S = cos_sin[0].shape[1]
         float_mask = torch.zeros(1, 1, S, S)
         _ = attn_q(embeds, cos_sin)  # observers collect
-    attn_q.freeze_qparams()
+
+convert(attn_q)
 
 assert attn_q._mode is Mode.QUANT, "Quantization mode should be active now."
 

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

@@ -31,10 +31,12 @@ import pathlib
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from tico.experimental.quantization.evaluation.metric import compute_peir
-from tico.experimental.quantization.evaluation.utils import plot_two_outputs
-from tico.experimental.quantization.ptq.mode import Mode
-from tico.experimental.quantization.ptq.wrappers.llama.quant_decoder_layer import (
+from tico.quantization import convert, prepare
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.evaluation.metric import compute_peir
+from tico.quantization.evaluation.utils import plot_two_outputs
+from tico.quantization.wrapq.mode import Mode
+from tico.quantization.wrapq.wrappers.llama.quant_decoder_layer import (
     QuantLlamaDecoderLayer,
 )
 from tico.utils.utils import SuppressWarning
@@ -50,12 +52,11 @@ rotary = model.model.rotary_emb  # RoPE helper
 # 1. Swap in the quant wrapper
 # -------------------------------------------------------------------------
 fp32_layer = model.model.layers[0]  # keep a reference for diff check
-model.model.layers[0] = QuantLlamaDecoderLayer(
-    fp32_layer
-)  # PTQWrapper(fp32_layer) is also fine
+model.model.layers[0] = prepare(fp32_layer, PTQConfig())
 model.eval()
 
 qlayer = model.model.layers[0]  # alias for brevity
+assert isinstance(qlayer.wrapped, QuantLlamaDecoderLayer)
 
 # -------------------------------------------------------------------------
 # 2. Single-pass calibration (gather activation ranges)
@@ -70,7 +71,6 @@ PROMPTS = [
 ]
 
 with torch.no_grad():
-    qlayer.enable_calibration()
     for prompt in PROMPTS:
         ids = tokenizer(prompt, return_tensors="pt")
         hidden = model.model.embed_tokens(ids["input_ids"])
@@ -78,7 +78,8 @@ with torch.no_grad():
         S = pos[0].shape[1]
         attn_mask = torch.zeros(1, 1, S, S)  # causal-mask placeholder
         _ = qlayer(hidden, attention_mask=attn_mask, position_embeddings=pos)
-    qlayer.freeze_qparams()
+
+convert(qlayer)
 
 assert qlayer._mode is Mode.QUANT, "Quantization mode should be active now."
 

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

@@ -18,13 +18,14 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 import tico
-from tico.experimental.quantization.evaluation.metric import compute_peir
-from tico.experimental.quantization.evaluation.utils import plot_two_outputs
-from tico.experimental.quantization.ptq.dtypes import INT16
-from tico.experimental.quantization.ptq.mode import Mode
-from tico.experimental.quantization.ptq.qscheme import QScheme
-from tico.experimental.quantization.ptq.quant_config import QuantConfig
-from tico.experimental.quantization.ptq.wrappers.llama.quant_mlp import QuantLlamaMLP
+from tico.quantization import convert, prepare
+from tico.quantization.config.ptq import PTQConfig
+from tico.quantization.evaluation.metric import compute_peir
+from tico.quantization.evaluation.utils import plot_two_outputs
+from tico.quantization.wrapq.dtypes import INT16
+from tico.quantization.wrapq.mode import Mode
+from tico.quantization.wrapq.qscheme import QScheme
+from tico.quantization.wrapq.wrappers.llama.quant_mlp import QuantLlamaMLP
 from tico.utils.utils import SuppressWarning
 
 name = "Maykeye/TinyLLama-v0"
@@ -36,13 +37,13 @@ model.eval()
 # 1. Replace layer-0’s MLP with QuantLlamaMLP
 # -------------------------------------------------------------------------
 fp32_mlp = model.model.layers[0].mlp
-model.model.layers[0].mlp = QuantLlamaMLP(
-    fp32_mlp,
-    qcfg=QuantConfig(default_dtype=INT16, default_qscheme=QScheme.PER_TENSOR_SYMM),
-)  # PTQWrapper(fp32_mlp) is also fine
+model.model.layers[0].mlp = prepare(
+    fp32_mlp, PTQConfig(default_dtype=INT16, default_qscheme=QScheme.PER_TENSOR_SYMM)
+)
 model.eval()
 
 mlp_q = model.model.layers[0].mlp
+assert isinstance(mlp_q.wrapped, QuantLlamaMLP)
 
 # -------------------------------------------------------------------------
 # 2. Single-pass calibration
@@ -57,13 +58,12 @@ PROMPTS = [
 ]
 
 with torch.no_grad():
-    mlp_q.enable_calibration()
     for prompt in PROMPTS:
         enc = tokenizer(prompt, return_tensors="pt")
         emb = model.model.embed_tokens(enc["input_ids"])
         _ = mlp_q(emb)
 
-    mlp_q.freeze_qparams()
+convert(mlp_q)
 
 assert mlp_q._mode is Mode.QUANT, "Quantization mode should be active now."