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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of tico might be problematic. Click here for more details.

Files changed (133) hide show
  1. tico/__init__.py +1 -1
  2. tico/config/v1.py +5 -0
  3. tico/passes/cast_mixed_type_args.py +2 -0
  4. tico/passes/convert_expand_to_slice_cat.py +153 -0
  5. tico/passes/convert_matmul_to_linear.py +312 -0
  6. tico/passes/convert_to_relu6.py +1 -1
  7. tico/passes/decompose_fake_quantize_tensor_qparams.py +4 -3
  8. tico/passes/ops.py +0 -1
  9. tico/passes/remove_redundant_expand.py +3 -1
  10. tico/quantization/__init__.py +6 -0
  11. tico/quantization/algorithm/fpi_gptq/fpi_gptq.py +161 -0
  12. tico/quantization/algorithm/fpi_gptq/quantizer.py +179 -0
  13. tico/{experimental/quantization → quantization}/algorithm/gptq/gptq.py +24 -3
  14. tico/{experimental/quantization → quantization}/algorithm/gptq/quantizer.py +14 -6
  15. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/annotator.py +6 -8
  16. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/adaptive_avg_pool2d.py +4 -6
  17. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/add.py +4 -6
  18. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/conv2d.py +4 -6
  19. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/div.py +4 -6
  20. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/linear.py +4 -6
  21. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/mean.py +4 -6
  22. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/mul.py +4 -6
  23. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/relu6.py +4 -6
  24. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/rsqrt.py +4 -6
  25. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/sub.py +4 -6
  26. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/spec.py +1 -3
  27. tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/utils.py +1 -1
  28. tico/{experimental/quantization → quantization}/algorithm/pt2e/quantizer.py +5 -2
  29. tico/{experimental/quantization → quantization}/algorithm/pt2e/utils.py +1 -3
  30. tico/{experimental/quantization → quantization}/algorithm/smoothquant/observer.py +26 -8
  31. tico/{experimental/quantization → quantization}/algorithm/smoothquant/quantizer.py +28 -9
  32. tico/quantization/algorithm/smoothquant/smooth_quant.py +327 -0
  33. tico/quantization/config/base.py +26 -0
  34. tico/quantization/config/fpi_gptq.py +29 -0
  35. tico/quantization/config/gptq.py +29 -0
  36. tico/quantization/config/pt2e.py +25 -0
  37. tico/{experimental/quantization/ptq/quant_config.py → quantization/config/ptq.py} +18 -10
  38. tico/{experimental/quantization/config.py → quantization/config/smoothquant.py} +9 -37
  39. tico/{experimental/quantization → quantization}/evaluation/evaluate.py +6 -12
  40. tico/{experimental/quantization → quantization}/evaluation/executor/circle_executor.py +3 -4
  41. tico/{experimental/quantization → quantization}/evaluation/executor/triv24_executor.py +2 -4
  42. tico/{experimental/quantization → quantization}/evaluation/utils.py +1 -1
  43. tico/{experimental/quantization → quantization}/public_interface.py +11 -18
  44. tico/{experimental/quantization → quantization}/quantizer.py +1 -1
  45. tico/quantization/quantizer_registry.py +73 -0
  46. tico/quantization/wrapq/examples/compare_ppl.py +230 -0
  47. tico/quantization/wrapq/examples/debug_quant_outputs.py +224 -0
  48. tico/{experimental/quantization/ptq → quantization/wrapq}/examples/quantize_linear.py +11 -10
  49. tico/{experimental/quantization/ptq → quantization/wrapq}/examples/quantize_llama_attn.py +10 -12
  50. tico/{experimental/quantization/ptq → quantization/wrapq}/examples/quantize_llama_decoder_layer.py +10 -9
  51. tico/{experimental/quantization/ptq → quantization/wrapq}/examples/quantize_llama_mlp.py +13 -13
  52. tico/quantization/wrapq/examples/quantize_with_gptq.py +265 -0
  53. tico/{experimental/quantization/ptq → quantization/wrapq}/observers/affine_base.py +3 -3
  54. tico/{experimental/quantization/ptq → quantization/wrapq}/observers/base.py +2 -2
  55. tico/{experimental/quantization/ptq → quantization/wrapq}/observers/ema.py +2 -2
  56. tico/{experimental/quantization/ptq → quantization/wrapq}/observers/identity.py +1 -1
  57. tico/{experimental/quantization/ptq → quantization/wrapq}/observers/minmax.py +2 -2
  58. tico/{experimental/quantization/ptq → quantization/wrapq}/observers/mx.py +1 -1
  59. tico/quantization/wrapq/quantizer.py +179 -0
  60. tico/{experimental/quantization/ptq → quantization/wrapq}/utils/introspection.py +3 -5
  61. tico/{experimental/quantization/ptq → quantization/wrapq}/utils/metrics.py +3 -2
  62. tico/quantization/wrapq/wrappers/fairseq/__init__.py +5 -0
  63. tico/quantization/wrapq/wrappers/fairseq/decoder_export_single_step.py +234 -0
  64. tico/quantization/wrapq/wrappers/fairseq/quant_decoder.py +429 -0
  65. tico/quantization/wrapq/wrappers/fairseq/quant_decoder_layer.py +492 -0
  66. tico/quantization/wrapq/wrappers/fairseq/quant_encoder.py +331 -0
  67. tico/quantization/wrapq/wrappers/fairseq/quant_encoder_layer.py +163 -0
  68. tico/quantization/wrapq/wrappers/fairseq/quant_mha.py +381 -0
  69. tico/quantization/wrapq/wrappers/llama/__init__.py +1 -0
  70. tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/llama/quant_attn.py +58 -21
  71. tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/llama/quant_decoder_layer.py +21 -13
  72. tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/llama/quant_mlp.py +5 -7
  73. tico/quantization/wrapq/wrappers/nn/__init__.py +1 -0
  74. tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/nn/quant_layernorm.py +6 -7
  75. tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/nn/quant_linear.py +7 -8
  76. tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/nn/quant_silu.py +8 -9
  77. tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/ptq_wrapper.py +4 -6
  78. tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/quant_elementwise.py +55 -17
  79. tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/quant_module_base.py +10 -9
  80. tico/{experimental/quantization/ptq → quantization/wrapq}/wrappers/registry.py +17 -10
  81. tico/serialize/circle_serializer.py +11 -4
  82. tico/serialize/operators/op_constant_pad_nd.py +41 -11
  83. tico/serialize/operators/op_le.py +54 -0
  84. tico/serialize/operators/op_mm.py +15 -132
  85. tico/utils/convert.py +20 -15
  86. tico/utils/register_custom_op.py +6 -4
  87. tico/utils/signature.py +7 -8
  88. tico/utils/validate_args_kwargs.py +12 -0
  89. {tico-0.1.0.dev250904.dist-info → tico-0.1.0.dev251109.dist-info}/METADATA +48 -2
  90. {tico-0.1.0.dev250904.dist-info → tico-0.1.0.dev251109.dist-info}/RECORD +128 -108
  91. tico/experimental/quantization/__init__.py +0 -6
  92. tico/experimental/quantization/algorithm/smoothquant/smooth_quant.py +0 -164
  93. tico/experimental/quantization/ptq/examples/compare_ppl.py +0 -121
  94. tico/experimental/quantization/ptq/examples/debug_quant_outputs.py +0 -129
  95. tico/experimental/quantization/ptq/examples/quantize_with_gptq.py +0 -165
  96. /tico/{experimental/quantization → quantization}/algorithm/__init__.py +0 -0
  97. /tico/{experimental/quantization/algorithm/gptq → quantization/algorithm/fpi_gptq}/__init__.py +0 -0
  98. /tico/{experimental/quantization/algorithm/pt2e → quantization/algorithm/gptq}/__init__.py +0 -0
  99. /tico/{experimental/quantization → quantization}/algorithm/gptq/quant.py +0 -0
  100. /tico/{experimental/quantization → quantization}/algorithm/gptq/utils.py +0 -0
  101. /tico/{experimental/quantization/algorithm/pt2e/annotation → quantization/algorithm/pt2e}/__init__.py +0 -0
  102. /tico/{experimental/quantization/algorithm/pt2e/transformation → quantization/algorithm/pt2e/annotation}/__init__.py +0 -0
  103. /tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/config.py +0 -0
  104. /tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/__init__.py +0 -0
  105. /tico/{experimental/quantization/algorithm/smoothquant → quantization/algorithm/pt2e/transformation}/__init__.py +0 -0
  106. /tico/{experimental/quantization → quantization}/algorithm/pt2e/transformation/convert_scalars_to_attrs.py +0 -0
  107. /tico/{experimental/quantization/evaluation → quantization/algorithm/smoothquant}/__init__.py +0 -0
  108. /tico/{experimental/quantization/evaluation/executor → quantization/config}/__init__.py +0 -0
  109. /tico/{experimental/quantization/passes → quantization/evaluation}/__init__.py +0 -0
  110. /tico/{experimental/quantization → quantization}/evaluation/backend.py +0 -0
  111. /tico/{experimental/quantization/ptq → quantization/evaluation/executor}/__init__.py +0 -0
  112. /tico/{experimental/quantization → quantization}/evaluation/executor/backend_executor.py +0 -0
  113. /tico/{experimental/quantization → quantization}/evaluation/metric.py +0 -0
  114. /tico/{experimental/quantization/ptq/examples → quantization/passes}/__init__.py +0 -0
  115. /tico/{experimental/quantization → quantization}/passes/fold_quant_ops.py +0 -0
  116. /tico/{experimental/quantization → quantization}/passes/insert_quantize_on_dtype_mismatch.py +0 -0
  117. /tico/{experimental/quantization → quantization}/passes/propagate_qparam_backward.py +0 -0
  118. /tico/{experimental/quantization → quantization}/passes/propagate_qparam_forward.py +0 -0
  119. /tico/{experimental/quantization → quantization}/passes/quantize_bias.py +0 -0
  120. /tico/{experimental/quantization → quantization}/passes/remove_weight_dequant_op.py +0 -0
  121. /tico/{experimental/quantization/ptq/observers → quantization/wrapq}/__init__.py +0 -0
  122. /tico/{experimental/quantization/ptq → quantization/wrapq}/dtypes.py +0 -0
  123. /tico/{experimental/quantization/ptq/utils → quantization/wrapq/examples}/__init__.py +0 -0
  124. /tico/{experimental/quantization/ptq → quantization/wrapq}/mode.py +0 -0
  125. /tico/{experimental/quantization/ptq/wrappers → quantization/wrapq/observers}/__init__.py +0 -0
  126. /tico/{experimental/quantization/ptq → quantization/wrapq}/qscheme.py +0 -0
  127. /tico/{experimental/quantization/ptq/wrappers/llama → quantization/wrapq/utils}/__init__.py +0 -0
  128. /tico/{experimental/quantization/ptq → quantization/wrapq}/utils/reduce_utils.py +0 -0
  129. /tico/{experimental/quantization/ptq/wrappers/nn → quantization/wrapq/wrappers}/__init__.py +0 -0
  130. {tico-0.1.0.dev250904.dist-info → tico-0.1.0.dev251109.dist-info}/LICENSE +0 -0
  131. {tico-0.1.0.dev250904.dist-info → tico-0.1.0.dev251109.dist-info}/WHEEL +0 -0
  132. {tico-0.1.0.dev250904.dist-info → tico-0.1.0.dev251109.dist-info}/entry_points.txt +0 -0
  133. {tico-0.1.0.dev250904.dist-info → tico-0.1.0.dev251109.dist-info}/top_level.txt +0 -0
@@ -12,43 +12,9 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from abc import ABC, abstractmethod
16
- from typing import Dict, Optional
15
+ from typing import Dict, Literal, Optional
17
16
 
18
-
19
- class BaseConfig(ABC):
20
- """
21
- Base configuration class for quantization.
22
- """
23
-
24
- @property
25
- @abstractmethod
26
- def name(self) -> str:
27
- pass
28
-
29
-
30
- class PT2EConfig(BaseConfig):
31
- """
32
- Configuration for pytorch 2.0 export quantization.
33
- """
34
-
35
- @property
36
- def name(self) -> str:
37
- return "pt2e"
38
-
39
-
40
- class GPTQConfig(BaseConfig):
41
- """
42
- Configuration for GPTQ.
43
- """
44
-
45
- def __init__(self, verbose: bool = False, show_progress: bool = True):
46
- self.verbose = verbose
47
- self.show_progress = show_progress
48
-
49
- @property
50
- def name(self) -> str:
51
- return "gptq"
17
+ from tico.quantization.config.base import BaseConfig
52
18
 
53
19
 
54
20
  class SmoothQuantConfig(BaseConfig):
@@ -60,10 +26,16 @@ class SmoothQuantConfig(BaseConfig):
60
26
  self,
61
27
  alpha: float = 0.5,
62
28
  custom_alpha_map: Optional[Dict[str, float]] = None,
29
+ acts_from: Literal["input", "output"] = "input",
63
30
  ):
64
31
  self.alpha = alpha
65
32
  self.custom_alpha_map = custom_alpha_map
33
+ # Where to collect activation statistics from:
34
+ # - "input": use forward-pre-hook (Tensor before the Linear op)
35
+ # - "output": use forward-hook (Tensor after the Linear op)
36
+ # Default is "input".
37
+ self.acts_from = acts_from
66
38
 
67
39
  @property
68
40
  def name(self) -> str:
69
- return "smooth_quant"
41
+ return "smoothquant"
@@ -20,18 +20,12 @@ import torch
20
20
  from circle_schema import circle
21
21
  from torch.utils import _pytree as pytree
22
22
 
23
- from tico.experimental.quantization.evaluation.backend import BACKEND
24
- from tico.experimental.quantization.evaluation.executor.backend_executor import (
25
- BackendExecutor,
26
- )
27
- from tico.experimental.quantization.evaluation.executor.circle_executor import (
28
- CircleExecutor,
29
- )
30
- from tico.experimental.quantization.evaluation.executor.triv24_executor import (
31
- Triv24Executor,
32
- )
33
- from tico.experimental.quantization.evaluation.metric import MetricCalculator
34
- from tico.experimental.quantization.evaluation.utils import (
23
+ from tico.quantization.evaluation.backend import BACKEND
24
+ from tico.quantization.evaluation.executor.backend_executor import BackendExecutor
25
+ from tico.quantization.evaluation.executor.circle_executor import CircleExecutor
26
+ from tico.quantization.evaluation.executor.triv24_executor import Triv24Executor
27
+ from tico.quantization.evaluation.metric import MetricCalculator
28
+ from tico.quantization.evaluation.utils import (
35
29
  ensure_list,
36
30
  find_invalid_types,
37
31
  get_graph_input_output,
@@ -19,9 +19,7 @@ from typing import List
19
19
  import numpy as np
20
20
  import torch
21
21
 
22
- from tico.experimental.quantization.evaluation.executor.backend_executor import (
23
- BackendExecutor,
24
- )
22
+ from tico.quantization.evaluation.executor.backend_executor import BackendExecutor
25
23
  from tico.utils.model import CircleModel
26
24
  from tico.utils.utils import run_bash_cmd
27
25
 
@@ -72,4 +70,5 @@ class CircleExecutor(BackendExecutor):
72
70
  return out
73
71
 
74
72
  def __del__(self):
75
- self.temp_dir.cleanup()
73
+ if hasattr(self, "temp_dir") and self.temp_dir:
74
+ self.temp_dir.cleanup()
@@ -20,10 +20,8 @@ import numpy as np
20
20
  import torch
21
21
  from circle_schema import circle
22
22
 
23
- from tico.experimental.quantization.evaluation.executor.backend_executor import (
24
- BackendExecutor,
25
- )
26
- from tico.experimental.quantization.evaluation.utils import (
23
+ from tico.quantization.evaluation.executor.backend_executor import BackendExecutor
24
+ from tico.quantization.evaluation.utils import (
27
25
  dequantize,
28
26
  get_graph_input_output,
29
27
  quantize,
@@ -44,7 +44,7 @@ def quantize(
44
44
  data = np.array(data)
45
45
  # Perfrom quantization
46
46
  if not scale:
47
- logger.warn("WARNING: scale value is 0. 1e-7 will be used instead.")
47
+ logger.warning("WARNING: scale value is 0. 1e-7 will be used instead.")
48
48
  scale = 1e-7
49
49
  rescaled = np.round(data / scale) + zero_point
50
50
  # Clamp the values
@@ -13,25 +13,17 @@
13
13
  # limitations under the License.
14
14
 
15
15
  import copy
16
- from typing import Any, Dict, Optional, Type
16
+ from typing import Any, Dict, Optional
17
17
 
18
18
  import torch
19
19
 
20
- from tico.experimental.quantization.algorithm.gptq.quantizer import GPTQQuantizer
21
- from tico.experimental.quantization.algorithm.pt2e.quantizer import PT2EQuantizer
22
- from tico.experimental.quantization.algorithm.smoothquant.quantizer import (
23
- SmoothQuantQuantizer,
24
- )
25
- from tico.experimental.quantization.config import BaseConfig
26
- from tico.experimental.quantization.quantizer import BaseQuantizer
20
+ from tico.quantization.algorithm.gptq.quantizer import GPTQQuantizer
21
+ from tico.quantization.algorithm.pt2e.quantizer import PT2EQuantizer
22
+ from tico.quantization.config.base import BaseConfig
23
+ from tico.quantization.quantizer import BaseQuantizer
24
+ from tico.quantization.quantizer_registry import get_quantizer
27
25
 
28
26
 
29
- config_to_quantizer: Dict[str, Type[BaseQuantizer]] = {
30
- "pt2e": PT2EQuantizer,
31
- "gptq": GPTQQuantizer,
32
- "smooth_quant": SmoothQuantQuantizer,
33
- }
34
-
35
27
  QUANTIZER_ATTRIBUTE_NAME = "tico_quantizer"
36
28
 
37
29
 
@@ -40,7 +32,7 @@ def prepare(
40
32
  quant_config: BaseConfig,
41
33
  args: Optional[Any] = None,
42
34
  kwargs: Optional[Dict[str, Any]] = None,
43
- inplace: Optional[bool] = False,
35
+ inplace: Optional[bool] = True,
44
36
  ):
45
37
  """
46
38
  Prepare the model for quantization using the provided configuration.
@@ -61,21 +53,22 @@ def prepare(
61
53
  """
62
54
  if hasattr(model, QUANTIZER_ATTRIBUTE_NAME):
63
55
  raise RuntimeError("prepare() already has been called.")
64
- if quant_config.name == "pt2e" and inplace:
56
+ quantizer = get_quantizer(quant_config)
57
+
58
+ if isinstance(quantizer, PT2EQuantizer) and inplace:
65
59
  raise RuntimeError(
66
60
  "In-place is not supported for PT2E quantization due to limitation in the underlying Torch APIs. Please set 'inplace=False' to proceed."
67
61
  )
68
62
 
69
63
  model = model if inplace else copy.deepcopy(model)
70
64
 
71
- quantizer = config_to_quantizer[quant_config.name](quant_config)
72
65
  model = quantizer.prepare(model, args, kwargs)
73
66
  setattr(model, QUANTIZER_ATTRIBUTE_NAME, quantizer)
74
67
 
75
68
  return model
76
69
 
77
70
 
78
- def convert(model, inplace: Optional[bool] = False):
71
+ def convert(model, inplace: Optional[bool] = True):
79
72
  """
80
73
  Convert the prepared model to a quantized model using the provided configuration.
81
74
 
@@ -17,7 +17,7 @@ from typing import Any, Dict, Optional
17
17
 
18
18
  import torch
19
19
 
20
- from tico.experimental.quantization.config import BaseConfig
20
+ from tico.quantization.config.base import BaseConfig
21
21
 
22
22
 
23
23
  class BaseQuantizer(ABC):
@@ -0,0 +1,73 @@
1
+ # Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import importlib
16
+ from typing import Dict, Optional, Type, TypeVar
17
+
18
+ from tico.quantization.config.base import BaseConfig
19
+ from tico.quantization.quantizer import BaseQuantizer
20
+
21
+ TQ = TypeVar("TQ", bound=BaseQuantizer)
22
+
23
+ # Mapping: Config type -> Quantizer type
24
+ _REGISTRY: Dict[Type[BaseConfig], Type[BaseQuantizer]] = {}
25
+
26
+
27
+ def register_quantizer(config_cls: Type[BaseConfig]):
28
+ """
29
+ Decorator to register a quantizer for a given config class.
30
+ Usage:
31
+ @register_quantizer(GPTQConfig)
32
+ class GPTQQuantizer(BaseQuantizer): ...
33
+ """
34
+
35
+ def wrapper(quantizer_cls: Type[TQ]) -> Type[TQ]:
36
+ _REGISTRY[config_cls] = quantizer_cls
37
+ return quantizer_cls
38
+
39
+ return wrapper
40
+
41
+
42
+ def _lookup(cfg: BaseConfig) -> Optional[Type[BaseQuantizer]]:
43
+ """Return a quantizer class only if the exact config type is registered."""
44
+ return _REGISTRY.get(type(cfg))
45
+
46
+
47
+ def get_quantizer(cfg: BaseConfig) -> BaseQuantizer:
48
+ """Factory to return a quantizer instance for the given config."""
49
+ qcls = _lookup(cfg)
50
+ if qcls is not None:
51
+ return qcls(cfg)
52
+
53
+ # Lazy import by naming convention
54
+ name = getattr(cfg, "name", None)
55
+ if name:
56
+ if name == "ptq":
57
+ importlib.import_module(f"tico.quantization.wrapq.quantizer")
58
+ else:
59
+ try:
60
+ importlib.import_module(f"tico.quantization.algorithm.{name}.quantizer")
61
+ except Exception as e:
62
+ raise RuntimeError(
63
+ f"Failed to import quantizer module for config name='{name}': {e}"
64
+ )
65
+
66
+ qcls = _lookup(cfg)
67
+ if qcls is not None:
68
+ return qcls(cfg)
69
+
70
+ raise RuntimeError(
71
+ f"No quantizer registered for config type {type(cfg).__name__} "
72
+ f"(name='{getattr(cfg,'name',None)}')."
73
+ )
@@ -0,0 +1,230 @@
1
+ # Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ # =============================================================================
16
+ # QUICK PTQ WORKFLOW (OPTIONAL FP32 BASELINE)
17
+ # -----------------------------------------------------------------------------
18
+ # Toggle RUN_FP to choose between:
19
+ # • FP32 perplexity measurement only, OR
20
+ # • Full post-training UINT-8 flow (wrap → calibrate → eval).
21
+ # =============================================================================
22
+
23
+ import argparse
24
+ import sys
25
+
26
+ import torch
27
+ import tqdm
28
+ from datasets import load_dataset
29
+ from transformers import AutoModelForCausalLM, AutoTokenizer
30
+
31
+ from tico.quantization import convert, prepare
32
+ from tico.quantization.config.ptq import PTQConfig
33
+ from tico.quantization.wrapq.utils.metrics import perplexity
34
+
35
+ # Token-budget presets for activation calibration
36
+ TOKENS: dict[str, int] = {
37
+ # Smoke test (<1 min turnaround on CPU/GPU)
38
+ "debug": 2_000, # ≈16 × 128-seq batches
39
+ # Good default for 1-7B models (≲3 % ppl delta)
40
+ "baseline": 50_000,
41
+ # Production / 4-bit observer smoothing
42
+ "production": 200_000,
43
+ }
44
+
45
+ DTYPE_MAP = {
46
+ "float32": torch.float32,
47
+ "bfloat16": torch.bfloat16,
48
+ "float16": torch.float16,
49
+ }
50
+
51
+ # Hardcoded dataset settings
52
+ DATASET_NAME = "wikitext"
53
+ DATASET_CONFIG = "wikitext-2-raw-v1"
54
+ TRAIN_SPLIT = "train"
55
+ TEST_SPLIT = "test"
56
+
57
+
58
+ def main():
59
+ parser = argparse.ArgumentParser(description="Quick PTQ example (FP or UINT8)")
60
+ parser.add_argument(
61
+ "--mode",
62
+ choices=["fp", "uint8"],
63
+ default="fp",
64
+ help="Choose FP baseline only or full UINT8 PTQ path.",
65
+ )
66
+ parser.add_argument(
67
+ "--model", type=str, required=True, help="HF repo name or local path."
68
+ )
69
+ parser.add_argument(
70
+ "--device",
71
+ type=str,
72
+ default="cuda" if torch.cuda.is_available() else "cpu",
73
+ help="Device to run on (cuda|cpu).",
74
+ )
75
+ parser.add_argument(
76
+ "--dtype",
77
+ choices=list(DTYPE_MAP.keys()),
78
+ default="float32",
79
+ help=f"Model dtype for load.",
80
+ )
81
+ parser.add_argument(
82
+ "--stride", type=int, default=512, help="Sliding-window stride for perplexity."
83
+ )
84
+ parser.add_argument("--seed", type=int, default=42, help="Random seed.")
85
+ parser.add_argument(
86
+ "--trust-remote-code",
87
+ action="store_true",
88
+ help="Enable only if you trust the model repo code.",
89
+ )
90
+ parser.add_argument(
91
+ "--hf-token",
92
+ type=str,
93
+ default=None,
94
+ help="Optional HF token for gated/private models.",
95
+ )
96
+ parser.add_argument(
97
+ "--use-cache",
98
+ dest="use_cache",
99
+ action="store_true",
100
+ default=False,
101
+ help="Use model KV cache if enabled (off by default).",
102
+ )
103
+ parser.add_argument(
104
+ "--no-tqdm", action="store_true", help="Disable tqdm progress bars."
105
+ )
106
+ # 2) calib-preset default = debug
107
+ parser.add_argument(
108
+ "--calib-preset",
109
+ choices=list(TOKENS.keys()),
110
+ default="debug",
111
+ help="Calibration token budget preset.",
112
+ )
113
+
114
+ args = parser.parse_args()
115
+
116
+ # Basic setup
117
+ torch.manual_seed(args.seed)
118
+ device = torch.device(args.device)
119
+ dtype = DTYPE_MAP[args.dtype]
120
+
121
+ print("=== Config ===")
122
+ print(f"Mode : {args.mode}")
123
+ print(f"Model : {args.model}")
124
+ print(f"Device : {device.type}")
125
+ print(f"DType : {args.dtype}")
126
+ print(f"Stride : {args.stride}")
127
+ print(f"Use HF cache? : {args.use_cache}")
128
+ print(
129
+ f"Calib preset : {args.calib_preset} ({TOKENS[args.calib_preset]:,} tokens)"
130
+ )
131
+ print()
132
+
133
+ # -------------------------------------------------------------------------
134
+ # 1. Load model and tokenizer
135
+ # -------------------------------------------------------------------------
136
+ tokenizer = AutoTokenizer.from_pretrained(
137
+ args.model,
138
+ trust_remote_code=args.trust_remote_code,
139
+ token=args.hf_token,
140
+ )
141
+
142
+ model = (
143
+ AutoModelForCausalLM.from_pretrained(
144
+ args.model,
145
+ torch_dtype=dtype,
146
+ trust_remote_code=args.trust_remote_code,
147
+ token=args.hf_token,
148
+ )
149
+ .to(device)
150
+ .eval()
151
+ )
152
+
153
+ model.config.use_cache = args.use_cache
154
+
155
+ if args.mode == "fp":
156
+ fp_model = model
157
+ else:
158
+ # INT8 PTQ path
159
+ uint8_model = model
160
+
161
+ CALIB_TOKENS = TOKENS[args.calib_preset]
162
+ print(f"Calibrating with {CALIB_TOKENS:,} tokens.\n")
163
+
164
+ # ---------------------------------------------------------------------
165
+ # 2. Wrap every Transformer layer with PTQWrapper
166
+ # ---------------------------------------------------------------------
167
+ qcfg = PTQConfig() # all-uint8 defaults
168
+ prepare(uint8_model, qcfg)
169
+
170
+ # ---------------------------------------------------------------------
171
+ # 3. Single-pass activation calibration
172
+ # ---------------------------------------------------------------------
173
+ print("Calibrating UINT-8 observers …")
174
+ calib_txt = " ".join(
175
+ load_dataset(DATASET_NAME, DATASET_CONFIG, split=TRAIN_SPLIT)["text"]
176
+ )[:CALIB_TOKENS]
177
+ ids = tokenizer(calib_txt, return_tensors="pt").input_ids.to(device)
178
+
179
+ # Run inference to collect ranges
180
+ iterator = range(0, ids.size(1) - 1, args.stride)
181
+ if not args.no_tqdm:
182
+ iterator = tqdm.tqdm(iterator, desc="Calibration")
183
+ with torch.no_grad():
184
+ for i in iterator:
185
+ uint8_model(ids[:, i : i + args.stride])
186
+
187
+ # Freeze (scale, zero-point)
188
+ convert(uint8_model)
189
+
190
+ # -------------------------------------------------------------------------
191
+ # 4. Evaluate perplexity
192
+ # -------------------------------------------------------------------------
193
+ print("\nCalculating perplexities …")
194
+ test_ds = load_dataset(DATASET_NAME, DATASET_CONFIG, split=TEST_SPLIT)
195
+ enc = tokenizer("\n\n".join(test_ds["text"]), return_tensors="pt")
196
+
197
+ if args.mode == "fp":
198
+ ppl_fp = perplexity(
199
+ fp_model,
200
+ enc,
201
+ args.device,
202
+ stride=args.stride,
203
+ show_progress=not args.no_tqdm,
204
+ )
205
+ else:
206
+ ppl_int8 = perplexity(
207
+ uint8_model,
208
+ enc,
209
+ args.device,
210
+ stride=args.stride,
211
+ show_progress=not args.no_tqdm,
212
+ )
213
+
214
+ # -------------------------------------------------------------------------
215
+ # 5. Report
216
+ # -------------------------------------------------------------------------
217
+ print("\n┌── Wikitext-2 test perplexity ─────────────")
218
+ if args.mode == "fp":
219
+ print(f"│ FP : {ppl_fp:8.2f}")
220
+ else:
221
+ print(f"│ UINT-8 : {ppl_int8:8.2f}")
222
+ print("└───────────────────────────────────────────")
223
+
224
+
225
+ if __name__ == "__main__":
226
+ try:
227
+ main()
228
+ except Exception as e:
229
+ print(f"\n[Error] {e}", file=sys.stderr)
230
+ sys.exit(1)