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

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

@@ -17,15 +17,17 @@ import types
 from typing import Any, Callable, Dict, List, Optional
 
 import torch
+from tqdm.auto import tqdm
 
-from tico.experimental.quantization.algorithm.gptq.gptq import GPTQ
-from tico.experimental.quantization.algorithm.gptq.utils import (
+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.experimental.quantization.config import BaseConfig, GPTQConfig
-from tico.experimental.quantization.quantizer import BaseQuantizer
+from tico.quantization.config.gptq import GPTQConfig
+from tico.quantization.quantizer import BaseQuantizer
+from tico.quantization.quantizer_registry import register_quantizer
 
 
 class StopForward(Exception):
@@ -34,6 +36,7 @@ class StopForward(Exception):
     pass
 
 
+@register_quantizer(GPTQConfig)
 class GPTQQuantizer(BaseQuantizer):
     """
     Quantizer for applying the GPTQ algorithm (typically for weight quantization).
@@ -43,7 +46,7 @@ class GPTQQuantizer(BaseQuantizer):
         3) convert(model) to consume the collected data and apply GPTQ.
     """
 
-    def __init__(self, config: BaseConfig):
+    def __init__(self, config: GPTQConfig):
         super().__init__(config)
 
         # cache_args[i] -> list of the i-th positional argument for each batch
@@ -181,7 +184,14 @@ class GPTQQuantizer(BaseQuantizer):
             target_layers = [model]
 
         quantizers: Dict[str, Any] = {}
-        for l_idx, layer in enumerate(target_layers):
+        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())]
@@ -210,7 +220,13 @@ class GPTQQuantizer(BaseQuantizer):
 
                 # Run layer forward over all cached batches to build Hessian/statistics
                 batch_num = self.num_batches
-                for batch_idx in range(batch_num):
+                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
                     )
@@ -238,7 +254,13 @@ class GPTQQuantizer(BaseQuantizer):
                     gptq[name].free()
 
             # 4) After quantization, re-run the layer to produce outputs for the next layer
-            for batch_idx in range(batch_num):
+            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
                 )

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

@@ -25,14 +25,12 @@ from torch.ao.quantization.observer import MinMaxObserver, PerChannelMinMaxObser
 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

@@ -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 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

@@ -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/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:

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

@@ -16,20 +16,37 @@ from typing import Any, Dict, Optional
 
 import torch
 
-from tico.experimental.quantization.algorithm.smoothquant.observer import (
-    ChannelwiseMaxActsObserver,
-)
+from tico.quantization.algorithm.smoothquant.observer import ChannelwiseMaxActsObserver
 
-from tico.experimental.quantization.algorithm.smoothquant.smooth_quant import (
-    apply_smoothing,
-)
-from tico.experimental.quantization.config import SmoothQuantConfig
-from tico.experimental.quantization.quantizer import BaseQuantizer
+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)
 class SmoothQuantQuantizer(BaseQuantizer):
     """
     Quantizer for applying the SmoothQuant algorithm
+
+    Q) Why allow choosing between input and output activations?
+
+    SmoothQuant relies on channel-wise activation statistics to balance
+    weights and activations. In practice, there are two natural sources:
+
+    - "input": captures the tensor right before a Linear layer
+      (forward-pre-hook). This matches the original SmoothQuant paper
+      and focuses on scaling the raw hidden state.
+
+    - "output": captures the tensor right after a Linear layer
+      (forward-hook). This can better reflect post-weight dynamics,
+      especially when subsequent operations (bias, activation functions)
+      dominate the dynamic range.
+
+    Allowing both options provides flexibility: depending on model
+    architecture and calibration data, one may yield lower error than
+    the other. The default remains "input" for compatibility, but "output"
+    can be selected to empirically reduce error or runtime overhead.
     """
 
     def __init__(self, config: SmoothQuantConfig):
@@ -37,6 +54,7 @@ class SmoothQuantQuantizer(BaseQuantizer):
 
         self.alpha = config.alpha
         self.custom_alpha_map = config.custom_alpha_map
+        self.acts_from = config.acts_from  # "input" (default) or "output"
         self.observer: Optional[ChannelwiseMaxActsObserver] = None
 
     @torch.no_grad()
@@ -55,7 +73,8 @@ class SmoothQuantQuantizer(BaseQuantizer):
         Returns:
             The model prepared for SmoothQuant quantization.
         """
-        self.observer = ChannelwiseMaxActsObserver(model)
+        # Attach hooks according to `config.acts_from`
+        self.observer = ChannelwiseMaxActsObserver(model, acts_from=self.acts_from)
         self.observer.attach()
 
         return model