compressed-tensors-nightly 0.3.3.20240514__py3-none-any.whl

compressed_tensors/__init__.py

@@ -0,0 +1,21 @@
+# 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 .base import *
+
+# flake8: noqa
+from .compressors import *
+from .config import *
+from .quantization import QuantizationConfig, QuantizationStatus
+from .utils import *

compressed_tensors/base.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.
+
+SPARSITY_CONFIG_NAME = "sparsity_config"
+QUANTIZATION_CONFIG_NAME = "quantization_config"
+COMPRESSION_CONFIG_NAME = "compression_config"

compressed_tensors/compressors/__init__.py

@@ -0,0 +1,22 @@
+# 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 Compressor
+from .dense import DenseCompressor
+from .helpers import load_compressed, save_compressed, save_compressed_model
+from .int_quantized import IntQuantizationCompressor
+from .model_compressor import ModelCompressor
+from .sparse_bitmask import BitmaskCompressor, BitmaskTensor

compressed_tensors/compressors/base.py

@@ -0,0 +1,59 @@
+# 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 Dict, Generator, Tuple, Union
+
+from compressed_tensors.config import SparsityCompressionConfig
+from compressed_tensors.quantization import QuantizationConfig
+from compressed_tensors.registry import RegistryMixin
+from torch import Tensor
+
+
+__all__ = ["Compressor"]
+
+
+class Compressor(RegistryMixin):
+    """
+    Base class representing a model compression algorithm
+
+    :param config: config specifying compression parameters
+    """
+
+    def __init__(
+        self, config: Union[SparsityCompressionConfig, QuantizationConfig, None] = None
+    ):
+        self.config = config
+
+    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
+        :return: compressed state dict
+        """
+        raise NotImplementedError()
+
+    def decompress(
+        self, path_to_model_or_tensors: str, device: str = "cpu"
+    ) -> 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
+        :return: compressed state dict
+        """
+        raise NotImplementedError()

compressed_tensors/compressors/dense.py

@@ -0,0 +1,34 @@
+# 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 Dict, Generator, Tuple
+
+from compressed_tensors.compressors import Compressor
+from compressed_tensors.config import CompressionFormat
+from torch import Tensor
+
+
+@Compressor.register(name=CompressionFormat.dense.value)
+class DenseCompressor(Compressor):
+    """
+    Identity compressor for dense models, returns the original state_dict
+    """
+
+    def compress(self, model_state: Dict[str, Tensor], **kwargs) -> Dict[str, Tensor]:
+        return model_state
+
+    def decompress(
+        self, path_to_model_or_tensors: str, device: str = "cpu"
+    ) -> Generator[Tuple[str, Tensor], None, None]:
+        return iter([])

compressed_tensors/compressors/helpers.py

@@ -0,0 +1,137 @@
+# 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 pathlib import Path
+from typing import Dict, Generator, Optional, Tuple, Union
+
+import torch
+from compressed_tensors.compressors import Compressor
+from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig
+from compressed_tensors.utils.safetensors_load import get_weight_mappings
+from safetensors import safe_open
+from safetensors.torch import save_file
+from torch import Tensor
+
+
+__all__ = [
+    "load_compressed",
+    "save_compressed",
+    "save_compressed_model",
+]
+
+
+def save_compressed(
+    tensors: Dict[str, Tensor],
+    save_path: Union[str, Path],
+    compression_format: Optional[CompressionFormat] = None,
+):
+    """
+    Save compressed tensors to disk. If tensors are not compressed,
+    save them as is.
+
+    :param tensors: dictionary of tensors to compress
+    :param save_path: path to save compressed tensors
+    :param compression_format: compression format used for the tensors
+    :return: compression config, if tensors were compressed - None otherwise
+    """
+    if tensors is None or len(tensors) == 0:
+        raise ValueError("No tensors or empty tensors provided to compress")
+
+    # if no compression_format specified, default to `dense`
+    compression_format = compression_format or CompressionFormat.dense.value
+
+    if not (
+        compression_format in Compressor.registered_names()
+        or compression_format in Compressor.registered_aliases()
+    ):
+        raise ValueError(
+            f"Unknown compression format: {compression_format}. "
+            f"Must be one of {set(Compressor.registered_names() + Compressor.registered_aliases())}"  # noqa E501
+        )
+
+    # compress
+    compressor = Compressor.load_from_registry(compression_format)
+    # save compressed tensors
+    compressed_tensors = compressor.compress(tensors)
+    save_file(compressed_tensors, save_path)
+
+
+def load_compressed(
+    compressed_tensors: Union[str, Path],
+    compression_config: SparsityCompressionConfig = None,
+    device: Optional[str] = "cpu",
+) -> Generator[Tuple[str, Tensor], None, None]:
+    """
+    Load compressed tensors from disk.
+    If tensors are not compressed, load them as is.
+
+    :param compressed_tensors: path to compressed tensors.
+        This can be a path to a file or a directory containing
+        one or multiple safetensor files (if multiple - in the format
+        assumed by huggingface)
+    :param compression_config: compression config to use for decompressing tensors.
+    :param device: device to move tensors to. If None, tensors are loaded on CPU.
+    :param return_dict: if True, return a dictionary of decompressed tensors
+    :return a generator that yields the name and tensor of the decompressed tensor
+    """
+    if compressed_tensors is None or not Path(compressed_tensors).exists():
+        raise ValueError("No compressed tensors provided to load")
+
+    if (
+        compression_config is None
+        or compression_config.format == CompressionFormat.dense.value
+    ):
+        # if no compression_config specified, or `dense` format specified,
+        # assume tensors are not compressed on disk
+        weight_mappings = get_weight_mappings(compressed_tensors)
+        for weight_name, file_with_weight_name in weight_mappings.items():
+            with safe_open(file_with_weight_name, framework="pt", device=device) as f:
+                weight = f.get_tensor(weight_name)
+                yield weight_name, weight
+    else:
+        # decompress tensors
+        compression_format = compression_config.format
+        compressor = Compressor.load_from_registry(
+            compression_format, config=compression_config
+        )
+        yield from compressor.decompress(compressed_tensors, device=device)
+
+
+def save_compressed_model(
+    model: torch.nn.Module,
+    filename: str,
+    compression_format: Optional[CompressionFormat] = None,
+    force_contiguous: bool = True,
+):
+    """
+    Wrapper around safetensors `save_model` helper function, which allows for
+    saving compressed model to disk.
+
+    Note: The model is assumed to have a
+        state_dict with  unique entries
+
+    :param model: model to save on disk
+    :param filename: filename location to save the file
+    :param compression_format: compression format used for the model
+    :param force_contiguous: forcing the state_dict to be saved as contiguous tensors
+    """
+    state_dict = model.state_dict()
+    if force_contiguous:
+        state_dict = {k: v.contiguous() for k, v in state_dict.items()}
+    try:
+        save_compressed(state_dict, filename, compression_format=compression_format)
+    except ValueError as e:
+        msg = str(e)
+        msg += " Or use save_compressed_model(..., force_contiguous=True), read the docs for potential caveats."  # noqa E501
+        raise ValueError(msg)

compressed_tensors/compressors/int_quantized.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from typing import Dict, Generator, Tuple
+
+import torch
+from compressed_tensors.compressors import Compressor
+from compressed_tensors.config import CompressionFormat
+from compressed_tensors.quantization.lifecycle.forward import dequantize, quantize
+from compressed_tensors.utils import get_nested_weight_mappings, merge_names
+from safetensors import safe_open
+from torch import Tensor
+from tqdm import tqdm
+
+
+__all__ = ["IntQuantizationCompressor"]
+
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+@Compressor.register(name=CompressionFormat.int_quantized.value)
+class IntQuantizationCompressor(Compressor):
+    """
+    Integer compression for quantized models. Weight of each quantized layer is
+    converted from its original float type to the format specified by the layer's
+    quantization scheme.
+    """
+
+    COMPRESSION_PARAM_NAMES = ["weight", "weight_scale", "weight_zero_point"]
+
+    def compress(self, model_state: Dict[str, Tensor], **kwargs) -> Dict[str, Tensor]:
+        model_quant_args = kwargs["model_quant_args"]
+        compressed_dict = {}
+        _LOGGER.debug(
+            f"Compressing model with {len(model_state)} parameterized layers..."
+        )
+
+        for name, value in tqdm(model_state.items(), desc="Compressing model"):
+            if name.endswith(".weight"):
+                prefix = name.removesuffix(".weight")
+                scale = model_state.get(merge_names(prefix, "weight_scale"), None)
+                zp = model_state.get(merge_names(prefix, "weight_zero_point"), None)
+                if scale is not None and zp is not None:
+                    # weight is quantized, compress it
+                    quant_args = model_quant_args[prefix]
+                    try:
+                        bit_depth = torch.finfo(value.dtype).bits
+                    except TypeError:
+                        bit_depth = torch.iinfo(value.dtype).bits
+                    if bit_depth > quant_args.num_bits:
+                        # only quantize if not already quantized
+                        value = quantize(
+                            x=value,
+                            scale=scale,
+                            zero_point=zp,
+                            args=quant_args,
+                            dtype=torch.int8,
+                        )
+
+            compressed_dict[name] = value.to("cpu")
+
+        return compressed_dict
+
+    def decompress(
+        self, path_to_model_or_tensors: str, device: str = "cpu"
+    ) -> Generator[Tuple[str, Tensor], None, None]:
+        weight_mappings = get_nested_weight_mappings(
+            path_to_model_or_tensors, self.COMPRESSION_PARAM_NAMES
+        )
+        for weight_name in weight_mappings.keys():
+            weight_data = {}
+            for param_name, safe_path in weight_mappings[weight_name].items():
+                full_name = merge_names(weight_name, param_name)
+                with safe_open(safe_path, framework="pt", device=device) as f:
+                    weight_data[param_name] = f.get_tensor(full_name)
+
+            if len(weight_data) == len(self.COMPRESSION_PARAM_NAMES):
+                decompressed = dequantize(
+                    x_q=weight_data["weight"],
+                    scale=weight_data["weight_scale"],
+                    zero_point=weight_data["weight_zero_point"],
+                )
+                yield merge_names(weight_name, "weight"), decompressed

compressed_tensors/compressors/model_compressor.py

@@ -0,0 +1,264 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import logging
+import operator
+import os
+from typing import Dict, Optional, Union
+
+from compressed_tensors.base import (
+    COMPRESSION_CONFIG_NAME,
+    QUANTIZATION_CONFIG_NAME,
+    SPARSITY_CONFIG_NAME,
+)
+from compressed_tensors.compressors import Compressor
+from compressed_tensors.config import SparsityCompressionConfig
+from compressed_tensors.quantization import (
+    QuantizationConfig,
+    QuantizationStatus,
+    apply_quantization_config,
+    load_pretrained_quantization,
+)
+from compressed_tensors.quantization.utils import (
+    is_module_quantized,
+    iter_named_leaf_modules,
+)
+from compressed_tensors.utils import get_safetensors_folder
+from torch import Tensor
+from torch.nn import Module, Parameter
+from tqdm import tqdm
+from transformers import AutoConfig
+from transformers.file_utils import CONFIG_NAME
+
+
+__all__ = ["ModelCompressor"]
+
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+class ModelCompressor:
+    """
+    Handles compression and decompression of a model with a sparsity config and/or
+    quantization config.
+
+    Compression LifeCycle
+        - compressor = ModelCompressor.from_pretrained_model(model)
+        - compressed_state_dict = compressor.compress(model, state_dict)
+            - compressor.quantization_compressor.compress(model, state_dict)
+            - compressor.sparsity_compressor.compress(model, state_dict)
+        - model.save_pretrained(output_dir, state_dict=compressed_state_dict)
+        - compressor.update_config(output_dir)
+
+    Decompression LifeCycle
+        - compressor = ModelCompressor.from_pretrained(comp_model_path)
+        - model = AutoModel.from_pretrained(comp_model_path)
+        - compressor.decompress(comp_model_path, model)
+            - compressor.sparsity_compressor.decompress(comp_model_path, model)
+            - compressor.quantization_compressor.decompress(comp_model_path, model)
+
+    :param sparsity_config: config specifying sparsity compression parameters
+    :param quantization_config: config specifying quantization compression parameters
+    """
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str,
+    ) -> Optional["ModelCompressor"]:
+        """
+        Given a path to a model config, extract the sparsity and/or quantization
+        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
+        """
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
+        compression_config = getattr(config, COMPRESSION_CONFIG_NAME, None)
+        if compression_config is None:
+            return None
+
+        sparsity_config = compression_config.get(SPARSITY_CONFIG_NAME, None)
+        quantization_config = compression_config.get(QUANTIZATION_CONFIG_NAME, None)
+
+        if sparsity_config is None and quantization_config is None:
+            return None
+
+        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:
+            quantization_config = QuantizationConfig.parse_obj(quantization_config)
+
+        return cls(
+            sparsity_config=sparsity_config, quantization_config=quantization_config
+        )
+
+    @classmethod
+    def from_pretrained_model(
+        cls,
+        model: Module,
+        sparsity_config: Union[SparsityCompressionConfig, str, None] = None,
+        quantization_format: Optional[str] = None,
+    ) -> Optional["ModelCompressor"]:
+        """
+        Given a pytorch model and optional sparsity and/or quantization configs,
+        load the appropriate compressors
+
+        :param model: pytorch model to target for compression
+        :param sparsity_config: a filled in sparsity config or string corresponding
+            to a sparsity compression algorithm
+        :param quantization_format: string corresponding to a quantization compression
+            algorithm
+        :return: compressor for the extracted configs
+        """
+        quantization_config = QuantizationConfig.from_pretrained(
+            model, format=quantization_format
+        )
+
+        if isinstance(sparsity_config, str):  # we passed in a sparsity format
+            sparsity_config = SparsityCompressionConfig.load_from_registry(
+                sparsity_config
+            )
+
+        if sparsity_config is None and quantization_config is None:
+            return None
+
+        return cls(
+            sparsity_config=sparsity_config, quantization_config=quantization_config
+        )
+
+    def __init__(
+        self,
+        sparsity_config: Optional[SparsityCompressionConfig] = None,
+        quantization_config: Optional[QuantizationConfig] = None,
+    ):
+        self.sparsity_config = sparsity_config
+        self.quantization_config = quantization_config
+        self.sparsity_compressor = None
+        self.quantization_compressor = None
+
+        if sparsity_config is not None:
+            self.sparsity_compressor = Compressor.load_from_registry(
+                sparsity_config.format, config=sparsity_config
+            )
+        if quantization_config is not None:
+            self.quantization_compressor = Compressor.load_from_registry(
+                quantization_config.format, config=quantization_config
+            )
+
+    def compress(
+        self, model: Module, state_dict: Optional[Dict[str, Tensor]] = None
+    ) -> Dict[str, Tensor]:
+        """
+        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
+        :return: compressed state dict
+        """
+        if state_dict is None:
+            state_dict = model.state_dict()
+
+        compressed_state_dict = state_dict
+        quantized_modules_to_args = _get_weight_arg_mappings(model)
+        if self.quantization_compressor is not None:
+            compressed_state_dict = self.quantization_compressor.compress(
+                state_dict, model_quant_args=quantized_modules_to_args
+            )
+
+        if self.sparsity_compressor is not None:
+            compressed_state_dict = self.sparsity_compressor.compress(
+                compressed_state_dict
+            )
+
+        return compressed_state_dict
+
+    def decompress(self, model_path: str, model: Module):
+        """
+        Overwrites the weights in model with weights decompressed from model_path
+
+        :param model_path: path to compressed weights
+        :param model: pytorch model to load decompressed weights into
+        """
+        model_path = get_safetensors_folder(model_path)
+        if self.sparsity_compressor is not None:
+            dense_gen = self.sparsity_compressor.decompress(model_path)
+            self._replace_weights(dense_gen, model)
+            setattr(model, SPARSITY_CONFIG_NAME, self.sparsity_compressor.config)
+
+        if self.quantization_compressor is not None:
+            apply_quantization_config(model, self.quantization_config)
+            load_pretrained_quantization(model, model_path)
+            dense_gen = self.quantization_compressor.decompress(model_path)
+            self._replace_weights(dense_gen, model)
+
+            def update_status(module):
+                module.quantization_status = QuantizationStatus.FROZEN
+
+            model.apply(update_status)
+            setattr(model, QUANTIZATION_CONFIG_NAME, self.quantization_config)
+
+    def update_config(self, save_directory: str):
+        """
+        Update the model config located at save_directory with compression configs
+        for sparsity and/or quantization
+
+        :param save_directory: path to a folder containing a HF model config
+        """
+        config_file_path = os.path.join(save_directory, CONFIG_NAME)
+        if not os.path.exists(config_file_path):
+            _LOGGER.warning(
+                f"Could not find a valid model config file in "
+                f"{save_directory}. Compression config will not be saved."
+            )
+            return
+
+        with open(config_file_path, "r") as config_file:
+            config_data = json.load(config_file)
+
+        config_data[COMPRESSION_CONFIG_NAME] = {}
+        if self.quantization_config is not None:
+            quant_config_data = self.quantization_config.model_dump()
+            config_data[COMPRESSION_CONFIG_NAME][
+                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][
+                SPARSITY_CONFIG_NAME
+            ] = sparsity_config_data
+
+        with open(config_file_path, "w") as config_file:
+            json.dump(config_data, config_file, indent=2, sort_keys=True)
+
+    def _replace_weights(self, dense_weight_generator, model):
+        for name, data in tqdm(dense_weight_generator, desc="Decompressing model"):
+            # loading the decompressed weights into the model
+            model_device = operator.attrgetter(name)(model).device
+            data_new = Parameter(data.to(model_device))
+            data_old = operator.attrgetter(name)(model)
+            data_old.data = data_new.data
+
+
+def _get_weight_arg_mappings(model: Module) -> Dict:
+    quantized_modules_to_args = {}
+    for name, submodule in iter_named_leaf_modules(model):
+        if is_module_quantized(submodule):
+            if submodule.quantization_scheme.weights is not None:
+                quantized_modules_to_args[name] = submodule.quantization_scheme.weights
+
+    return quantized_modules_to_args