tico 0.1.0.dev250730__py3-none-any.whl → 0.1.0.dev250731__py3-none-any.whl

tico/__init__.py

@@ -29,7 +29,7 @@ __all__ = [
 ]
 
 # THIS LINE IS AUTOMATICALLY GENERATED BY setup.py
-__version__ = "0.1.0.dev250730"
+__version__ = "0.1.0.dev250731"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

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

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 import types
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 import torch
 
@@ -28,16 +28,34 @@ from tico.experimental.quantization.config import BaseConfig, GPTQConfig
 from tico.experimental.quantization.quantizer import BaseQuantizer
 
 
+class StopForward(Exception):
+    """Custom exception used to stop the forward pass after the first layer."""
+
+    pass
+
+
 class GPTQQuantizer(BaseQuantizer):
     """
-    Quantizer for applying the GPTQ algorithm (typically for weight quantization)
+    Quantizer for applying the GPTQ algorithm (typically for weight quantization).
+    This implementation expects:
+        1) prepare(model, ...) to only attach hooks/Catchers and NOT run the model internally.
+        2) The user runs the model with arbitrary number of batches to collect calibration data.
+        3) convert(model) to consume the collected data and apply GPTQ.
     """
 
     def __init__(self, config: BaseConfig):
         super().__init__(config)
 
-        self.cache_args: List[Any] = []
-        self.cache_kwargs: Dict[str, Any] = {"batch_num": 0}
+        # cache_args[i] -> list of the i-th positional argument for each batch
+        self.cache_args: List[List[Any]] = []
+        # cache_kwargs[k] -> list of the value for keyword k for each batch
+        self.cache_kwargs: Dict[str, List[Any]] = {}
+        self.num_batches: int = 0
+
+        # References to original forwards for restoration
+        self._orig_model_forward: Optional[Callable[..., Any]] = None
+        self._orig_layer_forward: Optional[Callable[..., Any]] = None
+        self._first_layer_ref: Optional[torch.nn.Module] = None
 
     @torch.no_grad()
     def prepare(
@@ -50,104 +68,125 @@ class GPTQQuantizer(BaseQuantizer):
         Overrides the forward method of the first LLaMA layer (layer 0) to capture the
          input required for calibration.
 
-        This method modifies the original forward pass of LLaMA layer 0 so that the
-         inputs used during inference are intercepted and recorded. These captured inputs
-        are then utilized to calibrate the quantization parameters for the GPTQ.
+        When the user calls `model(...)`, we intercept (and store) the inputs to that
+         layer, then raise an exception to stop the forward pass immediately. These
+        captured inputs are then utilized to calibrate the quantization parameters
+         for the GPTQ.
 
         Parameters:
-            model: The target PyTorch model.
-            args: Positional example inputs required for capturing graph.
-            kwargs: Keyword example inputs required for capturing graph.
+            model (torch.nn.Module): The target PyTorch model
+            args (Any, optional): Unused (kept for API compatibility)
+            kwargs (Dict[str, Any], optional): Unused (kept for API compatibility)
 
         Returns:
-            The model prepared for GPTQ quantization.
+            torch.nn.Module: The model with the catcher attached
         """
-        if args is None and kwargs is None:
-            raise RuntimeError(
-                "Either args or kwargs must be provided for captruing graph."
-            )
-        # Define a function to capture input activations and associated parameters.
+        # Define the catcher to store inputs/kwargs and stop the execution
         def forward(layer, *args, **kwargs):
-            self.cache_kwargs["batch_num"] += 1
+            """
+            Stores this batch's inputs and kwargs, then raises StopForward to stop computation.
+            """
+            # Store positional args
             for idx, item in enumerate(args):
                 if (idx + 1) > len(self.cache_args):
                     self.cache_args.append([])
                 self.cache_args[idx].append(item)
-            for arg in kwargs:
-                if self.cache_kwargs.get(arg, None) is None:
-                    self.cache_kwargs[arg] = []
-                self.cache_kwargs[arg].append(kwargs[arg])
-            # Raise an error to interrupt the forward pass after capturing data.
-            raise ValueError
+            # Store keyword args
+            for k, v in kwargs.items():
+                if self.cache_kwargs.get(k, None) is None:
+                    self.cache_kwargs[k] = []
+                self.cache_kwargs[k].append(v)
+
+            self.num_batches += 1
+            raise StopForward  # stop after the first layer
 
         # Replace the first layer with defined function to capture calibration data.
         if hasattr(model, "model"):
-            assert hasattr(model.model, "layers")
-            assert isinstance(model.model.layers, torch.nn.ModuleList)
-            layer_forward_cache = model.model.layers[0].forward
-            model.model.layers[0].forward = types.MethodType(
-                forward, model.model.layers[0]
-            )
+            if hasattr(model.model, "layers") and isinstance(
+                model.model.layers, torch.nn.ModuleList
+            ):
+                self._first_layer_ref = model.model.layers[0]
+            else:
+                raise RuntimeError(
+                    "GPTQ Quantizer assumes the model has a nested structure like `model.model.layers`, commonly found in LLaMA and other Hugging Face transformer models."
+                )
         else:
-            assert hasattr(model, "forward")
-            layer_forward_cache = model.forward
-            model.forward = types.MethodType(forward, model.forward)
-
-        model_forward_cache = model.forward
-        # Replace model's forward to avoid ValueError
-        def model_forward(model, *args, **kwargs):
-            nonlocal model_forward_cache
-            try:
-                model_forward_cache(*args, **kwargs)
-            except ValueError:
-                pass
+            # fallback if the model is not LLaMA-like; treat whole model as single layer
+            self._first_layer_ref = model
 
-        model.forward = types.MethodType(model_forward, model)
-        kwargs = kwargs or {}
-        model(*args, **kwargs)  # type: ignore[misc]
+        assert hasattr(self._first_layer_ref, "forward")
+        # Backup the original forward of the first layer
+        assert isinstance(self._first_layer_ref, torch.nn.Module)
+        self._orig_layer_forward = self._first_layer_ref.forward
+        self._first_layer_ref.forward = types.MethodType(forward, self._first_layer_ref)
 
-        # Recover original forward
-        model.forward = model_forward_cache
-        if hasattr(model, "model"):
-            assert hasattr(model.model, "layers")
-            assert isinstance(model.model.layers, torch.nn.ModuleList)
-            model.model.layers[0].forward = layer_forward_cache
-        else:
-            model.forward = layer_forward_cache
+        def model_forward_wrapper(_model, *m_args, **m_kwargs):
+            """
+            Wrapper to ignore StopForward exceptions so the user's training loop doesn't crash.
+            """
+            try:
+                assert self._orig_model_forward is not None
+                return self._orig_model_forward(*m_args, **m_kwargs)
+            except StopForward:
+                # We stopped after the first layer; return None or dummy output if needed.
+                return None
+
+        # Backup model.forward so we can suppress StopForward
+        self._orig_model_forward = model.forward
+        model.forward = types.MethodType(model_forward_wrapper, model)
 
         return model
 
     @torch.no_grad()
     def convert(self, model):
         """
-        Convert the prepared model to its GPTQ quantized version.
+        Perform GPTQ quantization using cached first-layer inputs.
 
-        Applies the GPTQ quantization on weights based on the collected statistics.
+        Steps:
+          1) Restore original forwards (no more catching).
+          2) Iterate through each Transformer layer sequentially:
+             a) For each layer, register forward hooks to collect (inp, out) stats for GPTQ.
+             b) Run the layer on cached inputs for all batches.
+             c) Apply GPTQ and update the weights.
+             d) Re-run the layer to produce outputs for the next layer; update cached inputs.
+          3) Restore model.config.use_cache if needed and clear internal caches.
 
         Parameters:
-            model: The prepared PyTorch model.
+            model (torch.nn.Module): The prepared model.
 
         Returns:
-            The quantized model.
+            torch.nn.Module: Quantized model.
         """
+        # Restore original forwards (we no longer want to stop after first layer)
+        assert self._orig_model_forward is not None
+        model.forward = self._orig_model_forward
+        assert (
+            self._first_layer_ref is not None and self._orig_layer_forward is not None
+        )
+        self._first_layer_ref.forward = self._orig_layer_forward
+
         gptq_conf = self.config
         assert isinstance(gptq_conf, GPTQConfig)
-
-        # Save the original cache setting and disable caching during calibration/inference.
-        if hasattr(model, "config"):
-            use_cache = model.config.use_cache
+        # Disable use_cache during calibration
+        if hasattr(model, "config") and hasattr(model.config, "use_cache"):
+            orig_use_cache = model.config.use_cache
             model.config.use_cache = False
+        else:
+            orig_use_cache = None
 
-        quantizers = {}
+        # Identify layers
         if hasattr(model, "model"):
             target_layers = model.model.layers
         else:
             target_layers = [model]
+
+        quantizers: Dict[str, Any] = {}
         for l_idx, layer in enumerate(target_layers):
-            # Identify quantizable submodules within the layer.
+            # 1) Identify quantizable submodules within the layer
             full = find_layers(layer)
-
             sequential = [list(full.keys())]
+
+            # 2) Set up GPTQ objects and gather stats
             for names in sequential:
                 subset = {n: full[n] for n in names}
 
@@ -155,20 +194,22 @@ class GPTQQuantizer(BaseQuantizer):
                 for name in subset:
                     gptq[name] = GPTQ(subset[name])
                     gptq[name].quantizer.configure(
-                        8, perchannel=True, sym=False, mse=False
+                        bits=8, perchannel=True, sym=False, mse=False
                     )
-                # Define a hook to collect input/output batches for quantizer calibration.
+
+                # Hook to collect (inp, out) for GPTQ
                 def add_batch(name):
-                    def tmp(_, inp, out):
+                    def _hook(_, inp, out):
                         gptq[name].add_batch(inp[0].data, out.data)
 
-                    return tmp
+                    return _hook
 
                 handles = []
                 for name in subset:
                     handles.append(subset[name].register_forward_hook(add_batch(name)))
-                # Run the current layer on the stored calibration inputs to capture activation stats.
-                batch_num = self.cache_kwargs.pop("batch_num")
+
+                # Run layer forward over all cached batches to build Hessian/statistics
+                batch_num = self.num_batches
                 for batch_idx in range(batch_num):
                     cache_args_batch = gather_single_batch_from_list(
                         self.cache_args, batch_idx
@@ -176,15 +217,16 @@ class GPTQQuantizer(BaseQuantizer):
                     cache_kwargs_batch = gather_single_batch_from_dict(
                         self.cache_kwargs, batch_idx
                     )
-                    layer(*cache_args_batch, **cache_kwargs_batch)[0]
-                self.cache_kwargs["batch_num"] = batch_num
+                    layer(*cache_args_batch, **cache_kwargs_batch)
+
+                # Remove handles
                 for h in handles:
                     h.remove()
-                # Quantize each submodule using the collected calibration data.
+
+                # 3) Quantize each submodule
                 for name in subset:
                     if gptq_conf.verbose:
-                        print(l_idx, name)
-                        print("Quantizing ...")
+                        print(f"[Layer {l_idx}] {name} -> Quantizing ...")
                     gptq[name].fasterquant(
                         percdamp=0.01,
                         groupsize=-1,
@@ -192,18 +234,10 @@ class GPTQQuantizer(BaseQuantizer):
                         static_groups=False,
                         verbose=gptq_conf.verbose,
                     )
-                    quantizers["model.layers.%d.%s" % (l_idx, name)] = gptq[
-                        name
-                    ].quantizer
+                    quantizers[f"model.layers.{l_idx}.{name}"] = gptq[name].quantizer
                     gptq[name].free()
-            """
-            Execute the quantized layer with the calibration inputs to obtain ouptuts
-             that will serve as inputs for the next layer.
 
-            This ensures that the quantization effects are correctly propagated to subsequent
-             layers.
-            """
-            batch_num = self.cache_kwargs.pop("batch_num")
+            # 4) After quantization, re-run the layer to produce outputs for the next layer
             for batch_idx in range(batch_num):
                 cache_args_batch = gather_single_batch_from_list(
                     self.cache_args, batch_idx
@@ -211,15 +245,26 @@ class GPTQQuantizer(BaseQuantizer):
                 cache_kwargs_batch = gather_single_batch_from_dict(
                     self.cache_kwargs, batch_idx
                 )
-                outs = layer(*cache_args_batch, **cache_kwargs_batch)[0]
+                outs = layer(*cache_args_batch, **cache_kwargs_batch)
+                # LLaMA's decoder layer return type differs across Transformers versions:
+                # some return a tuple (hidden_states, ...), others return just a tensor.
+                # This line ensures we always take the first element when it's a tuple.
+                outs = outs[0] if isinstance(outs, tuple) else outs
                 # Update inputs for next iteration.
                 self.cache_args[0][batch_idx] = outs
-            self.cache_kwargs["batch_num"] = batch_num
 
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
+
         # Restore the original cache configuration.
-        if hasattr(model, "config"):
-            model.config.use_cache = use_cache
+        if orig_use_cache is not None:
+            model.config.use_cache = orig_use_cache
+
+        # Clear caches to free memory
+        self.cache_args.clear()
+        self.cache_kwargs.clear()
+        self.num_batches = 0
+
+        model.quantizers = quantizers
 
         return model

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

@@ -58,7 +58,7 @@ def gather_single_batch_from_list(data_list, idx):
     Returns:
         list: single batch.
     """
-    # obtain a set of keyword input from cache
+    # obtain a set of positional input from cache
     single_batch = []
     for data_item in data_list:
         single_batch.append(data_item[idx])

tico/experimental/quantization/public_interface.py

@@ -59,6 +59,8 @@ def prepare(
     Returns:
         The model prepared for quantization.
     """
+    if hasattr(model, QUANTIZER_ATTRIBUTE_NAME):
+        raise RuntimeError("prepare() already has been called.")
     if quant_config.name == "pt2e" and inplace:
         raise RuntimeError(
             "In-place is not supported for PT2E quantization due to limitation in the underlying Torch APIs. Please set 'inplace=False' to proceed."
@@ -99,6 +101,12 @@ def convert(model, inplace: Optional[bool] = False):
         raise RuntimeError(
             "In-place is not supported for PT2E quantization due to limitation in the underlying Torch APIs. Please set 'inplace=False' to proceed."
         )
+    # deepcopy prevents the quantizer from restoring the catcher used for calibration.
+    # TODO Revisit `inplace` policy.
+    if isinstance(quantizer, GPTQQuantizer) and not inplace:
+        raise RuntimeError(
+            "GPTQ quantization only supports `in-place=True`. Please set 'inplace=True' to proceed."
+        )
 
     model = model if inplace else copy.deepcopy(model)
 

tico/utils/record_input.py

@@ -15,7 +15,6 @@
 import copy
 
 import inspect
-from contextlib import contextmanager
 from typing import Callable, List, Optional
 
 import torch.nn as nn

{tico-0.1.0.dev250730.dist-info → tico-0.1.0.dev250731.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tico
-Version: 0.1.0.dev250730
+Version: 0.1.0.dev250731
 Summary: Convert exported Torch module to circle
 Home-page: UNKNOWN
 License: UNKNOWN

{tico-0.1.0.dev250730.dist-info → tico-0.1.0.dev250731.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=azezzGizSQu7vHuaNsna2JC-5HPI5SV5QI_q-6STx14,1883
+tico/__init__.py,sha256=DnYmoCoxFZl3z2PPb_zOksdK5E8E2JAGnARIST2A4n0,1883
 tico/pt2_to_circle.py,sha256=gu3MD4Iqc0zMZcCZ2IT8oGbyj21CTSbT3Rgd9s2B_9A,2767
 tico/config/__init__.py,sha256=xZzCXjZ84qE-CsBi-dfaL05bqpQ3stKKfTXhnrJRyVs,142
 tico/config/base.py,sha256=q5xMqGxTUZs4mFqt5c7i_y9U00fYgdMGl9nUqIVMlCo,1248
@@ -7,14 +7,14 @@ tico/config/v1.py,sha256=O1jzpUBDwoWpLohEpI08pJNwVB-yz3ufPrQm2_XWq4Y,1108
 tico/experimental/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/experimental/quantization/__init__.py,sha256=IaJPZegVJp0P3luutBo907Kp5sOJensE1Mm-XBG_jBs,122
 tico/experimental/quantization/config.py,sha256=h01WpP8Y-dLj6yg12pMZm3PXJqUnU2sWip5jBRc5x9Q,1604
-tico/experimental/quantization/public_interface.py,sha256=OKW8UoBMjPwiTacrWgQY9ENCh8ucPnYMSrl2R-w0pJ0,3982
+tico/experimental/quantization/public_interface.py,sha256=4-v9VXsokRG2-UUYYHd_MlbHxChqdGI5iuySyYDY_Pw,4420
 tico/experimental/quantization/quantizer.py,sha256=_2pDtWFKDCuKfYF2bptOwIYsa0VFNFM1ZNgi8_OGvHM,2365
 tico/experimental/quantization/algorithm/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/experimental/quantization/algorithm/gptq/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/experimental/quantization/algorithm/gptq/gptq.py,sha256=Qn9b_2ki7B64DcVEY25NMkww3PdZ5EqYQQXfYhNDQ6I,5555
 tico/experimental/quantization/algorithm/gptq/quant.py,sha256=Rl4wAOCmlE0U09BtNCDbccaSNohRHCNLwFi3zCqZfNo,5127
-tico/experimental/quantization/algorithm/gptq/quantizer.py,sha256=icaFDXA1UibgRI0nBZ4N0Ij1ajVpShWUFw5pTDffOiE,8914
-tico/experimental/quantization/algorithm/gptq/utils.py,sha256=vDIW5ow5c1VSFpub7QumMWorHrV86c0kOtlBxMw2Y2Y,1808
+tico/experimental/quantization/algorithm/gptq/quantizer.py,sha256=KiaNcDkufbYPHdkkOGw9nAwLtk0yYwUDbyzFT3xRLOs,11066
+tico/experimental/quantization/algorithm/gptq/utils.py,sha256=leGKayf-xbSjVwwAGTA5RsxUKrhDiklOQdlsLifjdrs,1811
 tico/experimental/quantization/algorithm/pt2e/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/experimental/quantization/algorithm/pt2e/quantizer.py,sha256=mdTvsG87bo8fu0GaWqSM8iBCs-4f4EfUlVtk-Ko6M34,2546
 tico/experimental/quantization/algorithm/pt2e/utils.py,sha256=URjTGgsnDdhUC2Nr0-YJ9GWbVOKmjElfLr83Y8eCz-M,4806
@@ -195,7 +195,7 @@ tico/utils/model.py,sha256=pPOIjD0qjQirLibiRxxfjOR6efimOcDAd9R-74eus-k,1282
 tico/utils/padding.py,sha256=qKke-dJeeLHiRaePjDS66txrGyiYuipLVQeqLYad8uk,3349
 tico/utils/passes.py,sha256=kGmDe__5cPaO6i5EDAoXSVe6yXEoX9hAny4ROb3ZEmQ,2409
 tico/utils/pytree_utils.py,sha256=jrk3N6X6LiUnBCX_gM1K9nywbVAJBVnszlTAgeIeDUc,5219
-tico/utils/record_input.py,sha256=AHDDxuAjFP2Jgl9CsmQuwjivr-lAZxek1rO18CjK6IY,3932
+tico/utils/record_input.py,sha256=QN-8D71G_WAX3QQQ5CIwbEfFJZTQ3CvL4wCMiVddua4,3894
 tico/utils/register_custom_op.py,sha256=3-Yl6iYmx1qQA2igNHt4hYhQhQMkdPb7gF50LIY8yvc,27350
 tico/utils/serialize.py,sha256=mEuusEzi82WFsz3AkowgWwxSLeo50JDxyOj6yYDQhEI,1914
 tico/utils/torch_compat.py,sha256=oc6PztVsXdHcQ3iaVR90wLLxrGaj6zFHWZ8K9rRS6q8,1795
@@ -206,9 +206,9 @@ tico/utils/mx/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/utils/mx/elemwise_ops.py,sha256=V6glyAHsVR1joqpsgnNytatCD_ew92xNWZ19UFDoMTA,10281
 tico/utils/mx/formats.py,sha256=uzNWyu-1onUlwQfX5cZ6fZSUfHMRqorper7_T1k3jfk,3404
 tico/utils/mx/mx_ops.py,sha256=RcfUTYVi-wilGB2sC35OeARdwDqnixv7dG5iyZ-fQT8,8555
-tico-0.1.0.dev250730.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
-tico-0.1.0.dev250730.dist-info/METADATA,sha256=2bGJV9nHJLnjsv3q2PB6pLR-uNgYal0dvsGpuaIwS9Y,8430
-tico-0.1.0.dev250730.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250730.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250730.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250730.dist-info/RECORD,,
+tico-0.1.0.dev250731.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
+tico-0.1.0.dev250731.dist-info/METADATA,sha256=oCmTMbqj-MxLLJHDtnZkmdHWT4C_g796A3bgpNEFxUA,8430
+tico-0.1.0.dev250731.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+tico-0.1.0.dev250731.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
+tico-0.1.0.dev250731.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
+tico-0.1.0.dev250731.dist-info/RECORD,,