compressed-tensors-nightly 0.7.1.20241031__py3-none-any.whl → 0.7.1.20241102__py3-none-any.whl

compressed_tensors/quantization/cache.py

@@ -1,200 +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 enum import Enum
-from typing import Any, Dict, List, Optional, Tuple
-
-from compressed_tensors.quantization.observers import Observer
-from compressed_tensors.quantization.quant_args import QuantizationArgs
-from torch import Tensor
-from transformers import DynamicCache as HFDyanmicCache
-
-
-class KVCacheScaleType(Enum):
-    KEY = "k_scale"
-    VALUE = "v_scale"
-
-
-class QuantizedKVParameterCache(HFDyanmicCache):
-    """
-    Quantized KV cache used in the forward call based on HF's dynamic cache.
-    Quantization strategy (tensor, group, channel) set from Quantization arg's strategy
-    Singleton, so that the same cache gets reused in all forward call of self_attn.
-    Each time forward is called, .update() is called, and ._quantize(), ._dequantize()
-     gets called appropriately.
-    The size of tensor is
-     `[batch_size, num_heads, seq_len - residual_length, head_dim]`.
-
-
-    Triggered by adding kv_cache_scheme in the recipe.
-
-    Example:
-
-    ```python3
-    recipe = '''
-    quant_stage:
-        quant_modifiers:
-            QuantizationModifier:
-                kv_cache_scheme:
-                    num_bits: 8
-                    type: float
-                    strategy: tensor
-                    dynamic: false
-                    symmetric: true
-    '''
-
-    """
-
-    _instance = None
-    _initialized = False
-
-    def __new__(cls, *args, **kwargs):
-        """Singleton"""
-        if cls._instance is None:
-            cls._instance = super(QuantizedKVParameterCache, cls).__new__(cls)
-        return cls._instance
-
-    def __init__(self, quantization_args: QuantizationArgs):
-        if not self._initialized:
-            super().__init__()
-
-            self.quantization_args = quantization_args
-
-            self.k_observers: List[Observer] = []
-            self.v_observers: List[Observer] = []
-
-            # each index corresponds to layer_idx of the attention layer
-            self.k_scales: List[Tensor] = []
-            self.v_scales: List[Tensor] = []
-
-            self.k_zps: List[Tensor] = []
-            self.v_zps: List[Tensor] = []
-
-            self._initialized = True
-
-    def update(
-        self,
-        key_states: Tensor,
-        value_states: Tensor,
-        layer_idx: int,
-        cache_kwargs: Optional[Dict[str, Any]] = None,
-    ) -> Tuple[Tensor, Tensor]:
-        """
-        Get the k_scale and v_scale and output the
-         fakequant-ed key_states and value_states
-        """
-
-        if len(self.k_observers) <= layer_idx:
-            k_observer = self.quantization_args.get_observer()
-            v_observer = self.quantization_args.get_observer()
-
-            self.k_observers.append(k_observer)
-            self.v_observers.append(v_observer)
-
-        q_key_states = self._quantize(
-            key_states.contiguous(), KVCacheScaleType.KEY, layer_idx
-        )
-        q_value_states = self._quantize(
-            value_states.contiguous(), KVCacheScaleType.VALUE, layer_idx
-        )
-
-        qdq_key_states = self._dequantize(q_key_states, KVCacheScaleType.KEY, layer_idx)
-        qdq_value_states = self._dequantize(
-            q_value_states, KVCacheScaleType.VALUE, layer_idx
-        )
-
-        keys_to_return, values_to_return = qdq_key_states, qdq_value_states
-
-        return keys_to_return, values_to_return
-
-    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
-        """
-        Returns the sequence length of the cached states.
-        A layer index can be optionally passed.
-        """
-        if len(self.key_cache) <= layer_idx:
-            return 0
-        # since we cannot get the seq_length of each layer directly and
-        # rely on `_seen_tokens` which is updated every "layer_idx" == 0,
-        # this is a hack to get the actual seq_length for the given layer_idx
-        # this part of code otherwise fails when used to
-        # verify attn_weight shape in some models
-        return self._seen_tokens if layer_idx == 0 else self._seen_tokens - 1
-
-    def reset_states(self):
-        """reset the kv states (used in calibration)"""
-        self.key_cache: List[Tensor] = []
-        self.value_cache: List[Tensor] = []
-        # Used in `generate` to keep tally of how many tokens the cache has seen
-        self._seen_tokens = 0
-        self._quantized_key_cache: List[Tensor] = []
-        self._quantized_value_cache: List[Tensor] = []
-
-    def reset(self):
-        """
-        Reset the instantiation, create new instance on init
-        """
-        QuantizedKVParameterCache._instance = None
-        QuantizedKVParameterCache._initialized = False
-
-    def _quantize(self, tensor, kv_type, layer_idx):
-        """Quantizes a key/value using a defined quantization method."""
-        from compressed_tensors.quantization.lifecycle.forward import quantize
-
-        if kv_type == KVCacheScaleType.KEY:  # key type
-            observer = self.k_observers[layer_idx]
-            scales = self.k_scales
-            zps = self.k_zps
-        else:
-            assert kv_type == KVCacheScaleType.VALUE
-            observer = self.v_observers[layer_idx]
-            scales = self.v_scales
-            zps = self.v_zps
-
-        scale, zp = observer(tensor)
-        if len(scales) <= layer_idx:
-            scales.append(scale)
-            zps.append(zp)
-        else:
-            scales[layer_idx] = scale
-            zps[layer_idx] = scale
-
-        q_tensor = quantize(
-            x=tensor,
-            scale=scale,
-            zero_point=zp,
-            args=self.quantization_args,
-        )
-        return q_tensor
-
-    def _dequantize(self, qtensor, kv_type, layer_idx):
-        """Dequantizes back the tensor that was quantized by `self._quantize()`"""
-        from compressed_tensors.quantization.lifecycle.forward import dequantize
-
-        if kv_type == KVCacheScaleType.KEY:
-            scale = self.k_scales[layer_idx]
-            zp = self.k_zps[layer_idx]
-        else:
-            assert kv_type == KVCacheScaleType.VALUE
-            scale = self.v_scales[layer_idx]
-            zp = self.v_zps[layer_idx]
-
-        qdq_tensor = dequantize(
-            x_q=qtensor,
-            scale=scale,
-            zero_point=zp,
-            args=self.quantization_args,
-        )
-        return qdq_tensor

compressed_tensors/quantization/lifecycle/calibration.py

@@ -1,80 +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 compressed_tensors.quantization.quant_config import QuantizationStatus
-from compressed_tensors.utils import is_module_offloaded, update_parameter_data
-from torch.nn import Module
-
-
-__all__ = [
-    "set_module_for_calibration",
-]
-
-
-_LOGGER = logging.getLogger(__name__)
-
-
-def set_module_for_calibration(module: Module, quantize_weights_upfront: bool = True):
-    """
-    marks a layer as ready for calibration which activates observers
-    to update scales and zero points on each forward pass
-
-    apply to full model with `model.apply(set_module_for_calibration)`
-
-    :param module: module to set for calibration
-    :param quantize_weights_upfront: whether to automatically
-       run weight quantization at the start of calibration
-    """
-    if not getattr(module, "quantization_scheme", None):
-        # no quantization scheme nothing to do
-        return
-    status = getattr(module, "quantization_status", None)
-    if not status or status != QuantizationStatus.INITIALIZED:
-        _LOGGER.warning(
-            f"Attempting set module with status {status} to calibration mode. "
-            f"but status is not {QuantizationStatus.INITIALIZED} - you may "
-            "be calibrating an uninitialized module which may fail or attempting "
-            "to re-calibrate a frozen module"
-        )
-
-    if quantize_weights_upfront and module.quantization_scheme.weights is not None:
-        # set weight scale and zero_point up front, calibration data doesn't affect it
-        if not hasattr(module, "weight_observer"):
-            from compressed_tensors.quantization.lifecycle.initialize import (
-                initialize_observers,
-            )
-
-            initialize_observers(
-                module=module,
-                base_name="weight",
-                quantization_args=module.quantization_scheme.weights,
-            )
-
-        offloaded = is_module_offloaded(module)
-        if offloaded:
-            module._hf_hook.pre_forward(module)
-
-        observer = module.weight_observer
-        g_idx = getattr(module, "weight_g_idx", None)
-        scale, zero_point = observer(module.weight, g_idx=g_idx)
-        update_parameter_data(module, scale, "weight_scale")
-        update_parameter_data(module, zero_point, "weight_zero_point")
-
-        if offloaded:
-            module._hf_hook.post_forward(module, None)
-
-    module.quantization_status = QuantizationStatus.CALIBRATION

compressed_tensors/quantization/lifecycle/frozen.py

@@ -1,50 +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 compressed_tensors.quantization.quant_config import QuantizationStatus
-from torch.nn import Module
-
-
-__all__ = [
-    "freeze_module_quantization",
-]
-
-
-def freeze_module_quantization(module: Module):
-    """
-    deletes observers so static quantization is completed.
-
-    apply to full model with `model.apply(freeze_module_quantization)`
-
-    :param module: module to freeze quantization for
-    """
-    scheme = getattr(module, "quantization_scheme", None)
-    if not scheme:
-        # no quantization scheme nothing to do
-        return
-
-    if module.quantization_status == QuantizationStatus.FROZEN:
-        # nothing to do, already frozen
-        return
-
-    # delete observers from module if not dynamic
-    if hasattr(module, "input_observer") and not scheme.input_activations.dynamic:
-        delattr(module, "input_observer")
-    if hasattr(module, "weight_observer") and not scheme.weights.dynamic:
-        delattr(module, "weight_observer")
-    if hasattr(module, "output_observer") and not scheme.output_activations.dynamic:
-        delattr(module, "output_observer")
-
-    module.quantization_status = QuantizationStatus.FROZEN

compressed_tensors/quantization/observers/__init__.py

@@ -1,21 +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.
-
-# flake8: noqa
-# isort: skip_file
-
-from .helpers import *
-from .base import *
-from .min_max import *
-from .mse import *

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