compressed-tensors-nightly 0.5.0.20240814__py3-none-any.whl → 0.5.0.20240830__py3-none-any.whl

compressed_tensors/quantization/lifecycle/forward.py

@@ -25,7 +25,7 @@ from compressed_tensors.quantization.quant_args import (
 )
 from compressed_tensors.quantization.quant_config import QuantizationStatus
 from compressed_tensors.quantization.quant_scheme import QuantizationScheme
-from compressed_tensors.utils import update_parameter_data
+from compressed_tensors.utils import safe_permute, update_parameter_data
 from torch.nn import Module
 
 
@@ -45,6 +45,7 @@ def quantize(
     zero_point: torch.Tensor,
     args: QuantizationArgs,
     dtype: Optional[torch.dtype] = None,
+    g_idx: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """
     Quantize the input tensor x using the QuantizationStrategy specified in args.
@@ -58,16 +59,9 @@ def quantize(
     :param zero_point: zero point tensor
     :param args: quantization args dictating how to quantize x
     :param dtype: optional dtype to cast the quantized output to
+    :param g_idx: optional mapping from column index to group index
     :return: fake quantized tensor
     """
-    # ensure all tensors are on the same device
-    # assumes that the target device is the input
-    # tensor's device
-    if x.device != scale.device:
-        scale = scale.to(x.device)
-    if x.device != zero_point.device:
-        zero_point = zero_point.to(x.device)
-
     return _process_quantization(
         x=x,
         scale=scale,
@@ -76,6 +70,7 @@ def quantize(
         dtype=dtype,
         do_quantize=True,
         do_dequantize=False,
+        g_idx=g_idx,
     )
 
 
@@ -86,6 +81,7 @@ def dequantize(
     zero_point: torch.Tensor = None,
     args: QuantizationArgs = None,
     dtype: Optional[torch.dtype] = None,
+    g_idx: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """
     Dequantize a quantized input tensor x_q based on the strategy specified in args. If
@@ -96,6 +92,7 @@ def dequantize(
     :param zero_point: zero point tensor
     :param args: quantization args used to quantize x_q
     :param dtype: optional dtype to cast the dequantized output to
+    :param g_idx: optional mapping from column index to group index
     :return: dequantized float tensor
     """
     if args is None:
@@ -126,6 +123,7 @@ def dequantize(
         do_quantize=False,
         do_dequantize=True,
         dtype=dtype,
+        g_idx=g_idx,
     )
 
 
@@ -135,6 +133,7 @@ def fake_quantize(
     scale: torch.Tensor,
     zero_point: torch.Tensor,
     args: QuantizationArgs,
+    g_idx: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """
     Fake quantize the input tensor x by quantizing then dequantizing with
@@ -147,6 +146,7 @@ def fake_quantize(
     :param scale: scale tensor
     :param zero_point: zero point tensor
     :param args: quantization args dictating how to quantize x
+    :param g_idx: optional mapping from column index to group index
     :return: fake quantized tensor
     """
     return _process_quantization(
@@ -156,6 +156,7 @@ def fake_quantize(
         args=args,
         do_quantize=True,
         do_dequantize=True,
+        g_idx=g_idx,
     )
 
 
@@ -164,21 +165,19 @@ def _process_quantization(
     x: torch.Tensor,
     scale: torch.Tensor,
     zero_point: torch.Tensor,
+    g_idx: Optional[torch.Tensor],
     args: QuantizationArgs,
     dtype: Optional[torch.dtype] = None,
     do_quantize: bool = True,
     do_dequantize: bool = True,
 ) -> torch.Tensor:
-
     q_min, q_max = calculate_range(args, x.device)
     group_size = args.group_size
 
     if args.strategy == QuantizationStrategy.GROUP:
         output_dtype = dtype if dtype is not None else x.dtype
         output = torch.zeros_like(x).to(output_dtype)
-
-        # TODO: vectorize the for loop
-        # TODO: fix genetric assumption about the tensor size for computing group
+        columns = output.shape[1]
 
         # TODO: make validation step for inputs
 
@@ -187,23 +186,38 @@ def _process_quantization(
             scale = scale.unsqueeze(1)
             zero_point = zero_point.unsqueeze(1) if zero_point is not None else None
 
-        columns = x.shape[1]
         if columns >= group_size:
             if columns % group_size != 0:
                 raise ValueError(
-                    "tesnor column shape must be divisble "
+                    "tensor column shape must be divisble "
                     f"by the given group_size {group_size}"
                 )
-        for i in range(ceil(columns / group_size)):
-            # scale.shape should be [nchan, ndim]
-            # sc.shape should be [nchan, 1] after unsqueeze
-            sc = scale[:, i].view(-1, 1)
-            zp = zero_point[:, i].view(-1, 1) if zero_point is not None else None
 
-            idx = i * group_size
+        # support column-order (default) quantization as well as other orderings
+        # such as activation ordering. Below checks if g_idx has been initialized
+        is_column_order = g_idx is None or -1 in g_idx
+        if is_column_order:
+            num_groups = int(ceil(columns / group_size))
+            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)
+            x = safe_permute(x, perm, dim=1)
+
+        # TODO: experiment with vectorizing for loop for performance
+        end = 0
+        for index, group_count in enumerate(group_sizes):
+            sc = scale[:, index].view(-1, 1)
+            zp = zero_point[:, index].view(-1, 1) if zero_point is not None else None
+
+            start = end
+            end = start + group_count
             if do_quantize:
-                output[:, idx : (idx + group_size)] = _quantize(
-                    x[:, idx : (idx + group_size)],
+                output[:, start:end] = _quantize(
+                    x[:, start:end],
                     sc,
                     zp,
                     q_min,
@@ -211,13 +225,13 @@ def _process_quantization(
                     args,
                     dtype=dtype,
                 )
+
             if do_dequantize:
-                input = (
-                    output[:, idx : (idx + group_size)]
-                    if do_quantize
-                    else x[:, idx : (idx + group_size)]
-                )
-                output[:, idx : (idx + group_size)] = _dequantize(input, sc, zp)
+                input = output[:, start:end] if do_quantize else x[:, start:end]
+                output[:, start:end] = _dequantize(input, sc, zp)
+
+        if not is_column_order:
+            output = safe_permute(output, torch.argsort(perm), dim=1)
 
     else:  # covers channel, token and tensor strategies
         if do_quantize:
@@ -252,6 +266,7 @@ def wrap_module_forward_quantized(module: Module, scheme: QuantizationScheme):
             return forward_func_orig.__get__(module, module.__class__)(*args, **kwargs)
 
         input_ = args[0]
+        compressed = module.quantization_status == QuantizationStatus.COMPRESSED
 
         if scheme.input_activations is not None:
             # calibrate and (fake) quantize input activations when applicable
@@ -259,7 +274,7 @@ def wrap_module_forward_quantized(module: Module, scheme: QuantizationScheme):
                 module, input_, "input", scheme.input_activations
             )
 
-        if scheme.weights is not None:
+        if scheme.weights is not None and not compressed:
             # calibrate and (fake) quantize weights when applicable
             unquantized_weight = self.weight.data.clone()
             self.weight.data = maybe_calibrate_or_quantize(
@@ -278,7 +293,7 @@ def wrap_module_forward_quantized(module: Module, scheme: QuantizationScheme):
             )
 
         # restore back to unquantized_value
-        if scheme.weights is not None:
+        if scheme.weights is not None and not compressed:
             self.weight.data = unquantized_weight
 
         return output
@@ -292,11 +307,16 @@ def wrap_module_forward_quantized(module: Module, scheme: QuantizationScheme):
 def maybe_calibrate_or_quantize(
     module: Module, value: torch.Tensor, base_name: str, args: "QuantizationArgs"
 ) -> torch.Tensor:
-    # only run quantized for the included stages
-    if module.quantization_status not in {
-        QuantizationStatus.CALIBRATION,
-        QuantizationStatus.FROZEN,
-    }:
+    # don't run quantization if we haven't entered calibration mode
+    if module.quantization_status == QuantizationStatus.INITIALIZED:
+        return value
+
+    # in compressed mode, the weight is already compressed and quantized so we don't
+    # need to run fake quantization
+    if (
+        module.quantization_status == QuantizationStatus.COMPRESSED
+        and base_name == "weight"
+    ):
         return value
 
     if value.numel() == 0:
@@ -304,14 +324,16 @@ def maybe_calibrate_or_quantize(
         # skip quantization
         return value
 
+    g_idx = getattr(module, "weight_g_idx", None)
+
     if args.dynamic:
         # dynamic quantization - get scale and zero point directly from observer
         observer = getattr(module, f"{base_name}_observer")
-        scale, zero_point = observer(value)
+        scale, zero_point = observer(value, g_idx=g_idx)
     else:
         # static quantization - get previous scale and zero point from layer
         scale = getattr(module, f"{base_name}_scale")
-        zero_point = getattr(module, f"{base_name}_zero_point")
+        zero_point = getattr(module, f"{base_name}_zero_point", None)
 
         if (
             module.quantization_status == QuantizationStatus.CALIBRATION
@@ -320,13 +342,13 @@ def maybe_calibrate_or_quantize(
             # calibration mode - get new quant params from observer
             observer = getattr(module, f"{base_name}_observer")
 
-            updated_scale, updated_zero_point = observer(value)
+            updated_scale, updated_zero_point = observer(value, g_idx=g_idx)
 
             # update scale and zero point
             update_parameter_data(module, updated_scale, f"{base_name}_scale")
             update_parameter_data(module, updated_zero_point, f"{base_name}_zero_point")
 
-    return fake_quantize(value, scale, zero_point, args)
+    return fake_quantize(value, scale, zero_point, args, g_idx=g_idx)
 
 
 @torch.no_grad()
@@ -340,7 +362,9 @@ def _quantize(
     dtype: Optional[torch.dtype] = None,
 ) -> torch.Tensor:
 
-    scaled = x / scale + zero_point.to(x.dtype)
+    scaled = x / scale
+    if zero_point is not None:
+        scaled += zero_point.to(x.dtype)
     # clamp first because cast isn't guaranteed to be saturated (ie for fp8)
     clamped_value = torch.clamp(
         scaled,
@@ -361,11 +385,11 @@ def _dequantize(
     zero_point: torch.Tensor = None,
     dtype: Optional[torch.dtype] = None,
 ) -> torch.Tensor:
+    dequant_value = x_q.to(scale.dtype)
 
-    dequant_value = x_q
     if zero_point is not None:
         dequant_value = dequant_value - zero_point.to(scale.dtype)
-    dequant_value = dequant_value.to(scale.dtype) * scale
+    dequant_value = dequant_value * scale
 
     if dtype is not None:
         dequant_value = dequant_value.to(dtype)

compressed_tensors/quantization/lifecycle/helpers.py

@@ -16,7 +16,9 @@
 Miscelaneous helpers for the quantization lifecycle
 """
 
+from typing import Optional
 
+import torch
 from torch.nn import Module
 
 
@@ -27,16 +29,41 @@ __all__ = [
 ]
 
 
-def update_layer_weight_quant_params(layer: Module):
-    weight = getattr(layer, "weight", None)
+def update_layer_weight_quant_params(
+    layer: Module,
+    weight: Optional[torch.Tensor] = None,
+    g_idx: Optional[torch.Tensor] = None,
+    reset_obs: bool = False,
+):
+    """
+    Update quantization parameters on layer
+
+    :param layer: input layer
+    :param weight: weight to update quant params with, defaults to layer weight
+    :param g_idx: optional mapping from column index to group index
+    :param reset_obs: reset the observer before calculating quant params,
+        defaults to False
+    """
+    attached_weight = getattr(layer, "weight", None)
+
+    if weight is None:
+        weight = attached_weight
     scale = getattr(layer, "weight_scale", None)
     zero_point = getattr(layer, "weight_zero_point", None)
+    if g_idx is None:
+        g_idx = getattr(layer, "weight_g_idx", None)
     observer = getattr(layer, "weight_observer", None)
 
     if weight is None or observer is None or scale is None or zero_point is None:
         # scale, zp, or observer not calibratable or weight not available
         return
 
+    if reset_obs:
+        observer.reset()
+
+    if attached_weight is not None:
+        weight = weight.to(attached_weight.dtype)
+
     updated_scale, updated_zero_point = observer(weight)
 
     # update scale and zero point

compressed_tensors/quantization/lifecycle/initialize.py

@@ -17,8 +17,6 @@ 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.lifecycle.forward import (
     wrap_module_forward_quantized,
 )
@@ -43,6 +41,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,6 +53,8 @@ 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:
@@ -61,14 +62,18 @@ def initialize_module_for_quantization(
         return
 
     if scheme.input_activations is not None:
-        _initialize_scale_zero_point_observer(module, "input", scheme.input_activations)
+        _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
+            weight_shape = module.weight.shape
             _initialize_scale_zero_point_observer(
-                module, "weight", scheme.weights, weight_shape=weight_shape
+                module,
+                "weight",
+                scheme.weights,
+                weight_shape=weight_shape,
+                force_zero_point=force_zero_point,
             )
         else:
             _LOGGER.warning(
@@ -78,7 +83,10 @@ def initialize_module_for_quantization(
             )
     if scheme.output_activations is not None:
         _initialize_scale_zero_point_observer(
-            module, "output", scheme.output_activations
+            module,
+            "output",
+            scheme.output_activations,
+            force_zero_point=force_zero_point,
         )
 
     module.quantization_scheme = scheme
@@ -86,6 +94,16 @@ def initialize_module_for_quantization(
 
     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
@@ -116,6 +134,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()
@@ -141,16 +160,31 @@ def _initialize_scale_zero_point_observer(
                 weight_shape[1] // quantization_args.group_size,
             )
 
-    # 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),
-        requires_grad=False,
-    )
-    module.register_parameter(f"{base_name}_zero_point", init_zero_point)
+    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)
+
+    # initialize with empty for actorder, to be populated by GPTQ or state_dict
+    if quantization_args.actorder:
+        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)

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 = {}