mct-nightly 2.2.0.20240917.426__py3-none-any.whl → 2.2.0.20240919.455__py3-none-any.whl

model_compression_toolkit/gptq/pytorch/gptq_training.py

@@ -21,6 +21,8 @@ import copy
 import torch
 
 from model_compression_toolkit.core.common.hessian import HessianInfoService
+from model_compression_toolkit.gptq.pytorch.quantizer.gradual_activation_quantization import \
+    get_gradual_activation_quantizer_wrapper_factory
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.pytorch.back2framework.pytorch_model_builder import PyTorchModelBuilder
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
@@ -36,6 +38,7 @@ from model_compression_toolkit.gptq.pytorch.graph_info import get_gptq_trainable
 from model_compression_toolkit.gptq.pytorch.quantizer.quantization_builder import quantization_builder
 from model_compression_toolkit.gptq.pytorch.quantizer.regularization_factory import get_regularization
 from mct_quantizers import PytorchQuantizationWrapper, PytorchActivationQuantizationHolder
+from model_compression_toolkit.trainable_infrastructure.pytorch.util import get_total_grad_steps
 
 
 class PytorchGPTQTrainer(GPTQTrainer):
@@ -66,6 +69,13 @@ class PytorchGPTQTrainer(GPTQTrainer):
             representative_data_gen: Dataset to use for inputs of the models.
             hessian_info_service: HessianInfoService to fetch info based on the hessian approximation of the float model.
         """
+        def _get_total_grad_steps():
+            return get_total_grad_steps(representative_data_gen) * gptq_config.n_epochs
+
+        # must be set prior to model building in the base class constructor
+        self.gradual_act_quantizer_wrapper_factory = get_gradual_activation_quantizer_wrapper_factory(
+            gptq_config, _get_total_grad_steps)
+
         super().__init__(graph_float,
                          graph_quant,
                          gptq_config,
@@ -98,7 +108,7 @@ class PytorchGPTQTrainer(GPTQTrainer):
 
         self.weights_for_average_loss = to_torch_tensor(self.compute_hessian_based_weights())
 
-        self.reg_func = get_regularization(self.gptq_config, representative_data_gen)
+        self.reg_func = get_regularization(self.gptq_config, _get_total_grad_steps)
 
     def _is_gptq_weights_trainable(self,
                                    node: BaseNode) -> bool:
@@ -145,7 +155,6 @@ class PytorchGPTQTrainer(GPTQTrainer):
     def get_activation_quantizer_holder(self, n: BaseNode) -> Callable:
         """
         Retrieve a PytorchActivationQuantizationHolder layer to use for activation quantization of a node.
-        If the layer is not supposed to be wrapped with an activation quantizer - return None.
         Args:
             n: Node to attach a PytorchActivationQuantizationHolder to its output.
         Returns:
@@ -153,13 +162,13 @@ class PytorchGPTQTrainer(GPTQTrainer):
         """
         _, activation_quantizers = quantization_builder(n, self.gptq_config)
         # Holder by definition uses a single quantizer for the activation quantization
-        # thus we make sure this is the only possible case (unless it's a node we no activation
-        # quantization, which in this case has an empty list).
-        if len(activation_quantizers) == 1:
-            return PytorchActivationQuantizationHolder(activation_quantizers[0])
-        Logger.critical(f"'PytorchActivationQuantizationHolder' requires exactly one quantizer, "
-                        f"but {len(activation_quantizers)} were found for node {n.name}. "
-                        f"Ensure the node is configured with a single activation quantizer.")
+        # thus we make sure this is the only possible case
+        if len(activation_quantizers) != 1:
+            Logger.critical(f"'PytorchActivationQuantizationHolder' requires exactly one quantizer, "
+                            f"but {len(activation_quantizers)} were found for node {n.name}. "
+                            f"Ensure the node is configured with a single activation quantizer.")
+        quantizer = self.gradual_act_quantizer_wrapper_factory(activation_quantizers[0])
+        return PytorchActivationQuantizationHolder(quantizer)
 
     def build_gptq_model(self):
         """

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -13,26 +13,26 @@
 # limitations under the License.
 # ==============================================================================
 import copy
+from typing import Callable, Union
 
-from typing import Callable
-from model_compression_toolkit.core import common
-from model_compression_toolkit.constants import ACT_HESSIAN_DEFAULT_BATCH_SIZE
-from model_compression_toolkit.verify_packages import FOUND_TORCH
+from model_compression_toolkit.constants import ACT_HESSIAN_DEFAULT_BATCH_SIZE, PYTORCH
+from model_compression_toolkit.core import CoreConfig
+from model_compression_toolkit.core.analyzer import analyzer_model_quantization
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
+    MixedPrecisionQuantizationConfig
+from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import \
+    ResourceUtilization
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
-from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT
-from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.constants import PYTORCH
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig, GPTQHessianScoresConfig
-from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
-from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.runner import core_runner
+from model_compression_toolkit.gptq.common.gptq_config import (
+    GradientPTQConfig, GPTQHessianScoresConfig, GradualActivationQuantizationConfig)
+from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT
 from model_compression_toolkit.gptq.keras.quantization_facade import GPTQ_MOMENTUM
 from model_compression_toolkit.gptq.runner import gptq_runner
-from model_compression_toolkit.core.analyzer import analyzer_model_quantization
-from model_compression_toolkit.core import CoreConfig
-from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
-    MixedPrecisionQuantizationConfig
-from model_compression_toolkit.metadata import get_versions_dict, create_model_metadata
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.metadata import create_model_metadata
+from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
+from model_compression_toolkit.verify_packages import FOUND_TORCH
 
 LR_DEFAULT = 1e-4
 LR_REST_DEFAULT = 1e-4
@@ -53,33 +53,38 @@ if FOUND_TORCH:
     DEFAULT_PYTORCH_TPC = get_target_platform_capabilities(PYTORCH, DEFAULT_TP_MODEL)
 
     def get_pytorch_gptq_config(n_epochs: int,
-                                optimizer: Optimizer = Adam([torch.Tensor([])], lr=LR_DEFAULT),
-                                optimizer_rest: Optimizer = Adam([torch.Tensor([])], lr=LR_REST_DEFAULT),
+                                optimizer: Optimizer = None,
+                                optimizer_rest: Optimizer = None,
                                 loss: Callable = multiple_tensors_mse_loss,
                                 log_function: Callable = None,
                                 use_hessian_based_weights: bool = True,
                                 regularization_factor: float = REG_DEFAULT,
-                                hessian_batch_size: int = ACT_HESSIAN_DEFAULT_BATCH_SIZE
+                                hessian_batch_size: int = ACT_HESSIAN_DEFAULT_BATCH_SIZE,
+                                gradual_activation_quantization: Union[bool, GradualActivationQuantizationConfig] = False,
                                 ) -> GradientPTQConfig:
         """
-        Create a GradientPTQConfigV2 instance for Pytorch models.
+        Create a GradientPTQConfig instance for Pytorch models.
 
         args:
             n_epochs (int): Number of epochs for running the representative dataset for fine-tuning.
             optimizer (Optimizer): Pytorch optimizer to use for fine-tuning for auxiliry variable.
             optimizer_rest (Optimizer): Pytorch optimizer to use for fine-tuning of the bias variable.
-            loss (Callable): loss to use during fine-tuning. should accept 4 lists of tensors. 1st list of quantized tensors, the 2nd list is the float tensors, the 3rd is a list of quantized weights and the 4th is a list of float weights.
+            loss (Callable): loss to use during fine-tuning. See the default loss function for the exact interface.
             log_function (Callable): Function to log information about the gptq process.
             use_hessian_based_weights (bool): Whether to use Hessian-based weights for weighted average loss.
             regularization_factor (float): A floating point number that defines the regularization factor.
             hessian_batch_size (int): Batch size for Hessian computation in Hessian-based weights GPTQ.
+            gradual_activation_quantization (bool, GradualActivationQuantizationConfig):
+              If False, GradualActivationQuantization is disabled.
+              If True, GradualActivationQuantization is enabled with the default settings.
+              GradualActivationQuantizationConfig object can be passed to use non-default settings.
 
         returns:
-            a GradientPTQConfigV2 object to use when fine-tuning the quantized model using gptq.
+            a GradientPTQConfig object to use when fine-tuning the quantized model using gptq.
 
         Examples:
 
-            Import MCT and Create a GradientPTQConfigV2 to run for 5 epochs:
+            Import MCT and Create a GradientPTQConfig to run for 5 epochs:
 
             >>> import model_compression_toolkit as mct
             >>> gptq_conf = mct.gptq.get_pytorch_gptq_config(n_epochs=5)
@@ -89,16 +94,31 @@ if FOUND_TORCH:
             >>> import torch
             >>> gptq_conf = mct.gptq.get_pytorch_gptq_config(n_epochs=3, optimizer=torch.optim.Adam([torch.Tensor(1)]))
 
-            The configuration can be passed to :func:`~model_compression_toolkit.pytorch_post_training_quantization` in order to quantize a pytorch model using gptq.
+            To enable Gradual Activation Quantization with non-default settings build GradualActivationQuantizationConfig:
+            >>> gradual_act_conf = mct.gptq.GradualActivationQuantizationConfig(mct.gptq.QFractionLinearAnnealingConfig(initial_q_fraction=0.2))
+            >>> gptq_conf = mct.gptq.get_pytorch_gptq_config(n_epochs=3, gradual_activation_quantization=gradual_act_conf)
+            The configuration can be passed to :func:`~model_compression_toolkit.pytorch_gradient_post_training_quantization` in order to quantize a pytorch model using gptq.
 
         """
+        optimizer = optimizer or Adam([torch.Tensor([])], lr=LR_DEFAULT)
+        optimizer_rest = optimizer_rest or Adam([torch.Tensor([])], lr=LR_REST_DEFAULT)
+
         bias_optimizer = torch.optim.SGD([torch.Tensor([])], lr=LR_BIAS_DEFAULT, momentum=GPTQ_MOMENTUM)
+
+        if isinstance(gradual_activation_quantization, bool):
+            gradual_quant_config = GradualActivationQuantizationConfig() if gradual_activation_quantization else None
+        elif isinstance(gradual_activation_quantization, GradualActivationQuantizationConfig):
+            gradual_quant_config = gradual_activation_quantization
+        else:
+            raise TypeError(f'gradual_activation_quantization argument should be bool or '
+                            f'GradualActivationQuantizationConfig, received {type(gradual_activation_quantization)}')    # pragma: no cover
+
         return GradientPTQConfig(n_epochs, optimizer, optimizer_rest=optimizer_rest, loss=loss,
                                  log_function=log_function, train_bias=True, optimizer_bias=bias_optimizer,
                                  use_hessian_based_weights=use_hessian_based_weights,
                                  regularization_factor=regularization_factor,
-                                 hessian_weights_config=GPTQHessianScoresConfig(hessian_batch_size=hessian_batch_size))
-
+                                 hessian_weights_config=GPTQHessianScoresConfig(hessian_batch_size=hessian_batch_size),
+                                 gradual_activation_quantization_config=gradual_quant_config)
 
     def pytorch_gradient_post_training_quantization(model: Module,
                                                     representative_data_gen: Callable,
@@ -222,11 +242,11 @@ if FOUND_TORCH:
 else:
     # If torch is not installed,
     # we raise an exception when trying to use these functions.
-    def get_pytorch_gptq_config(*args, **kwargs):
+    def get_pytorch_gptq_config(*args, **kwargs):    # pragma: no cover
         Logger.critical("PyTorch must be installed to use 'get_pytorch_gptq_config'. "
-                        "The 'torch' package is missing.")  # pragma: no cover
+                        "The 'torch' package is missing.")
 
 
-    def pytorch_gradient_post_training_quantization(*args, **kwargs):
+    def pytorch_gradient_post_training_quantization(*args, **kwargs):    # pragma: no cover
         Logger.critical("PyTorch must be installed to use 'pytorch_gradient_post_training_quantization'. "
-                        "The 'torch' package is missing.")  # pragma: no cover
+                        "The 'torch' package is missing.")

model_compression_toolkit/gptq/pytorch/quantizer/gradual_activation_quantization.py

@@ -0,0 +1,80 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from functools import partial
+from typing import Callable
+
+from model_compression_toolkit.gptq import GradientPTQConfig, QFractionLinearAnnealingConfig
+from model_compression_toolkit.trainable_infrastructure import BasePytorchTrainableQuantizer
+
+from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import LinearAnnealingScheduler
+
+
+def get_gradual_activation_quantizer_wrapper_factory(gptq_config: GradientPTQConfig,
+                                                     get_total_grad_steps_fn: Callable[[], int]) \
+        -> Callable[[BasePytorchTrainableQuantizer], 'GradualActivationQuantizerWrapper']:
+    """
+    Get a factory for 'GradualActivationQuantizerWrapper'.
+
+    Args:
+        gptq_config: GPTQ configuration.
+        get_total_grad_steps_fn: a callable to obtain the total expected number of gradient steps.
+
+    Returns:
+        A factory function to build 'GradualActivationQuantizerWrapper' from Quantizer.
+    """
+    if gptq_config.gradual_activation_quantization_config is None:
+        return lambda q: q
+
+    annealing_cfg = gptq_config.gradual_activation_quantization_config.q_fraction_scheduler_policy
+    if isinstance(annealing_cfg, QFractionLinearAnnealingConfig):
+        t_end = annealing_cfg.end_step or get_total_grad_steps_fn()
+        factor_scheduler = LinearAnnealingScheduler(t_start=annealing_cfg.start_step, t_end=t_end,
+                                                    initial_val=annealing_cfg.initial_q_fraction,
+                                                    target_val=annealing_cfg.target_q_fraction)
+    else:
+        raise ValueError(f'Unknown annealing policy {annealing_cfg}')
+
+    return partial(GradualActivationQuantizerWrapper, q_fraction_scheduler=factor_scheduler)
+
+
+class GradualActivationQuantizerWrapper:
+    # TODO update paper's url
+    """
+    Quantizer wrapper for Gradual Activation Quantization training (https://arxiv.org/abs/2309.11531).
+
+    It computes the weighted sum of the float activation 'x' and the quantized activation 'q(x)':
+
+      out = (1 - q_fraction) * x + q_fraction * q(x)
+
+    where 'q_fraction' is a tensor fraction to quantize in the range [0, 1] provided by a scheduler.
+
+    Args:
+        quantizer: quantizer to wrap.
+        q_fraction_scheduler: a callable that accepts a gradient step and returns the corresponding quantized fraction.
+    """
+    def __init__(self, quantizer: BasePytorchTrainableQuantizer, q_fraction_scheduler: Callable[[int], float]):
+        self.quantizer = quantizer
+        self.q_fraction_scheduler = q_fraction_scheduler
+        self.step_cnt = 0
+
+    def __call__(self, x, training: bool = True):
+        q_fraction = self.q_fraction_scheduler(self.step_cnt)
+        out_q = self.quantizer(x, training)
+        out = (1 - q_fraction) * x + q_fraction * out_q
+        self.step_cnt += 1
+        return out
+
+    def initialize_quantization(self, *args, **kwargs):
+        self.quantizer.initialize_quantization(*args, **kwargs)

model_compression_toolkit/gptq/pytorch/quantizer/regularization_factory.py

@@ -12,33 +12,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from tqdm import tqdm
 from typing import Callable
 
-from model_compression_toolkit.gptq import RoundingType, GradientPTQConfig, GradientPTQConfig
+from model_compression_toolkit.gptq import RoundingType, GradientPTQConfig
 from model_compression_toolkit.gptq.pytorch.quantizer.soft_rounding.soft_quantizer_reg import \
     SoftQuantizerRegularization
+from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import LinearAnnealingScheduler
 
 
-def get_regularization(gptq_config: GradientPTQConfig, representative_data_gen: Callable) -> Callable:
+WARMUP_STEP_FRACTION = 0.2
+
+def get_regularization(gptq_config: GradientPTQConfig, get_total_grad_steps_fn: Callable[[], int]) -> Callable:
     """
     Returns a function that computes the regularization term for GPTQ training based on the given
     rounding type in the GPTQ configuration.
 
     Args:
         gptq_config: A GPTQ configuration.
-        representative_data_gen: Dataset used for the GPTQ training.
+        get_total_grad_steps_fn: a callable to obtain the total expected number of gradient steps.
 
     Returns: A function for computing the regularization. If there is no regularization function defined for the given
         rounding type, then it returns a function that just returns 0.
 
     """
     if gptq_config.rounding_type == RoundingType.SoftQuantizer:
-        # dry run on the representative dataset to count number of batches
-        num_batches = 0
-        for _ in tqdm(representative_data_gen(), "GPTQ initialization"):
-            num_batches += 1
-
-        return SoftQuantizerRegularization(total_gradient_steps=num_batches * gptq_config.n_epochs)
+        total_gradient_steps = get_total_grad_steps_fn()
+        t_start = int(WARMUP_STEP_FRACTION * total_gradient_steps)
+        scheduler = LinearAnnealingScheduler(t_start=t_start, t_end=total_gradient_steps, initial_val=20, target_val=2)
+        return SoftQuantizerRegularization(scheduler)
     else:
         return lambda m, e_reg: 0

model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/soft_quantizer_reg.py

@@ -12,57 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import List
+from typing import List, Callable
 
 import torch
-import numpy as np
 from torch import nn
 
+from mct_quantizers import PytorchQuantizationWrapper
 from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
-from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
-from mct_quantizers import PytorchQuantizationWrapper
-
-
-class LinearTempDecay:
-    """
-    Annealing process for the soft quantizer regularization temperature term.
-    """
-
-    def __init__(self, t_max: int, rel_start_decay: float = 0.2, start_b: int = 20, end_b: int = 2):
-        """
-        Initializes a LinearTempDecay object.
-
-        Args:
-            t_max: maximal time step.
-            rel_start_decay: Decay step size at the beginning of the process.
-            start_b: Starting value of the regularization term.
-            end_b: Target value of the regularization term.
-        """
-
-        self.t_max = t_max
-        self.start_decay = rel_start_decay * t_max
-        self.start_b = start_b
-        self.end_b = end_b
-
-    def __call__(self, t: float) -> float:
-        """
-        Cosine annealing scheduler for soft quantizer regularization temperature term.
-
-        Args:
-            t: The current time step.
-
-        Returns: Scheduled temperature.
-        """
-
-        is_before_start_decay = (t < self.start_decay)
-
-        rel_t = (t - self.start_decay) / (self.t_max - self.start_decay)
-
-        return self.start_b * is_before_start_decay + \
-               (1 - is_before_start_decay) * \
-               (self.end_b + (self.start_b - self.end_b) * torch.maximum(to_torch_tensor(np.array([0.0])),
-                                                                         to_torch_tensor(np.array((1 - rel_t)))))
 
 
 class SoftQuantizerRegularization:
@@ -70,16 +27,16 @@ class SoftQuantizerRegularization:
     A class to handle the computation of soft quantizer regularization for GPTQ training.
     """
 
-    def __init__(self, total_gradient_steps: int):
+    def __init__(self, beta_scheduler: Callable[[int], float]):
         """
         Initializes the regularization computation object with a LinearDecay object.
 
         Args:
-            total_gradient_steps: The number of gradient steps during optimization.
+            beta_scheduler: a callable that accepts current time step and returns a corresponding beta value.
         """
 
         # Initializing the temperature decay according to the number of expected gradient steps
-        self.linear_decay = LinearTempDecay(total_gradient_steps)
+        self.beta_scheduler = beta_scheduler
 
         self.count_iter = 0
 
@@ -95,7 +52,7 @@ class SoftQuantizerRegularization:
         """
 
         soft_reg_aux: List[torch.Tensor] = []
-        b = self.linear_decay(self.count_iter)
+        b = self.beta_scheduler(self.count_iter)
         for layer in model.modules():
             if isinstance(layer, PytorchQuantizationWrapper):
                 kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),

model_compression_toolkit/trainable_infrastructure/pytorch/annealing_schedulers.py

@@ -0,0 +1,39 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
+
+
+class LinearAnnealingScheduler:
+    def __init__(self, t_start: int, t_end: int, initial_val: float, target_val: float):
+        """
+        Linear annealing scheduler. Returns the corresponding annealed value per time step.
+
+        Args:
+            t_start: time step to begin annealing.
+            t_end: time step to complete annealing.
+            initial_val: initial value.
+            target_val: target value.
+        """
+        if not (0 <= t_start < t_end):
+            raise ValueError(f'Expected 0 <= t_start < t_end, actual {t_end=} {t_start=}')
+
+        self.t_start = t_start
+        self.t_end = t_end
+        self.initial_val = initial_val
+        self.target_val = target_val
+
+    def __call__(self, t: int) -> float:
+        factor = to_torch_tensor((t - self.t_start) / (self.t_end - self.t_start)).clip(0, 1)
+        return self.initial_val + factor * (self.target_val - self.initial_val)

model_compression_toolkit/trainable_infrastructure/pytorch/util.py

@@ -0,0 +1,29 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from functools import lru_cache
+from typing import Callable
+
+from tqdm import tqdm
+
+
+@lru_cache
+def get_total_grad_steps(representative_data_gen: Callable) -> int:
+    # dry run on the representative dataset to count number of batches
+    num_batches = 0
+    for _ in tqdm(representative_data_gen(), "Estimating representative dataset size"):
+        num_batches += 1
+    return num_batches
+
+

tests_pytest/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. 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.
+# ==============================================================================

tests_pytest/pytorch/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. 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.
+# ==============================================================================

tests_pytest/pytorch/gptq/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. 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.
+# ==============================================================================

tests_pytest/pytorch/gptq/test_annealing_cfg.py

@@ -0,0 +1,40 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. 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 pytest
+
+from model_compression_toolkit.gptq import QFractionLinearAnnealingConfig
+
+
+def test_linear_annealing_cfg_validation():
+    with pytest.raises(ValueError, match='Expected.* target_q_fraction <= 1'):
+        QFractionLinearAnnealingConfig(initial_q_fraction=0.1, target_q_fraction=1.1, start_step=0, end_step=None)
+
+    with pytest.raises(ValueError, match='Expected.* 0 <= initial_q_fraction'):
+        QFractionLinearAnnealingConfig(initial_q_fraction=-0.1, target_q_fraction=-0.9, start_step=0, end_step=100)
+
+    with pytest.raises(ValueError, match='Expected.* initial_q_fraction < target_q_fraction'):
+        QFractionLinearAnnealingConfig(initial_q_fraction=0.1, target_q_fraction=0.1, start_step=0, end_step=100)
+
+    with pytest.raises(ValueError, match='Expected.* initial_q_fraction < target_q_fraction'):
+        QFractionLinearAnnealingConfig(initial_q_fraction=0.2, target_q_fraction=0.1, start_step=0, end_step=100)
+
+    with pytest.raises(ValueError, match='Expected.* start_step >= 0'):
+        QFractionLinearAnnealingConfig(initial_q_fraction=0, target_q_fraction=1, start_step=-1, end_step=100)
+
+    with pytest.raises(ValueError, match='Expected.* start_step < end_step'):
+        QFractionLinearAnnealingConfig(initial_q_fraction=0, target_q_fraction=1, start_step=100, end_step=100)
+
+    with pytest.raises(ValueError, match='Expected.* start_step < end_step'):
+        QFractionLinearAnnealingConfig(initial_q_fraction=0, target_q_fraction=1, start_step=100, end_step=99)