mct-nightly 1.7.1.31122022.post351__py3-none-any.whl → 1.8.0.1042023.post423__py3-none-any.whl

model_compression_toolkit/gptq/pytorch/quantizer/__init__.py

@@ -11,4 +11,8 @@
 # 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 model_compression_toolkit.gptq.pytorch.quantizer.ste_rounding.symmetric_ste
+import model_compression_toolkit.gptq.pytorch.quantizer.soft_rounding.symmetric_soft_quantizer
+import model_compression_toolkit.gptq.pytorch.quantizer.soft_rounding.uniform_soft_quantizer

model_compression_toolkit/gptq/pytorch/quantizer/base_pytorch_gptq_quantizer.py

@@ -0,0 +1,92 @@
+# Copyright 2023 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 abc import abstractmethod
+from typing import Union, Dict, List
+
+from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.core.common.constants import FOUND_TORCH
+from model_compression_toolkit.gptq.common.gptq_constants import WEIGHTS_QUANTIZATION_PARAMS
+
+from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig, \
+    TrainableQuantizerActivationConfig
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import \
+    BaseTrainableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.pytorch.base_pytorch_quantizer import \
+    BasePytorchTrainableQuantizer
+
+if FOUND_TORCH:
+    from torch import Tensor
+    from model_compression_toolkit.quantizers_infrastructure import PytorchQuantizationWrapper
+
+    class BasePytorchGPTQTrainableQuantizer(BasePytorchTrainableQuantizer):
+        """
+        A base class for trainable Pytorch quantizer for GPTQ.
+        """
+
+        def __init__(self,
+                     quantization_config: Union[TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig]):
+            """
+            Initializes BasePytorchGPTQTrainableQuantizer object.
+
+            Args:
+                quantization_config: quantizer config class contains all the information about a quantizer configuration.
+            """
+
+            super().__init__(quantization_config)
+
+        def update_layer_quantization_params(self, layer: PytorchQuantizationWrapper
+                                             ) -> (Dict[str, Tensor], Dict[str, Dict], Dict):
+            """
+            A Function to calculate the needed change in attributes in NodeQuantizationConfig after retraining.
+
+            Args:
+                layer: A wrapped Pytorch layer.
+
+            Returns:
+                3 dictionaries describing the change in layer's weights, weights config, activation config
+                that changed during GPTQ retraining.
+                Keys must match NodeQuantizationConfig attributes
+
+            """
+            weights = {}
+            for weight, quantizer_vars, quantizer in layer.get_weights_vars():
+                if not isinstance(quantizer, BaseTrainableQuantizer):
+                    Logger.error(f"Expecting a GPTQ trainable quantizer, "  # pragma: no cover
+                                 f"but got {type(quantizer)} which is not callable.")
+                weights.update({weight: quantizer(training=False, inputs=quantizer_vars)})
+
+            quant_config = {WEIGHTS_QUANTIZATION_PARAMS: self.get_quant_config()}
+
+            return weights, quant_config, {}
+
+        @abstractmethod
+        def get_quant_config(self):
+            """
+            Returns the config used to edit NodeQuantizationConfig after GPTQ retraining.
+
+            Returns:
+                A dictionary of attributes the quantize_config retraining has changed during GPTQ retraining.
+                Keys must match NodeQuantizationConfig attributes.
+
+            """
+            raise NotImplemented(f'{self.__class__.__name__} have to implement the '  # pragma: no cover
+                                 f'quantizer\'s get_quant_config.')
+
+else:
+    class BasePytorchGPTQTrainableQuantizer:  # pragma: no cover
+        def __init__(self, *args, **kwargs):
+            Logger.critical('Installing Pytorch is mandatory '
+                            'when using BasePytorchGPTQTrainableQuantizer. '
+                            'Could not find torch package.')  # pragma: no cover

model_compression_toolkit/gptq/pytorch/quantizer/quant_utils.py

@@ -30,11 +30,20 @@ def calculate_delta(max_tensor: torch.Tensor,
                     num_bits: int,
                     signed: bool) -> torch.Tensor:
     """
-    Compute the step size for the quantization.
+    Compute the step size for the symmetric quantization.
     """
     return max_tensor / (2 ** (num_bits - int(signed)))
 
 
+def calculate_delta_uniform(min_tensor: torch.Tensor,
+                            max_tensor: torch.Tensor,
+                            num_bits: int) -> torch.Tensor:
+    """
+    Compute the step size for the uniform quantization.
+    """
+    return (max_tensor-min_tensor) / (2 ** num_bits - 1)
+
+
 def ste_ceil(x: torch.Tensor) -> torch.Tensor:
     """
     Return the ceil values of a tensor.
@@ -66,93 +75,6 @@ def ste_clip(x: torch.Tensor, min_val=-1.0, max_val=1.0) -> torch.Tensor:
     return (torch.clip(x, min=min_val, max=max_val) - x).detach() + x
 
 
-def gumbel_softmax(x: torch.Tensor, tau: Union[torch.Tensor,float], gumbel_tensor: Union[torch.Tensor,float], eps: float = 1e-6, axis=0,
-                   gumbel_scale: float = 1.0) -> torch.Tensor:
-    """
-    A gumbel softmax function.
-    Args:
-        x: A tensor of log probability.
-        tau: A temperature tensor.
-        gumbel_tensor: A tensor of gumbel random variable.
-        eps: A small number for numeric stability.
-        axis: A integer representing the axis of which the gumbel softmax applyed on.
-        gumbel_scale: A normalization factor for the gumbel tensor values
-
-    Returns: A gumbel softmax probability tensor.
-
-    """
-    return softmax((log_softmax(x, dim=axis) + gumbel_tensor * gumbel_scale) / (tau + eps), dim=axis)
-
-
-def select_gumbel(prob: torch.Tensor) -> torch.Tensor:
-    """
-    This function apply ste on the output of the gumbel softmax.
-    Args:
-        prob: A tensor of probability.
-
-    Returns: A Tensor of ohe hot vector
-
-    """
-    max_index = torch.argmax(prob, dim=0)
-    axis_list = [i for i in range(len(max_index.shape))]
-    axis_list.insert(0, len(max_index.shape))
-    one_hot_prob = torch.permute(one_hot(max_index, num_classes=prob.shape[0]), axis_list)
-    return one_hot_prob + 0*prob
-
-
-def ste_gumbel(prob: torch.Tensor) -> torch.Tensor:
-    """
-    This function apply ste on the output of the gumbel softmax.
-    Args:
-        prob:A tensor of probability
-
-    Returns: A Tensor of ohe hot vector with STE.
-
-    """
-    delta = (select_gumbel(prob) - prob).detach()
-    return prob + delta
-
-
-def sample_gumbel(shape, eps=1e-6) -> torch.Tensor:
-    """
-    A function that sample a tensor of i.i.d gumbel random variable.
-    Args:
-        shape: The tensor output shape
-        eps: A small number for numeric stability.
-
-    Returns: A tensor of i.i.d gumbel random variable.
-
-    """
-    u = to_torch_tensor(torch.rand(shape))
-    return -torch.log(-torch.log(u + eps) + eps)
-
-
-def symmetric_quantizer(input_tensor: torch.Tensor,
-                        max_tensor: torch.Tensor,
-                        num_bits: int,
-                        signed: bool,
-                        power_of_two: bool = False) -> torch.Tensor:
-    """
-    Quantize a tensor symmetrically.
-    Args:
-        input_tensor: Tensor to quantize. values of this tensor are not changed during gptq.
-        max_tensor: Tensor with max values to compute the threshold.
-        num_bits: Num of bits to use.
-        signed: Signedness of the quantization range.
-        power_of_two: Whether the threshold should be constrained or not.
-    Returns:
-        A quantized tensor.
-    """
-
-    if power_of_two:
-        max_tensor = power_of_two_max(max_tensor)
-    delta_tensor = calculate_delta(max_tensor, num_bits, signed)
-    tensor_q = ste_round(input_tensor / delta_tensor)
-    min_int = -int(signed) * (2 ** (num_bits - int(signed)))
-    max_int = (2 ** (num_bits - int(signed))) - 1
-    return delta_tensor * ste_clip(tensor_q, min_val=min_int, max_val=max_int)
-
-
 def fix_range_to_include_zero(range_min: torch.Tensor,
                               range_max: torch.Tensor,
                               n_bits: int) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -180,34 +102,3 @@ def fix_range_to_include_zero(range_min: torch.Tensor,
     min_range_adj = min_range_adj * mid_range + max_negative * range_min
     max_range_adj = max_range_adj * mid_range + min_positive * range_max
     return min_range_adj, max_range_adj
-
-
-def uniform_quantizer(tensor_data: torch.Tensor,
-                       range_min: torch.Tensor,
-                       range_max: torch.Tensor,
-                       n_bits: int) -> torch.Tensor:
-    """
-    Quantize a tensor according to given range (min, max) and number of bits.
-    Args:
-        tensor_data: Tensor values to quantize.
-        range_min: minimum bound of the range for quantization (or array of min values per channel).
-        range_max: maximum bound of the range for quantization (or array of max values per channel).
-        n_bits: Number of bits to quantize the tensor.
-    Returns:
-        Quantized data.
-    """
-    # adjusts the quantization rage so the quantization grid include zero.
-    a, b = fix_range_to_include_zero(range_min, range_max, n_bits)
-
-    # Compute the step size of quantized values.
-    delta_tensor = (b - a) / (2 ** n_bits - 1)
-
-    # Apply rounding
-    input_tensor_int = ste_round((tensor_data - a) / delta_tensor)
-
-    # Clip data in range
-    clipped_tensor = ste_clip(input_tensor_int, min_val=0, max_val=2 ** n_bits - 1)
-
-    # Quantize the data between min/max of quantization range.
-    q = delta_tensor * clipped_tensor + a
-    return q

model_compression_toolkit/gptq/pytorch/quantizer/quantization_builder.py

@@ -0,0 +1,75 @@
+# Copyright 2023 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 typing import List, Dict, Tuple
+
+from model_compression_toolkit.gptq import GradientPTQConfigV2
+from model_compression_toolkit.core import common
+from model_compression_toolkit.core.pytorch.constants import KERNEL
+from model_compression_toolkit.exporter.model_wrapper.pytorch.builder.node_to_quantizer import \
+    get_activation_inferable_quantizer_kwargs
+from model_compression_toolkit.gptq.pytorch.quantizer.base_pytorch_gptq_quantizer import \
+    BasePytorchGPTQTrainableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.get_quantizers import \
+    get_inferable_quantizer_class
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.pytorch.quantizers import \
+    BasePyTorchInferableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.get_quantizer_config import \
+    get_trainable_quantizer_weights_config
+from model_compression_toolkit.qat.pytorch.quantizer.base_pytorch_qat_quantizer import BasePytorchQATTrainableQuantizer
+from model_compression_toolkit.quantizers_infrastructure import QuantizationTarget
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.get_quantizers import \
+    get_trainable_quantizer_class
+
+
+def quantization_builder(n: common.BaseNode,
+                         gptq_config: GradientPTQConfigV2,
+                         ) -> Tuple[Dict[str, BasePytorchQATTrainableQuantizer],
+                                    List[BasePyTorchInferableQuantizer]]:
+    """
+    Build quantizers for a node according to its quantization configuration and
+    a global NoOpQuantizeConfig object.
+
+    Args:
+        n: Node to build its QuantizeConfig.
+        gptq_config (GradientPTQConfigV2): GradientPTQConfigV2 configuration.
+
+    Returns:
+        A dictionary which maps the weights kernel attribute to a quantizer for GPTQ training.
+        Note that we return a dictionary although there is only a single attribute that is being mapped to a quantizer,
+        to be compatible with the quantization infrastructure template.
+    """
+
+    weights_quantizers = {}
+    if n.is_weights_quantization_enabled():
+        quant_method = n.final_weights_quantization_cfg.weights_quantization_method
+        quantizer_class = get_trainable_quantizer_class(quant_target=QuantizationTarget.Weights,
+                                                        quantizer_type=gptq_config.rounding_type,
+                                                        quant_method=quant_method,
+                                                        quantizer_base_class=BasePytorchGPTQTrainableQuantizer)
+        weights_quantizers.update({KERNEL: quantizer_class(get_trainable_quantizer_weights_config(n),
+                                                           **gptq_config.gptq_quantizer_params_override)})
+    activation_quantizers = []
+    if n.is_activation_quantization_enabled():
+        quant_method = n.final_activation_quantization_cfg.activation_quantization_method
+
+        quantizer_class = get_inferable_quantizer_class(quant_target=QuantizationTarget.Activation,
+                                                        quant_method=quant_method,
+                                                        quantizer_base_class=BasePyTorchInferableQuantizer)
+
+        kwargs = get_activation_inferable_quantizer_kwargs(n)
+
+        activation_quantizers.append(quantizer_class(**kwargs))
+
+    return weights_quantizers, activation_quantizers

model_compression_toolkit/gptq/pytorch/quantizer/regularization_factory.py

@@ -0,0 +1,45 @@
+# Copyright 2023 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 typing import Callable
+
+from model_compression_toolkit.gptq import RoundingType, GradientPTQConfigV2, GradientPTQConfig
+from model_compression_toolkit.gptq.pytorch.quantizer.soft_rounding.soft_quantizer_reg import \
+    SoftQuantizerRegularization
+
+
+def get_regularization(gptq_config: GradientPTQConfig, representative_data_gen: Callable) -> 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.
+
+    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 representative_data_gen():
+            num_batches += 1
+
+        n_epochs = GradientPTQConfigV2.from_v1(n_ptq_iter=num_batches, config_v1=gptq_config).n_epochs if \
+            not type(gptq_config) == GradientPTQConfigV2 else gptq_config.n_epochs
+        return SoftQuantizerRegularization(total_gradient_steps=num_batches * n_epochs)
+    else:
+        return lambda m, e_reg: 0

model_compression_toolkit/{exporter/model_wrapper/keras/quantizers → gptq/pytorch/quantizer/soft_rounding}/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Sony Semiconductor Israel, Inc. All rights reserved.
+# Copyright 2023 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.

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

@@ -0,0 +1,115 @@
+# Copyright 2023 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 typing import List
+
+import torch
+import numpy as np
+from torch import nn
+
+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 model_compression_toolkit.quantizers_infrastructure 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:
+    """
+    A class to handle the computation of soft quantizer regularization for GPTQ training.
+    """
+
+    def __init__(self, total_gradient_steps: int):
+        """
+        Initializes the regularization computation object with a LinearDecay object.
+
+        Args:
+            total_gradient_steps: The number of gradient steps during optimization.
+        """
+
+        # Initializing the temperature decay according to the number of expected gradient steps
+        self.linear_decay = LinearTempDecay(total_gradient_steps)
+
+        self.count_iter = 0
+
+    def __call__(self, model: nn.Module, entropy_reg: float):
+        """
+        Returns the soft quantizer regularization value for SoftRounding.
+
+        Args:
+            model: A model to be quantized with SoftRounding.
+            entropy_reg: Entropy value to scale the quantizer regularization.
+
+        Returns: Regularization value.
+        """
+
+        soft_reg_aux: List[torch.Tensor] = []
+        for layer in model.modules():
+            if isinstance(layer, PytorchQuantizationWrapper):
+                kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),
+                                                                      fw_info=DEFAULT_PYTORCH_INFO)
+
+                st = layer.weights_quantizers[kernel_attribute].get_soft_targets()
+                b = self.linear_decay(self.count_iter)
+
+                soft_reg_aux.append((1 - torch.pow(torch.abs(st - .5) * 2, b)).sum())
+
+        reg = 0
+
+        for sq in soft_reg_aux:
+            reg += sq
+
+        self.count_iter += 1
+
+        return entropy_reg * reg