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

compressed_tensors/__init__.py

@@ -19,3 +19,4 @@ from .compressors import *
 from .config import *
 from .quantization import QuantizationConfig, QuantizationStatus
 from .utils import *
+from .version import *

compressed_tensors/base.py

@@ -16,3 +16,5 @@ SPARSITY_CONFIG_NAME = "sparsity_config"
 QUANTIZATION_CONFIG_NAME = "quantization_config"
 COMPRESSION_CONFIG_NAME = "compression_config"
 KV_CACHE_SCHEME_NAME = "kv_cache_scheme"
+COMPRESSION_VERSION_NAME = "version"
+QUANTIZATION_METHOD_NAME = "quant_method"

compressed_tensors/compressors/__init__.py

@@ -14,15 +14,9 @@
 
 # flake8: noqa
 
-from .base import Compressor
-from .dense import DenseCompressor
-from .helpers import load_compressed, save_compressed, save_compressed_model
-from .marlin_24 import Marlin24Compressor
-from .model_compressor import ModelCompressor, map_modules_to_quant_args
-from .naive_quantized import (
-    FloatQuantizationCompressor,
-    IntQuantizationCompressor,
-    QuantizationCompressor,
-)
-from .pack_quantized import PackedQuantizationCompressor
-from .sparse_bitmask import BitmaskCompressor, BitmaskTensor
+from .base import *
+from .helpers import *
+from .model_compressors import *
+from .quantized_compressors import *
+from .sparse_compressors import *
+from .sparse_quantized_compressors import *

compressed_tensors/compressors/base.py

@@ -12,20 +12,47 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Generator, Tuple, Union
+from abc import ABC, abstractmethod
+from typing import Dict, Generator, Optional, Tuple, Union
 
+import torch
 from compressed_tensors.config import SparsityCompressionConfig
-from compressed_tensors.quantization import QuantizationConfig
+from compressed_tensors.quantization import QuantizationArgs, QuantizationConfig
 from compressed_tensors.registry import RegistryMixin
 from torch import Tensor
+from torch.nn import Module
 
 
-__all__ = ["Compressor"]
+__all__ = ["BaseCompressor"]
 
 
-class Compressor(RegistryMixin):
+class BaseCompressor(RegistryMixin, ABC):
     """
-    Base class representing a model compression algorithm
+    Base class representing a model compression algorithm. Each child class should
+    implement compression_param_info, compress_weight and decompress_weight.
+
+    Compressors support compressing/decompressing a full module state dict or a single
+    quantized PyTorch leaf module.
+
+    Model Load Lifecycle (run_compressed=False):
+        - ModelCompressor.decompress()
+            - apply_quantization_config()
+            - BaseCompressor.decompress()
+
+    Model Save Lifecycle:
+        - ModelCompressor.compress()
+            - BaseCompressor.compress()
+
+
+    Module Lifecycle (run_compressed=True):
+        - apply_quantization_config()
+        - compressed_module = CompressedLinear(module)
+            - initialize_module_for_quantization()
+            - BaseCompressor.compression_param_info()
+            - register_parameters()
+        - compressed_module.forward()
+            -compressed_module.decompress()
+
 
     :param config: config specifying compression parameters
     """
@@ -35,26 +62,127 @@ class Compressor(RegistryMixin):
     ):
         self.config = config
 
-    def compress(self, model_state: Dict[str, Tensor], **kwargs) -> Dict[str, Tensor]:
+    def compression_param_info(
+        self,
+        weight_shape: torch.Size,
+        quantization_args: Optional[QuantizationArgs] = None,
+    ) -> Dict[str, Tuple[torch.Size, torch.dtype]]:
+        """
+        Creates a dictionary of expected shapes and dtypes for each compression
+            parameter used by the compressor
+
+        :param weight_shape: uncompressed weight shape
+        :param quantization_args: quantization parameters for the weight
+        :return: dictionary mapping compressed parameter names to shape and dtype
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def compress(
+        self,
+        model_state: Dict[str, Tensor],
+        **kwargs,
+    ) -> Dict[str, Tensor]:
         """
         Compresses a dense state dict
 
         :param model_state: state dict of uncompressed model
+        :param kwargs: additional arguments for compression
         :return: compressed state dict
         """
         raise NotImplementedError()
 
+    @abstractmethod
     def decompress(
-        self, path_to_model_or_tensors: str, device: str = "cpu", **kwargs
+        self,
+        path_to_model_or_tensors: str,
+        device: str = "cpu",
+        **kwargs,
     ) -> Generator[Tuple[str, Tensor], None, None]:
         """
         Reads a compressed state dict located at path_to_model_or_tensors
         and returns a generator for sequentially decompressing back to a
         dense state dict
 
-        :param model_path: path to compressed safetensors model (directory with
-            one or more safetensors files) or compressed tensors file
+        :param path_to_model_or_tensors: path to compressed safetensors model (directory
+            with one or more safetensors files) or compressed tensors file
+        :param names_to_scheme: quantization args for each quantized weight
         :param device: optional device to load intermediate weights into
         :return: compressed state dict
         """
         raise NotImplementedError()
+
+    def compress_module(self, module: Module) -> Optional[Dict[str, torch.Tensor]]:
+        """
+        Compresses a single quantized leaf PyTorch module. If the module is not
+        quantized, this function has no effect.
+
+        :param module: PyTorch module to compress
+        :return: dictionary of compressed weight data, or None if module is not
+            quantized
+        """
+        if not hasattr(module, "quantization_scheme"):
+            return None  # module is not quantized
+        quantization_scheme = module.quantization_scheme
+        if not hasattr(quantization_scheme, "weights"):
+            return None  # weights are not quantized
+
+        quantization_args = quantization_scheme.weights
+        weight = getattr(module, "weight", None)
+        weight_scale = getattr(module, "weight_scale", None)
+        weight_zero_point = getattr(module, "weight_zero_point", None)
+
+        return self.compress_weight(
+            weight=weight,
+            scale=weight_scale,
+            zero_point=weight_zero_point,
+            quantization_args=quantization_args,
+        )
+
+    def compress_weight(
+        self,
+        weight: Tensor,
+        **kwargs,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Compresses a single uncompressed weight
+
+        :param weight: uncompressed weight tensor
+        :param kwargs: additional arguments for compression
+        """
+        raise NotImplementedError()
+
+    def decompress_module(self, module: Module):
+        """
+        Decompresses a single compressed leaf PyTorch module. If the module is not
+        quantized, this function has no effect.
+
+        :param module: PyTorch module to decompress
+        :return: tensor of the decompressed weight, or None if module is not quantized
+        """
+        if not hasattr(module, "quantization_scheme"):
+            return None  # module is not quantized
+        quantization_scheme = module.quantization_scheme
+        if not hasattr(quantization_scheme, "weights"):
+            return None  # weights are not quantized
+
+        quantization_args = quantization_scheme.weights
+        compressed_data = {}
+        for name, parameter in module.named_parameters():
+            compressed_data[name] = parameter
+
+        return self.decompress_weight(
+            compressed_data=compressed_data, quantization_args=quantization_args
+        )
+
+    def decompress_weight(
+        self, compressed_data: Dict[str, Tensor], **kwargs
+    ) -> torch.Tensor:
+        """
+        Decompresses a single compressed weight
+
+        :param compressed_data: dictionary of data needed for decompression
+        :param kwargs: additional arguments for decompression
+        :return: tensor of the decompressed weight
+        """
+        raise NotImplementedError()

compressed_tensors/compressors/helpers.py

@@ -16,7 +16,7 @@ from pathlib import Path
 from typing import Dict, Generator, Optional, Tuple, Union
 
 import torch
-from compressed_tensors.compressors import Compressor
+from compressed_tensors.compressors import BaseCompressor
 from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig
 from compressed_tensors.utils.safetensors_load import get_weight_mappings
 from safetensors import safe_open
@@ -52,16 +52,16 @@ def save_compressed(
     compression_format = compression_format or CompressionFormat.dense.value
 
     if not (
-        compression_format in Compressor.registered_names()
-        or compression_format in Compressor.registered_aliases()
+        compression_format in BaseCompressor.registered_names()
+        or compression_format in BaseCompressor.registered_aliases()
     ):
         raise ValueError(
             f"Unknown compression format: {compression_format}. "
-            f"Must be one of {set(Compressor.registered_names() + Compressor.registered_aliases())}"  # noqa E501
+            f"Must be one of {set(BaseCompressor.registered_names() + BaseCompressor.registered_aliases())}"  # noqa E501
         )
 
     # compress
-    compressor = Compressor.load_from_registry(compression_format)
+    compressor = BaseCompressor.load_from_registry(compression_format)
     # save compressed tensors
     compressed_tensors = compressor.compress(tensors)
     save_file(compressed_tensors, save_path)
@@ -102,7 +102,7 @@ def load_compressed(
     else:
         # decompress tensors
         compression_format = compression_config.format
-        compressor = Compressor.load_from_registry(
+        compressor = BaseCompressor.load_from_registry(
             compression_format, config=compression_config
         )
         yield from compressor.decompress(compressed_tensors, device=device)

compressed_tensors/compressors/model_compressors/__init__.py

@@ -0,0 +1,17 @@
+# 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
+
+
+from .model_compressor import *

compressed_tensors/compressors/{model_compressor.py → model_compressors/model_compressor.py}

@@ -18,18 +18,22 @@ import operator
 import os
 import re
 from copy import deepcopy
-from typing import Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, TypeVar, Union
 
+import compressed_tensors
 import torch
 import transformers
 from compressed_tensors.base import (
     COMPRESSION_CONFIG_NAME,
+    COMPRESSION_VERSION_NAME,
     QUANTIZATION_CONFIG_NAME,
+    QUANTIZATION_METHOD_NAME,
     SPARSITY_CONFIG_NAME,
 )
-from compressed_tensors.compressors import Compressor
-from compressed_tensors.config import SparsityCompressionConfig
+from compressed_tensors.compressors.base import BaseCompressor
+from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig
 from compressed_tensors.quantization import (
+    DEFAULT_QUANTIZATION_METHOD,
     QuantizationConfig,
     QuantizationStatus,
     apply_quantization_config,
@@ -40,7 +44,10 @@ from compressed_tensors.quantization.utils import (
     iter_named_leaf_modules,
 )
 from compressed_tensors.utils import get_safetensors_folder, update_parameter_data
-from compressed_tensors.utils.helpers import fix_fsdp_module_name
+from compressed_tensors.utils.helpers import (
+    fix_fsdp_module_name,
+    is_compressed_tensors_config,
+)
 from torch import Tensor
 from torch.nn import Module
 from tqdm import tqdm
@@ -53,6 +60,11 @@ __all__ = ["ModelCompressor", "map_modules_to_quant_args"]
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 
+if TYPE_CHECKING:
+    # dummy type if not available from transformers
+    CompressedTensorsConfig = TypeVar("CompressedTensorsConfig")
+
+
 class ModelCompressor:
     """
     Handles compression and decompression of a model with a sparsity config and/or
@@ -88,45 +100,41 @@ class ModelCompressor:
         configs and load a ModelCompressor
 
         :param pretrained_model_name_or_path: path to model config on disk or HF hub
-        :return: compressor for the extracted configs
+        :return: compressor for the configs, or None if model is not compressed
         """
         config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
         compression_config = getattr(config, COMPRESSION_CONFIG_NAME, None)
         return cls.from_compression_config(compression_config)
 
     @classmethod
-    def from_compression_config(cls, compression_config: Dict[str, Any]):
+    def from_compression_config(
+        cls, compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"]
+    ):
         """
-        :param compression_config: compression/quantization config dictionary
-            found under key "quantization_config" in HF model config
-        :return: compressor for the extracted configs
+        :param compression_config:
+            A compression or quantization config
+
+            The type is one of the following:
+            1. A Dict found under either "quantization_config" or "compression_config"
+                keys in the config.json
+            2. A CompressedTensorsConfig found under key "quantization_config" in HF
+                model config
+        :return: compressor for the configs, or None if model is not compressed
         """
         if compression_config is None:
             return None
 
-        try:
-            from transformers.utils.quantization_config import CompressedTensorsConfig
-
-            if isinstance(compression_config, CompressedTensorsConfig):
-                compression_config = compression_config.to_dict()
-        except ImportError:
-            pass
-
         sparsity_config = cls.parse_sparsity_config(compression_config)
         quantization_config = cls.parse_quantization_config(compression_config)
         if sparsity_config is None and quantization_config is None:
             return None
 
-        if sparsity_config is not None and not isinstance(
-            sparsity_config, SparsityCompressionConfig
-        ):
+        if sparsity_config is not None:
             format = sparsity_config.get("format")
             sparsity_config = SparsityCompressionConfig.load_from_registry(
                 format, **sparsity_config
             )
-        if quantization_config is not None and not isinstance(
-            quantization_config, QuantizationConfig
-        ):
+        if quantization_config is not None:
             quantization_config = QuantizationConfig.parse_obj(quantization_config)
 
         return cls(
@@ -149,7 +157,7 @@ class ModelCompressor:
             to a sparsity compression algorithm
         :param quantization_format: string corresponding to a quantization compression
             algorithm
-        :return: compressor for the extracted configs
+        :return: compressor for the configs, or None if model is not compressed
         """
         quantization_config = QuantizationConfig.from_pretrained(
             model, format=quantization_format
@@ -168,32 +176,60 @@ class ModelCompressor:
         )
 
     @staticmethod
-    def parse_sparsity_config(compression_config: Dict) -> Union[Dict, None]:
+    def parse_sparsity_config(
+        compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"]
+    ) -> Union[Dict[str, Any], None]:
+        """
+        Parse sparsity config from quantization/compression config. Sparsity
+        config is nested inside q/c config
+
+        :param compression_config: quantization/compression config
+        :return: sparsity config
+        """
         if compression_config is None:
             return None
-        if SPARSITY_CONFIG_NAME not in compression_config:
-            return None
-        if hasattr(compression_config, SPARSITY_CONFIG_NAME):
-            # for loaded HFQuantizer config
-            return getattr(compression_config, SPARSITY_CONFIG_NAME)
 
-        # SparseAutoModel format
+        if is_compressed_tensors_config(compression_config):
+            s_config = compression_config.sparsity_config
+            return s_config.dict() if s_config is not None else None
+
         return compression_config.get(SPARSITY_CONFIG_NAME, None)
 
     @staticmethod
-    def parse_quantization_config(compression_config: Dict) -> Union[Dict, None]:
+    def parse_quantization_config(
+        compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"]
+    ) -> Union[Dict[str, Any], None]:
+        """
+        Parse quantization config from quantization/compression config. The
+        quantization are all the fields that are not the sparsity config or
+        metadata fields
+
+        :param compression_config: quantization/compression config
+        :return: quantization config without sparsity config or metadata fields
+        """
         if compression_config is None:
             return None
 
-        if hasattr(compression_config, QUANTIZATION_CONFIG_NAME):
-            # for loaded HFQuantizer config
-            return getattr(compression_config, QUANTIZATION_CONFIG_NAME)
+        if is_compressed_tensors_config(compression_config):
+            q_config = compression_config.quantization_config
+            return q_config.dict() if q_config is not None else None
 
-        # SparseAutoModel format
         quantization_config = deepcopy(compression_config)
         quantization_config.pop(SPARSITY_CONFIG_NAME, None)
+
+        # some fields are required, even if a qconfig is not present
+        # pop them off and if nothing remains, then there is no qconfig
+        quant_method = quantization_config.pop(QUANTIZATION_METHOD_NAME, None)
+        _ = quantization_config.pop(COMPRESSION_VERSION_NAME, None)
+
         if len(quantization_config) == 0:
-            quantization_config = None
+            return None
+
+        # replace popped off values
+        # note that version is discarded for now
+        if quant_method is not None:
+            quantization_config[QUANTIZATION_METHOD_NAME] = quant_method
+
         return quantization_config
 
     def __init__(
@@ -207,11 +243,11 @@ class ModelCompressor:
         self.quantization_compressor = None
 
         if sparsity_config is not None:
-            self.sparsity_compressor = Compressor.load_from_registry(
+            self.sparsity_compressor = BaseCompressor.load_from_registry(
                 sparsity_config.format, config=sparsity_config
             )
         if quantization_config is not None:
-            self.quantization_compressor = Compressor.load_from_registry(
+            self.quantization_compressor = BaseCompressor.load_from_registry(
                 quantization_config.format, config=quantization_config
             )
 
@@ -222,7 +258,7 @@ class ModelCompressor:
         Compresses a dense state dict or model with sparsity and/or quantization
 
         :param model: uncompressed model to compress
-        :param model_state: optional uncompressed state_dict to insert into model
+        :param state_dict: optional uncompressed state_dict to insert into model
         :return: compressed state dict
         """
         if state_dict is None:
@@ -234,6 +270,10 @@ class ModelCompressor:
             compressed_state_dict = self.quantization_compressor.compress(
                 state_dict, names_to_scheme=quantized_modules_to_args
             )
+            if self.quantization_config.format != CompressionFormat.dense.value:
+                self.quantization_config.quantization_status = (
+                    QuantizationStatus.COMPRESSED
+                )
 
         if self.sparsity_compressor is not None:
             compressed_state_dict = self.sparsity_compressor.compress(
@@ -281,6 +321,9 @@ class ModelCompressor:
 
         :param save_directory: path to a folder containing a HF model config
         """
+        if self.quantization_config is None and self.sparsity_config is None:
+            return
+
         config_file_path = os.path.join(save_directory, CONFIG_NAME)
         if not os.path.exists(config_file_path):
             _LOGGER.warning(
@@ -292,13 +335,26 @@ class ModelCompressor:
         with open(config_file_path, "r") as config_file:
             config_data = json.load(config_file)
 
-        config_data[COMPRESSION_CONFIG_NAME] = {}
+        # required metadata whenever a quantization or sparsity config is present
+        # overwrite previous config and version if already existing
+        config_data[QUANTIZATION_CONFIG_NAME] = {}
+        config_data[QUANTIZATION_CONFIG_NAME][
+            COMPRESSION_VERSION_NAME
+        ] = compressed_tensors.__version__
+        if self.quantization_config is not None:
+            self.quantization_config.quant_method = DEFAULT_QUANTIZATION_METHOD
+        else:
+            config_data[QUANTIZATION_CONFIG_NAME][
+                QUANTIZATION_METHOD_NAME
+            ] = DEFAULT_QUANTIZATION_METHOD
+
+        # quantization and sparsity configs
         if self.quantization_config is not None:
             quant_config_data = self.quantization_config.model_dump()
-            config_data[COMPRESSION_CONFIG_NAME] = quant_config_data
+            config_data[QUANTIZATION_CONFIG_NAME] = quant_config_data
         if self.sparsity_config is not None:
             sparsity_config_data = self.sparsity_config.model_dump()
-            config_data[COMPRESSION_CONFIG_NAME][
+            config_data[QUANTIZATION_CONFIG_NAME][
                 SPARSITY_CONFIG_NAME
             ] = sparsity_config_data
 

compressed_tensors/compressors/quantized_compressors/__init__.py

@@ -0,0 +1,18 @@
+# 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
+
+from .base import *
+from .naive_quantized import *
+from .pack_quantized import *