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

compressed_tensors/quantization/lifecycle/initialize.py

@@ -17,17 +17,19 @@ import logging
 from typing import Optional
 
 import torch
-from accelerate.hooks import add_hook_to_module, remove_hook_from_module
-from accelerate.utils import PrefixedDataset
+from compressed_tensors.quantization.cache import KVCacheScaleType
 from compressed_tensors.quantization.lifecycle.forward import (
     wrap_module_forward_quantized,
+    wrap_module_forward_quantized_attn,
 )
 from compressed_tensors.quantization.quant_args import (
+    ActivationOrdering,
     QuantizationArgs,
     QuantizationStrategy,
 )
 from compressed_tensors.quantization.quant_config import QuantizationStatus
 from compressed_tensors.quantization.quant_scheme import QuantizationScheme
+from compressed_tensors.quantization.utils import is_kv_cache_quant_scheme
 from compressed_tensors.utils import get_execution_device, is_module_offloaded
 from torch.nn import Module, Parameter
 
@@ -43,6 +45,7 @@ _LOGGER = logging.getLogger(__name__)
 def initialize_module_for_quantization(
     module: Module,
     scheme: Optional[QuantizationScheme] = None,
+    force_zero_point: bool = True,
 ):
     """
     attaches appropriate scales, zero points, and observers to a layer
@@ -54,61 +57,93 @@ def initialize_module_for_quantization(
     :param scheme: scheme to use for quantization. if None is provided,
         will attempt to use scheme stored in the module under `quantization_scheme`,
         if not provided, the layer will be skipped
+    :param force_zero_point: whether to force initialization of a zero point for
+        symmetric quantization
     """
     scheme = scheme or getattr(module, "quantization_scheme", None)
     if scheme is None:
         # no scheme passed and layer not targeted for quantization - skip
         return
 
-    if scheme.input_activations is not None:
-        _initialize_scale_zero_point_observer(module, "input", scheme.input_activations)
-    if scheme.weights is not None:
-        if hasattr(module, "weight"):
-            weight_shape = None
-            if isinstance(module, torch.nn.Linear):
-                weight_shape = module.weight.shape
-            _initialize_scale_zero_point_observer(
-                module, "weight", scheme.weights, weight_shape=weight_shape
-            )
-        else:
-            _LOGGER.warning(
-                f"module type {type(module)} targeted for weight quantization but "
-                "has no attribute weight, skipping weight quantization "
-                f"for {type(module)}"
-            )
-    if scheme.output_activations is not None:
-        _initialize_scale_zero_point_observer(
-            module, "output", scheme.output_activations
-        )
+    if is_attention_module(module):
+        # wrap forward call of module to perform
+        # quantized actions based on calltime status
+        wrap_module_forward_quantized_attn(module, scheme)
+        _initialize_attn_scales(module)
 
-    module.quantization_scheme = scheme
-    module.quantization_status = QuantizationStatus.INITIALIZED
+    else:
 
-    offloaded = False
-    if is_module_offloaded(module):
-        offloaded = True
-        hook = module._hf_hook
-        prefix_dict = module._hf_hook.weights_map
-        new_prefix = {}
-
-        # recreate the prefix dict (since it is immutable)
-        # and add quantization parameters
-        for key, data in module.named_parameters():
-            if key not in prefix_dict:
-                new_prefix[f"{prefix_dict.prefix}{key}"] = data
+        if scheme.input_activations is not None:
+            _initialize_scale_zero_point_observer(
+                module,
+                "input",
+                scheme.input_activations,
+                force_zero_point=force_zero_point,
+            )
+        if scheme.weights is not None:
+            if hasattr(module, "weight"):
+                weight_shape = None
+                if isinstance(module, torch.nn.Linear):
+                    weight_shape = module.weight.shape
+                _initialize_scale_zero_point_observer(
+                    module,
+                    "weight",
+                    scheme.weights,
+                    weight_shape=weight_shape,
+                    force_zero_point=force_zero_point,
+                )
             else:
-                new_prefix[f"{prefix_dict.prefix}{key}"] = prefix_dict[key]
-        new_prefix_dict = PrefixedDataset(new_prefix, prefix_dict.prefix)
-        remove_hook_from_module(module)
-
-    # wrap forward call of module to perform quantized actions based on calltime status
-    wrap_module_forward_quantized(module, scheme)
-
-    if offloaded:
-        # we need to re-add the hook for offloading now that we've wrapped forward
-        add_hook_to_module(module, hook)
-        if prefix_dict is not None:
-            module._hf_hook.weights_map = new_prefix_dict
+                _LOGGER.warning(
+                    f"module type {type(module)} targeted for weight quantization but "
+                    "has no attribute weight, skipping weight quantization "
+                    f"for {type(module)}"
+                )
+
+        if scheme.output_activations is not None:
+            if not is_kv_cache_quant_scheme(scheme):
+                _initialize_scale_zero_point_observer(
+                    module, "output", scheme.output_activations
+                )
+
+        module.quantization_scheme = scheme
+        module.quantization_status = QuantizationStatus.INITIALIZED
+
+        offloaded = False
+        if is_module_offloaded(module):
+            try:
+                from accelerate.hooks import add_hook_to_module, remove_hook_from_module
+                from accelerate.utils import PrefixedDataset
+            except ModuleNotFoundError:
+                raise ModuleNotFoundError(
+                    "Offloaded model detected. To use CPU offloading with "
+                    "compressed-tensors the `accelerate` package must be installed, "
+                    "run `pip install compressed-tensors[accelerate]`"
+                )
+
+            offloaded = True
+            hook = module._hf_hook
+            prefix_dict = module._hf_hook.weights_map
+            new_prefix = {}
+
+            # recreate the prefix dict (since it is immutable)
+            # and add quantization parameters
+            for key, data in module.named_parameters():
+                if key not in prefix_dict:
+                    new_prefix[f"{prefix_dict.prefix}{key}"] = data
+                else:
+                    new_prefix[f"{prefix_dict.prefix}{key}"] = prefix_dict[key]
+            new_prefix_dict = PrefixedDataset(new_prefix, prefix_dict.prefix)
+            remove_hook_from_module(module)
+
+        # wrap forward call of module to perform
+        # quantized actions based on calltime status
+        wrap_module_forward_quantized(module, scheme)
+
+        if offloaded:
+            # we need to re-add the hook for offloading now that we've wrapped forward
+            add_hook_to_module(module, hook)
+            if prefix_dict is not None:
+                module._hf_hook.weights_map = new_prefix_dict
 
 
 def _initialize_scale_zero_point_observer(
@@ -116,6 +151,7 @@ def _initialize_scale_zero_point_observer(
     base_name: str,
     quantization_args: QuantizationArgs,
     weight_shape: Optional[torch.Size] = None,
+    force_zero_point: bool = True,
 ):
     # initialize observer module and attach as submodule
     observer = quantization_args.get_observer()
@@ -136,21 +172,68 @@ def _initialize_scale_zero_point_observer(
             # (output_channels, 1)
             expected_shape = (weight_shape[0], 1)
         elif quantization_args.strategy == QuantizationStrategy.GROUP:
+            num_groups = weight_shape[1] // quantization_args.group_size
             expected_shape = (
                 weight_shape[0],
-                weight_shape[1] // quantization_args.group_size,
+                max(num_groups, 1)
             )
 
-    # initializes empty scale and zero point parameters for the module
+    scale_dtype = module.weight.dtype
+    if scale_dtype not in [torch.float16, torch.bfloat16, torch.float32]:
+        scale_dtype = torch.float16
+
+    # initializes empty scale, zero point, and g_idx parameters for the module
     init_scale = Parameter(
-        torch.empty(expected_shape, dtype=module.weight.dtype, device=device),
+        torch.empty(expected_shape, dtype=scale_dtype, device=device),
         requires_grad=False,
     )
     module.register_parameter(f"{base_name}_scale", init_scale)
 
-    zp_dtype = quantization_args.pytorch_dtype()
-    init_zero_point = Parameter(
-        torch.empty(expected_shape, device=device, dtype=zp_dtype),
+    if force_zero_point or not quantization_args.symmetric:
+        zp_dtype = quantization_args.pytorch_dtype()
+        init_zero_point = Parameter(
+            torch.zeros(expected_shape, device=device, dtype=zp_dtype),
+            requires_grad=False,
+        )
+        module.register_parameter(f"{base_name}_zero_point", init_zero_point)
+
+    # only grouped activation ordering has g_idx
+    if quantization_args.actorder == ActivationOrdering.GROUP:
+        g_idx_shape = (weight_shape[1],)
+        g_idx_dtype = torch.int
+        init_g_idx = Parameter(
+            torch.full(g_idx_shape, -1, device=device, dtype=g_idx_dtype),
+            requires_grad=False,
+        )
+        module.register_parameter(f"{base_name}_g_idx", init_g_idx)
+
+
+def is_attention_module(module: Module):
+    return "attention" in module.__class__.__name__.lower() and (
+        hasattr(module, "k_proj")
+        or hasattr(module, "v_proj")
+        or hasattr(module, "qkv_proj")
+    )
+
+
+def _initialize_attn_scales(module: Module) -> None:
+    """Initlaize k_scale, v_scale for  self_attn"""
+
+    expected_shape = 1  # per tensor
+
+    param = next(module.parameters())
+    scale_dtype = param.dtype
+    device = param.device
+
+    init_scale = Parameter(
+        torch.empty(expected_shape, dtype=scale_dtype, device=device),
+        requires_grad=False,
+    )
+
+    module.register_parameter(KVCacheScaleType.KEY.value, init_scale)
+
+    init_scale = Parameter(
+        torch.empty(expected_shape, dtype=scale_dtype, device=device),
         requires_grad=False,
     )
-    module.register_parameter(f"{base_name}_zero_point", init_zero_point)
+    module.register_parameter(KVCacheScaleType.VALUE.value, init_scale)

compressed_tensors/quantization/observers/__init__.py

@@ -19,3 +19,4 @@ from .helpers import *
 from .base import *
 from .memoryless import *
 from .min_max import *
+from .mse import *

compressed_tensors/quantization/observers/base.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+from math import ceil
 from typing import Any, Iterable, Optional, Tuple, Union
 
 import torch
@@ -21,6 +22,7 @@ from compressed_tensors.quantization.quant_args import (
     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
 
@@ -46,15 +48,18 @@ class Observer(Module, RegistryMixin):
         self._num_observed_tokens = None
 
     @torch.no_grad()
-    def forward(self, observed: Tensor) -> Tuple[FloatTensor, IntTensor]:
+    def forward(
+        self, observed: Tensor, g_idx: Optional[Tensor] = None
+    ) -> Tuple[FloatTensor, IntTensor]:
         """
         maps directly to get_qparams
-        :param observed: optional observed tensor to calculate quantization parameters
-            from
+        :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)
+        return self.get_qparams(observed=observed, g_idx=g_idx)
 
     def calculate_qparams(
         self,
@@ -77,7 +82,9 @@ class Observer(Module, RegistryMixin):
         ...
 
     def get_qparams(
-        self, observed: Optional[Tensor] = None
+        self,
+        observed: Optional[Tensor] = None,
+        g_idx: Optional[Tensor] = None,
     ) -> Tuple[FloatTensor, IntTensor]:
         """
         Convenience function to wrap overwritten calculate_qparams
@@ -86,6 +93,7 @@ class Observer(Module, RegistryMixin):
 
         :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:
@@ -97,20 +105,42 @@ class Observer(Module, RegistryMixin):
                 self._scale, self._zero_point = self.calculate_qparams(observed)
 
             elif self.quantization_args.strategy == QuantizationStrategy.GROUP:
+                rows = observed.shape[0]
                 columns = observed.shape[1]
-                scales, zero_points = [], []
-                group_idxs = range(0, columns, self.quantization_args.group_size)
-                for group_id, group_idx in enumerate(group_idxs):
+                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[:, group_idx : (group_idx + group_size)],
+                        observed[:, start:end],
                         0,
-                        tensor_id=group_id,
+                        tensor_id=group_index,
                     )
-                    scales.append(scale)
-                    zero_points.append(zero_point)
 
-                self._scale = torch.cat(scales, dim=1, out=self._scale)
-                self._zero_point = torch.cat(zero_points, dim=1, out=self._zero_point)
+                    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
@@ -132,6 +162,8 @@ class Observer(Module, RegistryMixin):
         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)
@@ -171,3 +203,11 @@ class Observer(Module, RegistryMixin):
         # 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/min_max.py

@@ -94,3 +94,11 @@ class MovingAverageMinMaxObserver(Observer):
         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

@@ -0,0 +1,162 @@
+# 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 = {}