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

compressed_tensors/quantization/quant_config.py

@@ -13,25 +13,31 @@
 # limitations under the License.
 
 from enum import Enum
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
-from compressed_tensors.base import QUANTIZATION_CONFIG_NAME
-from compressed_tensors.quantization.quant_scheme import QuantizationScheme
+from compressed_tensors.config import CompressionFormat
+from compressed_tensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.quantization.quant_scheme import (
+    QuantizationScheme,
+    preset_name_to_scheme,
+)
 from compressed_tensors.quantization.utils import (
     calculate_compression_ratio,
     is_module_quantized,
     iter_named_leaf_modules,
     module_type,
+    parse_out_kv_cache_args,
 )
 from pydantic import BaseModel, Field
 from torch.nn import Module
-from transformers import AutoConfig
 
 
 __all__ = [
     "QuantizationStatus",
     "QuantizationConfig",
     "LIFECYCLE_ORDER",
+    "DEFAULT_QUANTIZATION_METHOD",
+    "DEFAULT_QUANTIZATION_FORMAT",
 ]
 
 
@@ -62,10 +68,33 @@ class QuantizationStatus(str, Enum):
         return
 
     def __ge__(self, other):
+        if other is None:
+            return True
         if not isinstance(other, self.__class__):
             raise NotImplementedError
         return LIFECYCLE_ORDER.index(self) >= LIFECYCLE_ORDER.index(other)
 
+    def __gt__(self, other):
+        if other is None:
+            return True
+        if not isinstance(other, self.__class__):
+            raise NotImplementedError
+        return LIFECYCLE_ORDER.index(self) > LIFECYCLE_ORDER.index(other)
+
+    def __lt__(self, other):
+        if other is None:
+            return False
+        if not isinstance(other, self.__class__):
+            raise NotImplementedError
+        return LIFECYCLE_ORDER.index(self) < LIFECYCLE_ORDER.index(other)
+
+    def __le__(self, other):
+        if other is None:
+            return False
+        if not isinstance(other, self.__class__):
+            raise NotImplementedError
+        return LIFECYCLE_ORDER.index(self) <= LIFECYCLE_ORDER.index(other)
+
 
 LIFECYCLE_ORDER = [
     QuantizationStatus.INITIALIZED,
@@ -74,6 +103,9 @@ LIFECYCLE_ORDER = [
     QuantizationStatus.COMPRESSED,
 ]
 
+DEFAULT_QUANTIZATION_METHOD = "compressed-tensors"
+DEFAULT_QUANTIZATION_FORMAT = "fakequant"
+
 
 class QuantizationConfig(BaseModel):
     """
@@ -81,45 +113,62 @@ class QuantizationConfig(BaseModel):
     mapped to a QuantizationScheme in config_groups.
 
     :param config_groups: dict of QuantizationSchemes specifying the quantization
-    settings for each quantized layer
+    settings for each quantized layer. A group could also be a reference to
+    a predefined scheme name, mapped to a list of its target layers/classes
     :param quant_method: a constant used to differentiate sparseML quantization from
     other quantization configs
     :param format: specifies how the quantized model is stored on disk
     :quantization_status: specifies the current status of all quantized layers. It is
-    assumed all layers are in the same state.
+        assumed all layers are in the same state.
+    :param kv_cache_scheme: optional QuantizationArgs, that specify the
+        quantization of the kv cache. If None, kv cache is not quantized.
+        When applying kv cache quantization to transformer AutoModelForCausalLM,
+        the kv_cache_scheme gets converted into a QuantizationScheme that:
+            - targets the `q_proj` and `k_proj` modules of the model. The outputs
+              of those modules are the keys and values that might be cached
+            - quantizes the outputs of the aformentioned layers, so that
+              keys and values are compressed before storing them in the cache
+        There is an explicit assumption that the model contains modules with
+        `k_proj` and `v_proj` in their names. If this is not the case
+        and kv_cache_scheme != None, the quantization of kv cache will fail
     :global_compression_ratio: optional informational config to report the model
     compression ratio acheived by the quantization config
     :ignore: optional list of layers to ignore from config_groups. Layers in this list
     are not quantized even if they match up with a target in config_groups
     """
 
-    config_groups: Dict[str, QuantizationScheme]
-    quant_method: str = "sparseml"
-    format: str = "fakequant"
+    config_groups: Dict[str, Union[QuantizationScheme, List[str]]]
+    quant_method: str = DEFAULT_QUANTIZATION_METHOD
+    kv_cache_scheme: Optional[QuantizationArgs] = None
+    format: str = DEFAULT_QUANTIZATION_FORMAT
     quantization_status: QuantizationStatus = QuantizationStatus.INITIALIZED
     global_compression_ratio: Optional[float] = None
     ignore: Optional[List[str]] = Field(default_factory=list)
 
-    @staticmethod
-    def from_model_config(model_name_or_path) -> "QuantizationConfig":
+    def model_post_init(self, __context):
         """
-        Given a path to a model config, extract a quantization config if it exists
-
-        :param pretrained_model_name_or_path: path to model config on disk or HF hub
-        :return: instantiated QuantizationConfig if config contains a quant config
+        updates any quantization schemes defined as presets to be fully loaded
+        schemes
         """
-        config = AutoConfig.from_pretrained(model_name_or_path)
-        quantization_config = getattr(config, QUANTIZATION_CONFIG_NAME, None)
-        if quantization_config is None:
-            return None
-
-        return QuantizationConfig.parse_obj(quantization_config)
+        for group_name, targets_or_scheme in self.config_groups.items():
+            if isinstance(targets_or_scheme, QuantizationScheme):
+                continue  # scheme already defined
+            self.config_groups[group_name] = preset_name_to_scheme(
+                name=group_name,
+                targets=targets_or_scheme,
+            )
+
+    def to_dict(self):
+        # for compatibility with HFQuantizer
+        return self.dict()
 
     @staticmethod
-    def from_pretrained(model: Module) -> "QuantizationConfig":
+    def from_pretrained(
+        model: Module, format: Optional[str] = None
+    ) -> Optional["QuantizationConfig"]:
         """
         Converts a model into its associated QuantizationConfig based on the
-        QuantizationScheme attached to each quanitzed module
+        QuantizationScheme attached to each quantized module
 
         :param model: model to calculate quantization scheme of
         :return: filled out QuantizationScheme for the input model
@@ -147,6 +196,9 @@ class QuantizationConfig(BaseModel):
                 if not match_found:
                     quant_scheme_to_layers.append(scheme)
 
+        if len(quant_scheme_to_layers) == 0:  # No quantized layers
+            return None
+
         # clean up ignore list, we can leave out layers types if none of the
         # instances are quantized
         consolidated_ignore = []
@@ -157,15 +209,44 @@ class QuantizationConfig(BaseModel):
             # else we leave it off the ignore list, doesn't fall under any of the
             # existing quantization schemes so it won't be quantized
 
+        kv_cache_args, quant_scheme_to_layers = parse_out_kv_cache_args(
+            quant_scheme_to_layers
+        )
+        kv_cache_scheme = (
+            kv_cache_args.model_dump() if kv_cache_args is not None else kv_cache_args
+        )
+
         config_groups = {}
         for idx, scheme in enumerate(quant_scheme_to_layers):
             group_name = "group_" + str(idx)
             config_groups[group_name] = scheme
 
+        # TODO: this is incorrect in compressed mode, since we are overwriting the
+        # original weight we lose the uncompressed bit_depth indo
         compression_ratio = calculate_compression_ratio(model)
+
+        if format is None:
+            if quantization_status == QuantizationStatus.COMPRESSED:
+                format = CompressionFormat.int_quantized.value
+            else:
+                format = CompressionFormat.dense.value
+
         return QuantizationConfig(
             config_groups=config_groups,
             quantization_status=quantization_status,
+            kv_cache_scheme=kv_cache_scheme,
             global_compression_ratio=compression_ratio,
+            format=format,
             ignore=consolidated_ignore,
         )
+
+    def requires_calibration_data(self):
+        for _, scheme in self.config_groups.items():
+            if scheme.input_activations is not None:
+                if not scheme.input_activations.dynamic:
+                    return True
+            if scheme.output_activations is not None:
+                if not scheme.output_activations.dynamic:
+                    return True
+
+        return False

compressed_tensors/quantization/quant_scheme.py

@@ -12,13 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from copy import deepcopy
 from typing import List, Optional
 
-from compressed_tensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.quantization.quant_args import (
+    QuantizationArgs,
+    QuantizationStrategy,
+    QuantizationType,
+)
 from pydantic import BaseModel
 
 
-__all__ = ["QuantizationScheme"]
+__all__ = [
+    "QuantizationScheme",
+    "preset_name_to_scheme",
+    "is_preset_scheme",
+]
 
 
 class QuantizationScheme(BaseModel):
@@ -37,3 +46,175 @@ class QuantizationScheme(BaseModel):
     weights: Optional[QuantizationArgs] = None
     input_activations: Optional[QuantizationArgs] = None
     output_activations: Optional[QuantizationArgs] = None
+
+    @classmethod
+    def default_scheme(
+        cls,
+        targets: Optional[List[str]] = None,
+    ):
+
+        if targets is None:
+            # default to quantizing all Linear layers
+            targets = ["Linear"]
+
+        # default to 8 bit integer symmetric quantization
+        # for weights
+        weights = QuantizationArgs(num_bits=8, symmetric=True)
+
+        # default to 8 bit integer asymmetric quantization
+        input_activations = QuantizationArgs(num_bits=8, symmetric=True)
+
+        # Do not quantize the output activations
+        # by default
+        output_activations = None
+
+        return cls(
+            targets=targets,
+            weights=weights,
+            input_activations=input_activations,
+            output_activations=output_activations,
+        )
+
+
+"""
+Pre-Set Quantization Scheme Args
+"""
+
+
+def preset_name_to_scheme(name: str, targets: List[str]) -> QuantizationScheme:
+    """
+    :param name: preset quantization settings name. must exist in upper case in
+        PRESET_SCHEMES
+    :param targets: list of quantization targets to be passed to the Scheme
+    :return: new QuantizationScheme for a given name with the given targets
+    """
+    name = name.upper()
+
+    if name not in PRESET_SCHEMES:
+        raise KeyError(
+            f"Unknown preset scheme name {name}, "
+            f"available names: {list(PRESET_SCHEMES.keys())}"
+        )
+
+    scheme_args = deepcopy(PRESET_SCHEMES[name])  # deepcopy to avoid args references
+    return QuantizationScheme(
+        targets=targets,
+        **scheme_args,
+    )
+
+
+def is_preset_scheme(name: str) -> bool:
+    """
+    :param name: preset quantization settings name
+    :return: True if the name is a preset scheme name
+    """
+    return name.upper() in PRESET_SCHEMES
+
+
+# 8 bit integer weights and 8 bit activations quantization
+W8A8 = dict(
+    weights=QuantizationArgs(
+        num_bits=8,
+        type=QuantizationType.INT,
+        strategy=QuantizationStrategy.CHANNEL,
+        symmetric=True,
+        dynamic=False,
+    ),
+    input_activations=QuantizationArgs(
+        num_bits=8,
+        type=QuantizationType.INT,
+        strategy=QuantizationStrategy.TOKEN,
+        symmetric=True,
+        dynamic=True,
+    ),
+)
+
+# 8 bit integer weights only quantization
+W8A16 = dict(
+    weights=QuantizationArgs(
+        num_bits=8,
+        type=QuantizationType.INT,
+        strategy=QuantizationStrategy.CHANNEL,
+        symmetric=True,
+        dynamic=False,
+    ),
+)
+
+# 4 bit integer weights only quantization
+W4A16 = dict(
+    weights=QuantizationArgs(
+        num_bits=4,
+        type=QuantizationType.INT,
+        strategy=QuantizationStrategy.GROUP,
+        group_size=128,
+        symmetric=True,
+        dynamic=False,
+    ),
+)
+
+# 4 bit integer weights and 8 bit activations quantization
+W4A8 = dict(
+    weights=QuantizationArgs(
+        num_bits=4,
+        type=QuantizationType.INT,
+        group_size=128,
+        strategy=QuantizationStrategy.GROUP,
+        symmetric=True,
+        dynamic=False,
+    ),
+    input_activations=QuantizationArgs(
+        num_bits=8,
+        type=QuantizationType.INT,
+        strategy=QuantizationStrategy.TOKEN,
+        symmetric=True,
+        dynamic=True,
+    ),
+)
+
+# FP8 weights and FP8 activations quantization
+FP8 = dict(
+    weights=QuantizationArgs(
+        num_bits=8,
+        type=QuantizationType.FLOAT,
+        strategy=QuantizationStrategy.TENSOR,
+        symmetric=True,
+        dynamic=False,
+    ),
+    input_activations=QuantizationArgs(
+        num_bits=8,
+        type=QuantizationType.FLOAT,
+        strategy=QuantizationStrategy.TENSOR,
+        symmetric=True,
+        dynamic=False,
+    ),
+)
+
+# FP8 weights and FP8 dynamic activations quantization
+FP8_DYNAMIC = dict(
+    weights=QuantizationArgs(
+        num_bits=8,
+        type=QuantizationType.FLOAT,
+        strategy=QuantizationStrategy.CHANNEL,
+        symmetric=True,
+        dynamic=False,
+    ),
+    input_activations=QuantizationArgs(
+        num_bits=8,
+        type=QuantizationType.FLOAT,
+        strategy=QuantizationStrategy.TOKEN,
+        symmetric=True,
+        dynamic=True,
+    ),
+)
+
+PRESET_SCHEMES = {
+    # Integer weight only schemes
+    "W8A16": W8A16,
+    "W4A16": W4A16,
+    # Integer weight and activation schemes
+    "W8A8": W8A8,
+    "W4A8": W4A8,
+    # Float weight and activation schemes
+    "FP8": FP8,
+    "FP8_DYNAMIC": FP8_DYNAMIC,
+}

compressed_tensors/quantization/utils/helpers.py

@@ -12,21 +12,50 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Tuple
+import logging
+import re
+from typing import List, Optional, Tuple
 
 import torch
+from compressed_tensors.quantization.observers.base import Observer
+from compressed_tensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.quantization.quant_scheme import QuantizationScheme
 from torch.nn import Module
 from tqdm import tqdm
 
 
 __all__ = [
+    "infer_quantization_status",
     "is_module_quantized",
     "is_model_quantized",
     "iter_named_leaf_modules",
     "module_type",
     "calculate_compression_ratio",
+    "get_torch_bit_depth",
+    "can_quantize",
+    "parse_out_kv_cache_args",
+    "KV_CACHE_TARGETS",
+    "is_kv_cache_quant_scheme",
 ]
 
+KV_CACHE_TARGETS = ["re:.*k_proj", "re:.*v_proj"]
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+def infer_quantization_status(model: Module) -> Optional["QuantizationStatus"]:  # noqa
+    """
+    Checks the quantization status of a model. Assumes all modules in the model have
+    the same status, so only the first quantized model is checked.
+
+    :param model: model to check quantization status for
+    :return: quantization status if the model is quantized, otherwise None
+    """
+    for module in model.modules():
+        status = getattr(module, "quantization_status", None)
+        if status is not None:
+            return status
+    return None
+
 
 def is_module_quantized(module: Module) -> bool:
     """
@@ -78,11 +107,60 @@ def module_type(module: Module) -> str:
 
 
 def iter_named_leaf_modules(model: Module) -> Tuple[str, Module]:
-    # yields modules that do not have any submodules
-    # TODO: potentially expand to add list of allowed submodules such as observers
+    """
+    Yields modules that do not have any submodules except observers. The observers
+    themselves are not yielded
+
+    :param model: model to get leaf modules of
+    :returns: generator tuple of (name, leaf_submodule)
+    """
     for name, submodule in model.named_modules():
-        if len(list(submodule.children())) == 0:
+        children = list(submodule.children())
+        if len(children) == 0 and not isinstance(submodule, Observer):
             yield name, submodule
+        else:
+            has_non_observer_children = False
+            for child in children:
+                if not isinstance(child, Observer):
+                    has_non_observer_children = True
+
+            if not has_non_observer_children:
+                yield name, submodule
+
+
+def get_torch_bit_depth(value: torch.Tensor) -> int:
+    """
+    Determine the number of bits used to represent the dtype of a tensor
+
+    :param value: tensor to check bit depth of
+    :return: bit depth of each element in the value tensor
+    """
+    try:
+        bit_depth = torch.finfo(value.dtype).bits
+    except TypeError:
+        bit_depth = torch.iinfo(value.dtype).bits
+
+    return bit_depth
+
+
+def can_quantize(value: torch.Tensor, quant_args: "QuantizationArgs") -> bool:  # noqa
+    """
+    Checks if value can be quantized by quant_args.
+
+    :param value: tensor to check for quantization
+    :param quant_args: QuantizationArgs to use for quantization
+    :return: False if value is already quantized to quant_args or value is incompatible
+    with quant_args, True if value can be quantized with quant_args
+    """
+    bit_depth = get_torch_bit_depth(value)
+    requested_depth = quant_args.num_bits
+    if bit_depth < quant_args.num_bits:
+        _LOGGER.warn(
+            f"Can't quantize tensor with bit depth {bit_depth} to {requested_depth}."
+            "The QuantizationArgs provided are not compatible with the input tensor."
+        )
+
+    return bit_depth > quant_args.num_bits
 
 
 def calculate_compression_ratio(model: Module) -> float:
@@ -101,10 +179,7 @@ def calculate_compression_ratio(model: Module) -> float:
         desc="Calculating quantization compression ratio",
     ):
         for parameter in model.parameters():
-            try:
-                uncompressed_bits = torch.finfo(parameter.dtype).bits
-            except TypeError:
-                uncompressed_bits = torch.iinfo(parameter.dtype).bits
+            uncompressed_bits = get_torch_bit_depth(parameter)
             compressed_bits = uncompressed_bits
             if is_module_quantized(submodule):
                 compressed_bits = submodule.quantization_scheme.weights.num_bits
@@ -114,3 +189,62 @@ def calculate_compression_ratio(model: Module) -> float:
             total_uncompressed += uncompressed_bits * num_weights
 
     return total_uncompressed / total_compressed
+
+
+def is_kv_cache_quant_scheme(scheme: QuantizationScheme) -> bool:
+    """
+    Check whether the QuantizationScheme targets the kv cache.
+    It does if all the following criteria are met:
+    - the scheme targets either exactly match the KV_CACHE_TARGETS
+        or the match KV_CACHE_TARGETS regex pattern
+    - the scheme quantizes output_activations (we want to quantize the
+        outputs from the KV_CACHE_TARGETS, as their correspond to the
+        keys and values that are to be saved in the cache)
+
+    :param scheme: The QuantizationScheme to investigate
+    :return: boolean flag
+    """
+    if len(scheme.targets) == 1:
+        # match on the KV_CACHE_TARGETS regex pattern
+        # if there is only one target
+        is_match_targets = any(
+            [re.match(pattern[3:], scheme.targets[0]) for pattern in KV_CACHE_TARGETS]
+        )
+    else:
+        # match on the exact KV_CACHE_TARGETS
+        # if there are multiple targets
+        is_match_targets = set(KV_CACHE_TARGETS) == set(scheme.targets)
+
+    is_match_output_activations = scheme.output_activations is not None
+    return is_match_targets and is_match_output_activations
+
+
+def parse_out_kv_cache_args(
+    quant_scheme_to_layers: List[QuantizationScheme],
+) -> Tuple[Optional[QuantizationArgs], List[QuantizationScheme]]:
+    """
+    If possible, parse out the kv cache specific QuantizationArgs
+    from the list of the QuantizationSchemes. If no kv cache
+    specific QuantizationArgs available, this function acts
+    as an identity function
+
+    :param quant_scheme_to_layers: list of QuantizationSchemes
+    :return: kv_cache_args (optional) and the (remaining or original)
+        list of the QuantizationSchemes
+    """
+    kv_cache_quant_scheme_to_layers = [
+        scheme for scheme in quant_scheme_to_layers if is_kv_cache_quant_scheme(scheme)
+    ]
+    quant_scheme_to_layers = [
+        scheme
+        for scheme in quant_scheme_to_layers
+        if not is_kv_cache_quant_scheme(scheme)
+    ]
+
+    if kv_cache_quant_scheme_to_layers:
+        kv_cache_quant_scheme_to_layers = kv_cache_quant_scheme_to_layers[0]
+        kv_cache_args = kv_cache_quant_scheme_to_layers.output_activations
+    else:
+        kv_cache_args = None
+
+    return kv_cache_args, quant_scheme_to_layers

compressed_tensors/utils/__init__.py

@@ -13,4 +13,8 @@
 # limitations under the License.
 # flake8: noqa
 
+from .helpers import *
+from .offload import *
+from .permutations_24 import *
 from .safetensors_load import *
+from .semi_structured_conversions import *