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

tico/quantization/algorithm/gptq/quantizer.py

@@ -0,0 +1,292 @@
+# Copyright (c) 2024 Intel Corporation
+# 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 types
+from typing import Any, Callable, Dict, List, Optional
+
+import torch
+from tqdm.auto import tqdm
+
+from tico.quantization.algorithm.gptq.gptq import GPTQ
+from tico.quantization.algorithm.gptq.utils import (
+    find_layers,
+    gather_single_batch_from_dict,
+    gather_single_batch_from_list,
+)
+from tico.quantization.config.gptq import GPTQConfig
+from tico.quantization.quantizer import BaseQuantizer
+from tico.quantization.quantizer_registry import register_quantizer
+
+
+class StopForward(Exception):
+    """Custom exception used to stop the forward pass after the first layer."""
+
+    pass
+
+
+@register_quantizer(GPTQConfig)
+class GPTQQuantizer(BaseQuantizer):
+    """
+    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: GPTQConfig):
+        super().__init__(config)
+
+        # 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(
+        self,
+        model: torch.nn.Module,
+        args: Optional[Any] = None,
+        kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Overrides the forward method of the first LLaMA layer (layer 0) to capture the
+         input required for calibration.
+
+        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 (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:
+            torch.nn.Module: The model with the catcher attached
+        """
+        # Define the catcher to store inputs/kwargs and stop the execution
+        def forward(layer, *args, **kwargs):
+            """
+            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)
+            # 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"):
+            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:
+            # fallback if the model is not LLaMA-like; treat whole model as single layer
+            self._first_layer_ref = model
+
+        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)
+
+        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):
+        """
+        Perform GPTQ quantization using cached first-layer inputs.
+
+        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 (torch.nn.Module): The prepared model.
+
+        Returns:
+            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)
+        # 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
+
+        # 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(
+            tqdm(
+                target_layers,
+                desc="Quantizing layers",
+                unit="layer",
+                disable=not gptq_conf.show_progress,
+            )
+        ):
+            # 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}
+
+                gptq: Dict[str, GPTQ] = {}
+                for name in subset:
+                    gptq[name] = GPTQ(subset[name])
+                    gptq[name].quantizer.configure(
+                        bits=8, perchannel=True, sym=False, mse=False
+                    )
+
+                # Hook to collect (inp, out) for GPTQ
+                def add_batch(name):
+                    def _hook(_, inp, out):
+                        gptq[name].add_batch(inp[0].data, out.data)
+
+                    return _hook
+
+                handles = []
+                for name in subset:
+                    handles.append(subset[name].register_forward_hook(add_batch(name)))
+
+                # Run layer forward over all cached batches to build Hessian/statistics
+                batch_num = self.num_batches
+                for batch_idx in tqdm(
+                    range(batch_num),
+                    desc=f"[L{l_idx}] collecting",
+                    leave=False,
+                    unit="batch",
+                    disable=not gptq_conf.show_progress,
+                ):
+                    cache_args_batch = gather_single_batch_from_list(
+                        self.cache_args, batch_idx
+                    )
+                    cache_kwargs_batch = gather_single_batch_from_dict(
+                        self.cache_kwargs, batch_idx
+                    )
+                    layer(*cache_args_batch, **cache_kwargs_batch)
+
+                # Remove handles
+                for h in handles:
+                    h.remove()
+
+                # 3) Quantize each submodule
+                for name in subset:
+                    if gptq_conf.verbose:
+                        print(f"[Layer {l_idx}] {name} -> Quantizing ...")
+                    gptq[name].fasterquant(
+                        percdamp=0.01,
+                        groupsize=-1,
+                        actorder=True,
+                        static_groups=False,
+                        verbose=gptq_conf.verbose,
+                    )
+                    quantizers[f"model.layers.{l_idx}.{name}"] = gptq[name].quantizer
+                    gptq[name].free()
+
+            # 4) After quantization, re-run the layer to produce outputs for the next layer
+            for batch_idx in tqdm(
+                range(batch_num),
+                desc=f"[L{l_idx}] re-forward",
+                leave=False,
+                unit="batch",
+                disable=not gptq_conf.show_progress,
+            ):
+                cache_args_batch = gather_single_batch_from_list(
+                    self.cache_args, batch_idx
+                )
+                cache_kwargs_batch = gather_single_batch_from_dict(
+                    self.cache_kwargs, batch_idx
+                )
+                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
+
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+        # Restore the original cache configuration.
+        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 → 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 → quantization}/algorithm/pt2e/annotation/annotator.py

@@ -21,23 +21,16 @@ if TYPE_CHECKING:
     import torch.fx
     from torch.ao.quantization.qconfig import _ObserverOrFakeQuantizeConstructor
 import torch
-from torch.ao.quantization.observer import (
-    MinMaxObserver,
-    MovingAverageMinMaxObserver,
-    MovingAveragePerChannelMinMaxObserver,
-    PerChannelMinMaxObserver,
-)
+from torch.ao.quantization.observer import MinMaxObserver, PerChannelMinMaxObserver
 from torch.ao.quantization.quantizer import QuantizationSpec, Quantizer
 from torch.ao.quantization.quantizer.utils import _get_module_name_filter
 
-from tico.experimental.quantization.algorithm.pt2e.annotation.op import *
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
-from tico.experimental.quantization.algorithm.pt2e.transformation.convert_scalars_to_attrs import (
+from tico.quantization.algorithm.pt2e.annotation.op import *
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
+from tico.quantization.algorithm.pt2e.transformation.convert_scalars_to_attrs import (
     convert_scalars_to_attrs,
 )
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/adaptive_avg_pool2d.py

@@ -19,12 +19,10 @@ if TYPE_CHECKING:
 import torch
 from torch.ao.quantization.quantizer import SharedQuantizationSpec
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import AdaptiveAvgPool2dArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/add.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import AddTensorArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/conv2d.py

@@ -19,12 +19,10 @@ if TYPE_CHECKING:
 import torch
 from torch.ao.quantization.quantizer import DerivedQuantizationSpec
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import Conv2DArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/div.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import DivTensorArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/linear.py

@@ -12,19 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, List, Optional, TYPE_CHECKING
+from typing import Callable, Optional, TYPE_CHECKING
 
 if TYPE_CHECKING:
     import torch.fx
 import torch
 from torch.ao.quantization.quantizer import DerivedQuantizationSpec
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import LinearArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/mean.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import MeanDimArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/mul.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import MulTensorArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/relu6.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import Relu6Args
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/rsqrt.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import RsqrtArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/sub.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import SubTensorArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/spec.py

@@ -18,9 +18,7 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 
 AnnotatorType = Callable[
     [

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

@@ -22,7 +22,7 @@ from torch.ao.quantization.quantizer import (
     SharedQuantizationSpec,
 )
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
 
 
 def annotate_input_qspec_map(node: torch.fx.Node, input_node: torch.fx.Node, qspec):

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

@@ -18,13 +18,16 @@ 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.quantizer import BaseQuantizer
+from tico.quantization.config.pt2e import PT2EConfig
+from tico.quantization.quantizer import BaseQuantizer
+from tico.quantization.quantizer_registry import register_quantizer
 
 
+@register_quantizer(PT2EConfig)
 class PT2EQuantizer(BaseQuantizer):
     """
     Quantizer for applying pytorch 2.0 export quantization (typically for activation quantization).

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

@@ -19,11 +19,8 @@ if TYPE_CHECKING:
 import torch
 from torch.ao.quantization.quantizer import QuantizationSpec
 from torch.ao.quantization.quantizer.utils import _get_module_name_filter
-from torch.utils import _pytree as pytree
 
-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/observer.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import functools
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Literal
 
 import torch
 
@@ -21,18 +21,24 @@ import torch
 class ChannelwiseMaxActsObserver:
     """
     Observer to calcuate channelwise maximum activation
+    It supports collecting activations from either module inputs or outputs.
     """
 
-    def __init__(self, model):
+    def __init__(
+        self, model: torch.nn.Module, acts_from: Literal["input", "output"] = "input"
+    ):
         """
         model
             A torch module whose activations are to be analyzed.
+        acts_from
+            Where to hook: "input" for forward-pre-hook, "output" for forward-hook.
         hooks
-            A list to store the hooks which are registered to collect activation statistics.
+            A list to store the hooks registered to collect activation statistics.
         max_acts
-            A dictionary to store the maximum activation values
+            A dictionary to store the per-channel maxima.
         """
         self.model = model
+        self.acts_from: Literal["input", "output"] = acts_from
         self.hooks: List[Any] = []
         self.max_acts: Dict[str, torch.Tensor] = {}
 
@@ -62,13 +68,25 @@ class ChannelwiseMaxActsObserver:
                 input = input[0]
             stat_tensor(name, input)
 
+        def stat_output_hook(m, input, output, name):
+            if isinstance(output, tuple):
+                output = output[0]
+            stat_tensor(name, output)
+
         for name, m in self.model.named_modules():
             if isinstance(m, torch.nn.Linear):
-                self.hooks.append(
-                    m.register_forward_pre_hook(
-                        functools.partial(stat_input_hook, name=name)
+                if self.acts_from == "input":
+                    self.hooks.append(
+                        m.register_forward_pre_hook(
+                            functools.partial(stat_input_hook, name=name)
+                        )
+                    )
+                else:  # "output"
+                    self.hooks.append(
+                        m.register_forward_hook(
+                            functools.partial(stat_output_hook, name=name)
+                        )
                     )
-                )
 
     def remove(self):
         for hook in self.hooks: