compressed-tensors 0.7.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

compressed_tensors/quantization/observers/base.py

@@ -1,213 +0,0 @@
-# Copyright (c) 2021 - present / Neuralmagic, 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 logging
-from math import ceil
-from typing import Any, Iterable, Optional, Tuple, Union
-
-import torch
-from compressed_tensors.quantization.quant_args import (
-    QuantizationArgs,
-    QuantizationStrategy,
-)
-from compressed_tensors.registry.registry import RegistryMixin
-from compressed_tensors.utils import safe_permute
-from torch import FloatTensor, IntTensor, Tensor
-from torch.nn import Module
-
-
-_LOGGER = logging.getLogger(__name__)
-
-
-__all__ = ["Observer"]
-
-
-class Observer(Module, RegistryMixin):
-    """
-    Base Observer class to be subclassed for specific implementation.
-    Subclasses should override `calculate_qparams` to return a scale, zero_point
-    pair
-    """
-
-    def __init__(self, quantization_args: QuantizationArgs):
-        self.quantization_args: QuantizationArgs = quantization_args
-        super().__init__()
-        self._scale = None
-        self._zero_point = None
-        self._num_observed_tokens = None
-
-    @torch.no_grad()
-    def forward(
-        self, observed: Tensor, g_idx: Optional[Tensor] = None
-    ) -> Tuple[FloatTensor, IntTensor]:
-        """
-        maps directly to get_qparams
-        :param observed: optional observed tensor from which to calculate
-            quantization parameters
-        :param g_idx: optional mapping from column index to group index
-        :return: tuple of scale and zero point based on last observed value
-        """
-        self.record_observed_tokens(observed)
-        return self.get_qparams(observed=observed, g_idx=g_idx)
-
-    def calculate_qparams(
-        self,
-        observed: Tensor,
-        reduce_dims: Optional[Tuple[int]] = None,
-    ) -> Tuple[FloatTensor, IntTensor]:
-        """
-        :param observed: observed tensor to calculate quantization parameters for
-        :param reduce_dims: optional tuple of dimensions to reduce along,
-            returned scale and zero point will be shaped (1,) along the
-            reduced dimensions
-        :return: tuple of scale and zero point derived from the observed tensor
-        """
-        raise NotImplementedError(f"{self.__class__} must implement calculate_qparams")
-
-    def post_calculate_qparams(self) -> None:
-        """
-        Run any logic specific to its observers after running calculate_qparams
-        """
-        ...
-
-    def get_qparams(
-        self,
-        observed: Optional[Tensor] = None,
-        g_idx: Optional[Tensor] = None,
-    ) -> Tuple[FloatTensor, IntTensor]:
-        """
-        Convenience function to wrap overwritten calculate_qparams
-        adds support to make observed tensor optional and support for tracking latest
-        calculated scale and zero point
-
-        :param observed: optional observed tensor to calculate quantization parameters
-            from
-        :param g_idx: optional mapping from column index to group index
-        :return: tuple of scale and zero point based on last observed value
-        """
-        if observed is not None:
-            group_size = self.quantization_args.group_size
-
-            if self.quantization_args.strategy == QuantizationStrategy.TENSOR:
-
-                # re-calculate scale and zero point, update the stored value
-                self._scale, self._zero_point = self.calculate_qparams(observed)
-
-            elif self.quantization_args.strategy == QuantizationStrategy.GROUP:
-                rows = observed.shape[0]
-                columns = observed.shape[1]
-                num_groups = int(ceil(columns / group_size))
-                self._scale = torch.empty(
-                    (rows, num_groups), dtype=observed.dtype, device=observed.device
-                )
-                zp_dtype = self.quantization_args.pytorch_dtype()
-                self._zero_point = torch.empty(
-                    (rows, num_groups), dtype=zp_dtype, device=observed.device
-                )
-
-                # support column-order (default) quantization as well as other orderings
-                # such as activation ordering. Below checks if g_idx has initialized
-                is_column_order = g_idx is None or -1 in g_idx
-                if is_column_order:
-                    group_sizes = torch.full((num_groups,), group_size, dtype=torch.int)
-                else:
-                    group_indices, group_sizes = torch.unique(g_idx, return_counts=True)
-                    group_sizes = group_sizes[torch.argsort(group_indices)]
-
-                    perm = torch.argsort(g_idx)
-                    observed = safe_permute(observed, perm, dim=1)
-
-                # TODO: experiment with vectorizing for loop for performance
-                end = 0
-                for group_index, group_count in enumerate(group_sizes):
-                    start = end
-                    end = start + group_count
-                    scale, zero_point = self.get_qparams_along_dim(
-                        observed[:, start:end],
-                        0,
-                        tensor_id=group_index,
-                    )
-
-                    self._scale[:, group_index] = scale.squeeze(1)
-                    self._zero_point[:, group_index] = zero_point.squeeze(1)
-
-            elif self.quantization_args.strategy == QuantizationStrategy.CHANNEL:
-                # assume observed is transposed, because its the output, hence use dim 0
-                self._scale, self._zero_point = self.get_qparams_along_dim(observed, 0)
-
-            elif self.quantization_args.strategy == QuantizationStrategy.TOKEN:
-                # use dim 1, assume the obsersed.shape = [batch, token, hidden]
-                # should be batch, token
-                self._scale, self._zero_point = self.get_qparams_along_dim(
-                    observed,
-                    dim={0, 1},
-                )
-
-        return self._scale, self._zero_point
-
-    def get_qparams_along_dim(
-        self,
-        observed,
-        dim: Union[int, Iterable[int]],
-        tensor_id: Optional[Any] = None,
-    ):
-        if isinstance(dim, int):
-            dim = [dim]
-        dim = set(dim)
-
-        reduce_dims = tuple(idx for idx in range(observed.ndim) if idx not in dim)
-        return self.calculate_qparams(
-            observed, reduce_dims=reduce_dims, tensor_id=tensor_id
-        )
-
-    def record_observed_tokens(self, batch_tensor: Tensor):
-        """
-        Counts the number of tokens observed during the
-        forward passes. The count is aggregated in the
-        _num_observed_tokens attribute of the class.
-
-        Note: The batch_tensor is expected to have two dimensions
-            (batch_size * sequence_length, num_features). This is the
-            general shape expected by the forward pass of the expert
-            layers in a MOE model. If the input tensor does not have
-            two dimensions, the _num_observed_tokens attribute will be set
-            to None.
-        """
-        if not isinstance(batch_tensor, Tensor):
-            raise ValueError(f"Expected value to be a tensor, got {type(batch_tensor)}")
-
-        if batch_tensor.ndim != 2:
-            _LOGGER.debug(
-                "The input tensor is expected to have two dimensions "
-                "(batch_size * sequence_length, num_features). "
-                f"The input tensor has {batch_tensor.ndim} dimensions."
-            )
-            return
-
-        if self._num_observed_tokens is None:
-            # initialize the count
-            self._num_observed_tokens = 0
-
-        # batch_tensor (batch_size * sequence_length, num_features)
-        # observed_tokens (batch_size * sequence_length)
-        observed_tokens, _ = batch_tensor.shape
-        self._num_observed_tokens += observed_tokens
-
-    def reset(self):
-        """
-        Reset the state of the observer
-        """
-        self._num_observed_tokens = None
-        self._scale = None
-        self._zero_point = None

compressed_tensors/quantization/observers/helpers.py

@@ -1,111 +0,0 @@
-# Copyright (c) 2021 - present / Neuralmagic, 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 collections import Counter
-from typing import Tuple
-
-import torch
-from compressed_tensors.quantization.quant_args import (
-    FP8_DTYPE,
-    QuantizationArgs,
-    QuantizationType,
-)
-from torch import FloatTensor, IntTensor, Tensor
-
-
-__all__ = ["calculate_qparams", "get_observer_token_count", "calculate_range"]
-
-
-def get_observer_token_count(module: torch.nn.Module) -> Counter:
-    """
-    Parse the module and return the number of tokens observed by
-    each module's observer.
-
-    :param module: module to parse
-    :return: counter with the number of tokens observed by each observer
-    """
-    token_counts = Counter()
-    for name, module in module.named_modules():
-        if name.endswith(".input_observer"):
-            token_counts[
-                name.replace(".input_observer", "")
-            ] = module._num_observed_tokens
-    return token_counts
-
-
-def calculate_qparams(
-    min_vals: Tensor, max_vals: Tensor, quantization_args: QuantizationArgs
-) -> Tuple[FloatTensor, IntTensor]:
-    """
-    :param min_vals: tensor of min value(s) to calculate scale(s) and zero point(s)
-        from
-    :param max_vals: tensor of max value(s) to calculate scale(s) and zero point(s)
-        from
-    :param quantization_args: settings to quantization
-    :return: tuple of the calculated scale(s) and zero point(s)
-    """
-    min_vals = torch.min(min_vals, torch.zeros_like(min_vals))
-    max_vals = torch.max(max_vals, torch.zeros_like(max_vals))
-    device = min_vals.device
-
-    bit_min, bit_max = calculate_range(quantization_args, device)
-    bit_range = bit_max - bit_min
-    zp_dtype = quantization_args.pytorch_dtype()
-
-    if quantization_args.symmetric:
-        max_val_pos = torch.max(torch.abs(min_vals), torch.abs(max_vals))
-        scales = max_val_pos / (float(bit_range) / 2)
-        scales = torch.clamp(scales, min=torch.finfo(torch.float32).eps)
-        zero_points = torch.zeros(scales.shape, device=device, dtype=min_vals.dtype)
-    else:
-        scales = (max_vals - min_vals) / float(bit_range)
-        scales = torch.clamp(scales, min=torch.finfo(torch.float32).eps)
-        zero_points = bit_min - (min_vals / scales)
-        zero_points = torch.clamp(zero_points, bit_min, bit_max)
-
-    # match zero-points to quantized type
-    zero_points = zero_points.to(zp_dtype)
-
-    if scales.ndim == 0:
-        scales = scales.reshape(1)
-        zero_points = zero_points.reshape(1)
-
-    return scales, zero_points
-
-
-def calculate_range(quantization_args: QuantizationArgs, device: str) -> Tuple:
-    """
-    Calculated the effective quantization range for the given Quantization Args
-
-    :param quantization_args: quantization args to get range of
-    :param device: device to store the range to
-    :return: tuple endpoints for the given quantization range
-    """
-    if quantization_args.type == QuantizationType.INT:
-        bit_range = 2**quantization_args.num_bits
-        q_max = torch.tensor(bit_range / 2 - 1, device=device)
-        q_min = torch.tensor(-bit_range / 2, device=device)
-    elif quantization_args.type == QuantizationType.FLOAT:
-        if quantization_args.num_bits != 8:
-            raise ValueError(
-                "Floating point quantization is only supported for 8 bits,"
-                f"got {quantization_args.num_bits}"
-            )
-        fp_range_info = torch.finfo(FP8_DTYPE)
-        q_max = torch.tensor(fp_range_info.max, device=device)
-        q_min = torch.tensor(fp_range_info.min, device=device)
-    else:
-        raise ValueError(f"Invalid quantization type {quantization_args.type}")
-
-    return q_min, q_max

compressed_tensors/quantization/observers/memoryless.py

@@ -1,56 +0,0 @@
-# Copyright (c) 2021 - present / Neuralmagic, 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 Any, Optional, Tuple
-
-import torch
-from compressed_tensors.quantization.observers.base import Observer
-from compressed_tensors.quantization.observers.helpers import calculate_qparams
-from torch import FloatTensor, IntTensor, Tensor
-
-
-__all__ = ["MemorylessObserver"]
-
-
-@Observer.register("memoryless", alias=["dynamic"])
-class MemorylessObserver(Observer):
-    """
-    Implements a quantization observer that sets the scale and
-    zero point based on the latest observed value without tracking state
-    """
-
-    def calculate_qparams(
-        self,
-        observed: Tensor,
-        tensor_id: Optional[Any] = None,
-        reduce_dims: Optional[Tuple[int]] = None,
-    ) -> Tuple[FloatTensor, IntTensor]:
-        """
-        Returns the min and max values of observed tensor
-
-        :param observed: observed tensor to calculate quantization parameters for
-        :param tensor_id: optional id for tensor; not used for memoryless
-        :param reduce_dims: optional tuple of dimensions to reduce along,
-            returned scale and zero point will be shaped (1,) along the
-            reduced dimensions
-        :return: tuple of scale and zero point derived from the observed tensor
-        """
-
-        if not reduce_dims:
-            min_val, max_val = torch.aminmax(observed)
-        else:
-            min_val = torch.amin(observed, dim=reduce_dims, keepdims=True)
-            max_val = torch.amax(observed, dim=reduce_dims, keepdims=True)
-
-        return calculate_qparams(min_val, max_val, self.quantization_args)

compressed_tensors/quantization/observers/min_max.py

@@ -1,104 +0,0 @@
-# Copyright (c) 2021 - present / Neuralmagic, 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 Any, Optional, Tuple
-
-import torch
-from compressed_tensors.quantization.observers.base import Observer
-from compressed_tensors.quantization.observers.helpers import calculate_qparams
-from compressed_tensors.quantization.quant_args import QuantizationArgs
-from torch import FloatTensor, IntTensor, Tensor
-
-
-__all__ = ["MovingAverageMinMaxObserver"]
-
-
-@Observer.register("minmax")
-class MovingAverageMinMaxObserver(Observer):
-    """
-    Implements a dynamic quantization observer that sets the scale and
-    zero point based on a moving average of the overall min and max observed values
-    """
-
-    def __init__(
-        self, quantization_args: QuantizationArgs, averaging_constant: float = 0.01
-    ):
-        super().__init__(quantization_args=quantization_args)
-
-        self.min_val = {}
-        self.max_val = {}
-        self.averaging_constant = averaging_constant
-
-    def calculate_qparams(
-        self,
-        observed: Tensor,
-        reduce_dims: Optional[Tuple[int]] = None,
-        tensor_id: Optional[Any] = None,
-    ) -> Tuple[FloatTensor, IntTensor]:
-        """
-        Updates the observed min and max using a moving average smoothed by the
-        averaging_constant
-
-        :param observed: observed tensor to calculate quantization parameters for
-        :param reduce_dims: optional tuple of dimensions to reduce along,
-            returned scale and zero point will be shaped (1,) along the
-            reduced dimensions
-        :param tensor_id: Optional id if different ranges of observed tensors are
-            passed, useful for sharding tensors by group_size
-        :return: tuple of scale and zero point derived from the observed tensor
-        """
-        tensor_id = tensor_id or "default"
-
-        if not reduce_dims:
-            min_val, max_val = torch.aminmax(observed)
-        else:
-            min_val = torch.amin(observed, dim=reduce_dims, keepdims=True)
-            max_val = torch.amax(observed, dim=reduce_dims, keepdims=True)
-
-        running_min_val = self.min_val.get(tensor_id, None)
-        running_max_val = self.max_val.get(tensor_id, None)
-
-        if running_min_val is None or running_max_val is None:
-            updated_min_val = min_val
-            updated_max_val = max_val
-        else:
-            updated_min_val = running_min_val + self.averaging_constant * (
-                min_val - running_min_val
-            )
-            updated_max_val = running_max_val + self.averaging_constant * (
-                max_val - running_max_val
-            )
-
-        self.min_val[tensor_id] = updated_min_val
-        self.max_val[tensor_id] = updated_max_val
-
-        return calculate_qparams(
-            updated_min_val, updated_max_val, self.quantization_args
-        )
-
-    def get_qparams_along_dim(
-        self, observed, dim: int, tensor_id: Optional[Any] = None
-    ):
-        reduce_dims = tuple(idx for idx in range(observed.ndim) if idx != dim)
-        return self.calculate_qparams(
-            observed, reduce_dims=reduce_dims, tensor_id=tensor_id
-        )
-
-    def reset(self):
-        """
-        Reset the state of the observer, including min and maximum values
-        """
-        super().reset()
-        self.min_val = {}
-        self.max_val = {}

compressed_tensors/quantization/observers/mse.py

@@ -1,162 +0,0 @@
-# Copyright (c) 2021 - present / Neuralmagic, 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 Any, Optional, Tuple
-
-import torch
-from compressed_tensors.quantization.observers.base import Observer
-from compressed_tensors.quantization.observers.helpers import calculate_qparams
-from compressed_tensors.quantization.quant_args import QuantizationArgs
-from torch import FloatTensor, IntTensor, Tensor
-
-
-__all__ = ["MovingAverageMSEObserver"]
-
-
-@Observer.register("mse")
-class MovingAverageMSEObserver(Observer):
-    """
-    Implements a dynamic quantization observer that sets the scale and
-    zero point based on a moving average of the mse-clipped min and max observed values
-    """
-
-    def __init__(
-        self,
-        quantization_args: QuantizationArgs,
-        averaging_constant: float = 0.01,
-        grid: float = 100.0,
-        maxshrink: float = 0.80,
-        norm: float = 2.4,
-    ):
-        super().__init__(quantization_args=quantization_args)
-
-        self.min_val = {}
-        self.max_val = {}
-        self.averaging_constant = averaging_constant
-        self.grid = grid
-        self.maxshrink = maxshrink
-        self.norm = norm
-
-    def calculate_mse_min_max(
-        self,
-        observed: Tensor,
-        reduce_dims: Optional[Tuple[int]] = None,
-    ):
-        """
-        Computes the mse-clipped min and max values of the observed tensor by
-        optimizing for quantization error
-
-        :param observed: observed tensor to calculate quantization parameters for
-        :param reduce_dims: optional tuple of dimensions to reduce along,
-            returned values will be shaped (1,) along the reduced dimensions
-        :return: tuple of min and max values derived from the observed tensor
-        """
-        from compressed_tensors.quantization.lifecycle import fake_quantize
-
-        if not reduce_dims:
-            absolute_min_val, absolute_max_val = torch.aminmax(observed)
-        else:
-            absolute_min_val = torch.amin(observed, dim=reduce_dims, keepdims=True)
-            absolute_max_val = torch.amax(observed, dim=reduce_dims, keepdims=True)
-
-        best = torch.full(absolute_min_val.shape, float("inf"))
-        min_val = torch.ones(absolute_min_val.shape)
-        max_val = torch.zeros(absolute_max_val.shape)
-        for i in range(int(self.maxshrink * self.grid)):
-            p = 1 - i / self.grid
-            shrinked_min_val = p * absolute_min_val
-            shrinked_max_val = p * absolute_max_val
-
-            candidate_scales, candidate_zero_points = calculate_qparams(
-                shrinked_min_val, shrinked_max_val, self.quantization_args
-            )
-            q = fake_quantize(
-                observed,
-                candidate_scales,
-                candidate_zero_points,
-                self.quantization_args,
-            )
-
-            q -= observed
-            q.abs_()
-            q.pow_(self.norm)
-            if not reduce_dims:
-                err = torch.sum(q)
-            else:
-                err = torch.sum(q, reduce_dims, keepdims=True)
-
-            tmp = err < best
-            if torch.any(tmp):
-                best[tmp] = err[tmp]
-                min_val[tmp] = shrinked_min_val[tmp]
-                max_val[tmp] = shrinked_max_val[tmp]
-        return min_val, max_val
-
-    def calculate_qparams(
-        self,
-        observed: Tensor,
-        reduce_dims: Optional[Tuple[int]] = None,
-        tensor_id: Optional[Any] = None,
-    ) -> Tuple[FloatTensor, IntTensor]:
-        """
-        Updates the mse-clipped min and max values of the observed tensor using
-        a moving average smoothed by the averaging_constant
-
-        :param observed: observed tensor to calculate quantization parameters for
-        :param reduce_dims: optional tuple of dimensions to reduce along,
-            returned scale and zero point will be shaped (1,) along the
-            reduced dimensions
-        :param tensor_id: Optional id if different ranges of observed tensors are
-            passed, useful for sharding tensors by group_size
-        :return: tuple of scale and zero point derived from the observed tensor
-        """
-        min_val, max_val = self.calculate_mse_min_max(observed, reduce_dims)
-
-        running_min_val = self.min_val.get(tensor_id, None)
-        running_max_val = self.max_val.get(tensor_id, None)
-
-        if running_min_val is None or running_max_val is None:
-            updated_min_val = min_val
-            updated_max_val = max_val
-        else:
-            updated_min_val = running_min_val + self.averaging_constant * (
-                min_val - running_min_val
-            )
-            updated_max_val = running_max_val + self.averaging_constant * (
-                max_val - running_max_val
-            )
-
-        tensor_id = tensor_id or "default"
-        self.min_val[tensor_id] = updated_min_val
-        self.max_val[tensor_id] = updated_max_val
-
-        return calculate_qparams(
-            updated_min_val, updated_max_val, self.quantization_args
-        )
-
-    def get_qparams_along_dim(
-        self, observed, dim: int, tensor_id: Optional[Any] = None
-    ):
-        reduce_dims = tuple(idx for idx in range(observed.ndim) if idx != dim)
-        return self.calculate_qparams(
-            observed, reduce_dims=reduce_dims, tensor_id=tensor_id
-        )
-
-    def reset(self):
-        """
-        Reset the state of the observer, including min and maximum values
-        """
-        super().reset()
-        self.min_val = {}
-        self.max_val = {}