compressed-tensors 0.5.0__tar.gz → 0.7.0__tar.gz

{compressed_tensors-0.5.0 → compressed-tensors-0.7.0}/PKG-INFO

@@ -1,51 +1,53 @@
 Metadata-Version: 2.1
 Name: compressed-tensors
-Version: 0.5.0
+Version: 0.7.0
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.
 Author-email: support@neuralmagic.com
 License: Apache 2.0
 Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: torch>=1.7.0
-Requires-Dist: transformers
-Requires-Dist: accelerate
-Requires-Dist: pydantic>=2.0
 Provides-Extra: dev
-Requires-Dist: black==22.12.0; extra == "dev"
-Requires-Dist: isort==5.8.0; extra == "dev"
-Requires-Dist: wheel>=0.36.2; extra == "dev"
-Requires-Dist: flake8>=3.8.3; extra == "dev"
-Requires-Dist: pytest>=6.0.0; extra == "dev"
-Requires-Dist: nbconvert>=7.16.3; extra == "dev"
-
-# compressed_tensors
-
-This repository extends a [safetensors](https://github.com/huggingface/safetensors) format to efficiently store sparse and/or quantized tensors on disk. `compressed-tensors` format supports multiple compression types to minimize the disk space and facilitate the tensor manipulation.
+Provides-Extra: accelerate
+License-File: LICENSE
 
-## Motivation
+# compressed-tensors
 
-### Reduce disk space by saving sparse tensors in a compressed format
+The `compressed-tensors` library extends the [safetensors](https://github.com/huggingface/safetensors) format, providing a versatile and efficient way to store and manage compressed tensor data. This library supports various quantization and sparsity schemes, making it a unified format for handling different model optimizations like GPTQ, AWQ, SmoothQuant, INT8, FP8, SparseGPT, and more.
 
-The compressed format stores the data much more efficiently by taking advantage of two properties of tensors:
+## Why `compressed-tensors`?
 
-- Sparse tensors -> due to a large number of entries that are equal to zero.
-- Quantized -> due to their low precision representation.
+As model compression becomes increasingly important for efficient deployment of LLMs, the landscape of quantization and compression techniques has become increasingly fragmented.
+Each method often comes with its own storage format and loading procedures, making it challenging to work with multiple techniques or switch between them.
+`compressed-tensors` addresses this by providing a single, extensible format that can represent a wide variety of compression schemes. 
 
-### Introduce an elegant interface to save/load compressed tensors
+* **Unified Checkpoint Format**: Supports various compression schemes in a single, consistent format.
+* **Wide Compatibility**: Works with popular quantization methods like GPTQ, SmoothQuant, and FP8. See [llm-compressor](https://github.com/vllm-project/llm-compressor)
+* **Flexible Quantization Support**: 
+  * Weight-only quantization (e.g., W4A16, W8A16, WnA16)
+  * Activation quantization (e.g., W8A8)
+  * KV cache quantization
+  * Non-uniform schemes (different layers can be quantized in different ways!)
+* **Sparsity Support**: Handles both unstructured and semi-structured (e.g., 2:4) sparsity patterns.
+* **Open-Source Integration**: Designed to work seamlessly with Hugging Face models and PyTorch.
 
-The library provides the user with the ability to compress/decompress tensors. The properties of tensors are defined by human-readable configs, allowing the users to understand the compression format at a quick glance.
+This allows developers and researchers to easily experiment with composing different quantization methods, simplify model deployment pipelines, and reduce the overhead of supporting multiple compression formats in inference engines.
 
 ## Installation
 
-### Pip
+### From [PyPI](https://pypi.org/project/compressed-tensors)
 
+Stable release:
 ```bash
 pip install compressed-tensors
 ```
 
-### From source
+Nightly release:
+```bash
+pip install compressed-tensors-nightly
+```
+
+### From Source
 
 ```bash
 git clone https://github.com/neuralmagic/compressed-tensors

{compressed_tensors-0.5.0 → compressed-tensors-0.7.0}/README.md

@@ -1,29 +1,40 @@
-# compressed_tensors
+# compressed-tensors
 
-This repository extends a [safetensors](https://github.com/huggingface/safetensors) format to efficiently store sparse and/or quantized tensors on disk. `compressed-tensors` format supports multiple compression types to minimize the disk space and facilitate the tensor manipulation.
+The `compressed-tensors` library extends the [safetensors](https://github.com/huggingface/safetensors) format, providing a versatile and efficient way to store and manage compressed tensor data. This library supports various quantization and sparsity schemes, making it a unified format for handling different model optimizations like GPTQ, AWQ, SmoothQuant, INT8, FP8, SparseGPT, and more.
 
-## Motivation
+## Why `compressed-tensors`?
 
-### Reduce disk space by saving sparse tensors in a compressed format
+As model compression becomes increasingly important for efficient deployment of LLMs, the landscape of quantization and compression techniques has become increasingly fragmented.
+Each method often comes with its own storage format and loading procedures, making it challenging to work with multiple techniques or switch between them.
+`compressed-tensors` addresses this by providing a single, extensible format that can represent a wide variety of compression schemes. 
 
-The compressed format stores the data much more efficiently by taking advantage of two properties of tensors:
+* **Unified Checkpoint Format**: Supports various compression schemes in a single, consistent format.
+* **Wide Compatibility**: Works with popular quantization methods like GPTQ, SmoothQuant, and FP8. See [llm-compressor](https://github.com/vllm-project/llm-compressor)
+* **Flexible Quantization Support**: 
+  * Weight-only quantization (e.g., W4A16, W8A16, WnA16)
+  * Activation quantization (e.g., W8A8)
+  * KV cache quantization
+  * Non-uniform schemes (different layers can be quantized in different ways!)
+* **Sparsity Support**: Handles both unstructured and semi-structured (e.g., 2:4) sparsity patterns.
+* **Open-Source Integration**: Designed to work seamlessly with Hugging Face models and PyTorch.
 
-- Sparse tensors -> due to a large number of entries that are equal to zero.
-- Quantized -> due to their low precision representation.
-
-### Introduce an elegant interface to save/load compressed tensors
-
-The library provides the user with the ability to compress/decompress tensors. The properties of tensors are defined by human-readable configs, allowing the users to understand the compression format at a quick glance.
+This allows developers and researchers to easily experiment with composing different quantization methods, simplify model deployment pipelines, and reduce the overhead of supporting multiple compression formats in inference engines.
 
 ## Installation
 
-### Pip
+### From [PyPI](https://pypi.org/project/compressed-tensors)
 
+Stable release:
 ```bash
 pip install compressed-tensors
 ```
 
-### From source
+Nightly release:
+```bash
+pip install compressed-tensors-nightly
+```
+
+### From Source
 
 ```bash
 git clone https://github.com/neuralmagic/compressed-tensors

{compressed_tensors-0.5.0 → compressed-tensors-0.7.0}/setup.py

@@ -46,10 +46,13 @@ def _setup_packages() -> List:
     )
     
 def _setup_install_requires() -> List:
-    return ["torch>=1.7.0", "transformers", "accelerate", "pydantic>=2.0"]
+    return ["torch>=1.7.0", "transformers", "pydantic>=2.0"]
 
 def _setup_extras() -> Dict:
-    return {"dev": ["black==22.12.0", "isort==5.8.0", "wheel>=0.36.2", "flake8>=3.8.3", "pytest>=6.0.0", "nbconvert>=7.16.3"]}
+    return {
+        "dev": ["black==22.12.0", "isort==5.8.0", "wheel>=0.36.2", "flake8>=3.8.3", "pytest>=6.0.0", "nbconvert>=7.16.3"],
+        "accelerate": ["accelerate"]
+    }
 
 setup(
     name=_PACKAGE_NAME,

{compressed_tensors-0.5.0 → compressed-tensors-0.7.0}/src/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-0.5.0 → compressed-tensors-0.7.0}/src/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-0.7.0/src/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 *
+from .helpers import *
+from .model_compressors import *
+from .quantized_compressors import *
+from .sparse_compressors import *
+from .sparse_quantized_compressors import *

compressed-tensors-0.7.0/src/compressed_tensors/compressors/base.py

@@ -0,0 +1,188 @@
+# 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 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 QuantizationArgs, QuantizationConfig
+from compressed_tensors.registry import RegistryMixin
+from torch import Tensor
+from torch.nn import Module
+
+
+__all__ = ["BaseCompressor"]
+
+
+class BaseCompressor(RegistryMixin, ABC):
+    """
+    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
+    """
+
+    def __init__(
+        self, config: Union[SparsityCompressionConfig, QuantizationConfig, None] = None
+    ):
+        self.config = config
+
+    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,
+    ) -> 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 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-0.5.0 → compressed-tensors-0.7.0}/src/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-0.7.0/src/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 *