compressed-tensors 0.3.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

compressed_tensors/compressors/__init__.py

@@ -16,10 +16,5 @@
 
 from .base import ModelCompressor
 from .dense import DenseCompressor
-from .helpers import (
-    infer_compressor_from_model_config,
-    load_compressed,
-    save_compressed,
-    save_compressed_model,
-)
+from .helpers import load_compressed, save_compressed, save_compressed_model
 from .sparse_bitmask import BitmaskCompressor, BitmaskTensor

compressed_tensors/compressors/base.py

@@ -22,6 +22,7 @@ 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
 
 
 __all__ = ["ModelCompressor"]
@@ -34,6 +35,29 @@ class ModelCompressor(RegistryMixin):
     :param config: config specifying compression parameters
     """
 
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_name_or_path: str
+    ) -> Optional["ModelCompressor"]:
+        """
+        Given a path to a model config, extract a sparsity config if it exists and
+        return the associated ModelCompressor
+
+        :param pretrained_model_name_or_path: path to model config on disk or HF hub
+        :return: matching compressor if config contains a sparsity config
+        """
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
+        sparsity_config = getattr(config, SPARSITY_CONFIG_NAME, None)
+        if sparsity_config is None:
+            return None
+
+        format = sparsity_config.get("format")
+        sparsity_config = CompressionConfig.load_from_registry(
+            format, **sparsity_config
+        )
+        compressor = cls.load_from_registry(format, config=sparsity_config)
+        return compressor
+
     def __init__(self, config: Optional[CompressionConfig] = None):
         self.config = config
 
@@ -47,7 +71,7 @@ class ModelCompressor(RegistryMixin):
         raise NotImplementedError()
 
     def decompress(
-        self, path_to_model_or_tensors: str
+        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

compressed_tensors/compressors/dense.py

@@ -29,6 +29,6 @@ class DenseCompressor(ModelCompressor):
         return model_state
 
     def decompress(
-        self, path_to_model_or_tensors: str, device: str
+        self, path_to_model_or_tensors: str, device: str = "cpu"
     ) -> Generator[Tuple[str, Tensor], None, None]:
         return iter([])

compressed_tensors/compressors/helpers.py

@@ -16,45 +16,21 @@ from pathlib import Path
 from typing import Dict, Generator, Optional, Tuple, Union
 
 import torch
-from compressed_tensors.base import SPARSITY_CONFIG_NAME
 from compressed_tensors.compressors import ModelCompressor
 from compressed_tensors.config import CompressionConfig, CompressionFormat
 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
-from transformers import AutoConfig
 
 
 __all__ = [
-    "infer_compressor_from_model_config",
     "load_compressed",
     "save_compressed",
     "save_compressed_model",
 ]
 
 
-def infer_compressor_from_model_config(
-    pretrained_model_name_or_path: str,
-) -> Optional[ModelCompressor]:
-    """
-    Given a path to a model config, extract a sparsity config if it exists and return
-    the associated ModelCompressor
-
-    :param pretrained_model_name_or_path: path to model config on disk or HF hub
-    :return: matching compressor if config contains a sparsity config
-    """
-    config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
-    sparsity_config = getattr(config, SPARSITY_CONFIG_NAME, None)
-    if sparsity_config is None:
-        return None
-
-    format = sparsity_config.get("format")
-    sparsity_config = CompressionConfig.load_from_registry(format, **sparsity_config)
-    compressor = ModelCompressor.load_from_registry(format, config=sparsity_config)
-    return compressor
-
-
 def save_compressed(
     tensors: Dict[str, Tensor],
     save_path: Union[str, Path],

compressed_tensors/compressors/sparse_bitmask.py

@@ -75,8 +75,9 @@ class BitmaskCompressor(ModelCompressor):
         self, path_to_model_or_tensors: str, device: str = "cpu"
     ) -> Generator[Tuple[str, Tensor], None, None]:
         """
-        Reads a bitmask compressed state dict located at path_to_model_or_tensors
-        and returns a generator for sequentially decompressing back to a dense state dict
+        Reads a bitmask 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

compressed_tensors/quantization/lifecycle/forward.py

@@ -111,7 +111,7 @@ def wrap_module_forward_quantized(module: Module, scheme: QuantizationScheme):
 
 
 def _maybe_calibrate_or_quantize(
-    module: Module, value: Module, base_name: str, args: "QuantizationArgs"
+    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 {
@@ -120,17 +120,23 @@ def _maybe_calibrate_or_quantize(
     }:
         return value
 
-    device = next(module.parameters()).device
-    scale = getattr(module, f"{base_name}_scale")
-    zero_point = getattr(module, f"{base_name}_zero_point")
-
-    if module.quantization_status == QuantizationStatus.CALIBRATION:
-        # get observer and get new quant params from observation
+    if args.dynamic:
+        # dynamic quantization - get scale and zero point directly from observer
         observer = getattr(module, f"{base_name}_observer")
-        updated_scale, updated_zero_point = observer(value)
-
-        # update scale and zero point
-        scale.data = updated_scale.to(device)
-        zero_point.data = updated_zero_point.to(device)
+        scale, zero_point = observer(value)
+    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")
+
+        if module.quantization_status == QuantizationStatus.CALIBRATION:
+            # calibration mode - get new quant params from observer
+            observer = getattr(module, f"{base_name}_observer")
+            updated_scale, updated_zero_point = observer(value)
+
+            # update scale and zero point
+            device = next(module.parameters()).device
+            scale.data = updated_scale.to(device)
+            zero_point.data = updated_zero_point.to(device)
 
     return fake_quantize(value, scale, zero_point, args)

compressed_tensors/quantization/lifecycle/frozen.py

@@ -30,17 +30,17 @@ def freeze_module_quantization(module: Module):
 
     :param module: module to freeze quantization for
     """
-    if not getattr(module, "quantization_scheme", None):
+    scheme = getattr(module, "quantization_scheme", None)
+    if not scheme:
         # no quantization scheme nothing to do
         return
 
-    # delete observers from module
-    observer_names = []
-    for submodule_name, _ in module.named_modules():
-        if "." not in submodule_name and submodule_name.endswith("_observer"):
-            # delete any observers that belong directly to this module
-            observer_names.append(submodule_name)
-    for observer_name in observer_names:
-        delattr(module, observer_name)
+    # delete observers from module if not dynamic
+    if scheme.input_activations and not scheme.input_activations.dynamic:
+        delattr(module, "input_observer")
+    if scheme.weights and not scheme.weights.dynamic:
+        delattr(module, "weight_observer")
+    if scheme.output_activations and not scheme.output_activations.dynamic:
+        delattr(module, "output_observer")
 
     module.quantization_status = QuantizationStatus.FROZEN

compressed_tensors/quantization/lifecycle/initialize.py

@@ -80,6 +80,13 @@ def initialize_module_for_quantization(
 def _initialize_scale_zero_point_observer(
     module: Module, base_name: str, quantization_args: QuantizationArgs
 ):
+    # initialize observer module and attach as submodule
+    observer = quantization_args.get_observer()
+    module.register_module(f"{base_name}_observer", observer)
+
+    if quantization_args.dynamic:
+        return  # no need to register a scale and zero point for a dynamic observer
+
     device = next(module.parameters()).device
 
     # initializes empty scale and zero point parameters for the module
@@ -90,7 +97,3 @@ def _initialize_scale_zero_point_observer(
         torch.empty(0, device=device, dtype=int), requires_grad=False
     )
     module.register_parameter(f"{base_name}_zero_point", init_zero_point)
-
-    # initialize observer module and attach as submodule
-    observer = quantization_args.get_observer()
-    module.register_module(f"{base_name}_observer", observer)

compressed_tensors/quantization/observers/memoryless.py

@@ -23,10 +23,10 @@ from torch import FloatTensor, IntTensor, Tensor
 __all__ = ["MemorylessObserver"]
 
 
-@Observer.register("memoryless")
+@Observer.register("memoryless", alias=["dynamic"])
 class MemorylessObserver(Observer):
     """
-    Implements a dynamic quantization observer that sets the scale and
+    Implements a quantization observer that sets the scale and
     zero point based on the latest observed value without tracking state
     """
 

compressed_tensors/quantization/quant_args.py

@@ -53,6 +53,11 @@ class QuantizationArgs(BaseModel):
     :param group_size: group length to use for the group strategy
     :param block_structure: 2d block structure to use for the block strategy, must be
     of the format "2x4", "8x16", etc.
+    :param dynamic: set True to perform dynamic quantization - values will not be
+        calibrated during calibration phase, instead during inference new quantization
+        ranges will be observed with every sample. Defaults to False for static
+        quantization. Note that enabling dynamic quantization will change the default
+        observer to a memoryless one
     """
 
     num_bits: int = 8
@@ -61,6 +66,7 @@ class QuantizationArgs(BaseModel):
     strategy: QuantizationStrategy = QuantizationStrategy.TENSOR
     group_size: Optional[int] = None
     block_structure: Optional[str] = None
+    dynamic: bool = False
     observer: str = Field(
         default="minmax",
         description=(
@@ -82,4 +88,9 @@ class QuantizationArgs(BaseModel):
         """
         from compressed_tensors.quantization.observers.base import Observer
 
+        if self.observer == "minmax" and self.dynamic:
+            # override defualt observer for dynamic, you never want minmax which
+            # keeps state across samples for dynamic
+            self.observer = "memoryless"
+
         return Observer.load_from_registry(self.observer, quantization_args=self)

compressed_tensors/utils/helpers.py

@@ -30,11 +30,25 @@ from transformers import AutoConfig
 
 __all__ = [
     "infer_compressor_from_model_config",
+    "infer_compression_config_from_model_config",
     "load_compressed",
     "save_compressed",
     "save_compressed_model",
 ]
 
+def infer_compressor_from_model_config(
+    pretrained_model_name_or_path: str,
+) -> Optional[CompressionConfig]:
+    """
+    Given a path to a model config, extract a sparsity config if it exists and return
+    the associated CompressionConfig
+    
+    :param pretrained_model_name_or_path: path to model config on disk or HF hub
+    :return: matching compression config if config contains a sparsity config
+    """
+    config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
+    return getattr(config, SPARSITY_CONFIG_NAME, None)
+
 
 def infer_compressor_from_model_config(
     pretrained_model_name_or_path: str,
@@ -46,14 +60,8 @@ def infer_compressor_from_model_config(
     :param pretrained_model_name_or_path: path to model config on disk or HF hub
     :return: matching compressor if config contains a sparsity config
     """
-    config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
-    sparsity_config = getattr(config, SPARSITY_CONFIG_NAME, None)
-    if sparsity_config is None:
-        return None
-
-    format = sparsity_config.get("format")
-    sparsity_config = CompressionConfig.load_from_registry(format, **sparsity_config)
-    compressor = ModelCompressor.load_from_registry(format, config=sparsity_config)
+    sparsity_config = infer_compressor_from_model_config(pretrained_model_name_or_path)
+    compressor = ModelCompressor.load_from_registry(sparsity_config.format, config=sparsity_config)
     return compressor
 
 

{compressed_tensors-0.3.1.dist-info → compressed_tensors-0.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: compressed-tensors
-Version: 0.3.1
+Version: 0.3.2
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.
@@ -103,4 +103,7 @@ save_compressed_model(model, "compressed_model.safetensors", compression_format=
 state_dict = dict(load_compressed("compressed_model.safetensors", compression_config))
 ```
 
+For more in-depth tutorial on bitmask compression, refer to the [notebook](https://github.com/neuralmagic/compressed-tensors/blob/d707c5b84bc3fef164aebdcd97cb6eaa571982f8/examples/bitmask_compression.ipynb).
+
+
 

{compressed_tensors-0.3.1.dist-info → compressed_tensors-0.3.2.dist-info}/RECORD

@@ -1,38 +1,38 @@
 compressed_tensors/__init__.py,sha256=SV1csvHUVCd8kHXz6UDZim1HZ_fAVG3vfk-j_4Bb6hY,789
 compressed_tensors/base.py,sha256=8zbgK87LpHkKoSknM55svXCT4E4dLLjPijwF9HfzmsQ,717
-compressed_tensors/compressors/__init__.py,sha256=3ZHKWSIWTjMx8XXgLtoP9JaVaCTvRecguLZTxLAAkKk,898
-compressed_tensors/compressors/base.py,sha256=F1smyJ6x2Sfq43tuP0QE9wZuhVqnewq-XUFPMtdU9yQ,2936
-compressed_tensors/compressors/dense.py,sha256=_VTusI3XjaY-zOdB_d7z4zOgPTJi9TJZZHF13g9ulS4,1263
-compressed_tensors/compressors/helpers.py,sha256=kSseqbwnu3JHZUKH8u4kQo5bmd87FvCcmWe0u2ikysA,6421
-compressed_tensors/compressors/sparse_bitmask.py,sha256=PYAK_Hcy2T57zlbpwl1FYkslluIr2x-d0Rh048YAtpI,8639
+compressed_tensors/compressors/__init__.py,sha256=UcHp0CwUBJoS2MBN6mLUT7B3uRf1TEoRGbME7gLPD38,841
+compressed_tensors/compressors/base.py,sha256=CqQo00ZIkAWpy0yVux5TXhK7WK_6Ws6qb5mCAvIoxB4,3902
+compressed_tensors/compressors/dense.py,sha256=ig9lItmyCX5-VzgMuUqea-s8fHsTjPj5-0VIsPLl0g0,1271
+compressed_tensors/compressors/helpers.py,sha256=wstgUEUYUCTMMu6G1YLF9G7vXqIJPj3MsWhqwU4J6Vw,5458
+compressed_tensors/compressors/sparse_bitmask.py,sha256=VbCGFC4Q3nfhKWqJdkM4hKQmrZZqLT8wLNxbVt4kLSs,8647
 compressed_tensors/config/__init__.py,sha256=ZBqWn3r6ku1qfmlHHYp0mQueY0i7Pwhr9rbQk9dDlMc,704
 compressed_tensors/config/base.py,sha256=IP-3Y416w-811WozDzKHycIBXjdlG4Ddy7vpbwhOPD8,1373
 compressed_tensors/config/dense.py,sha256=xtkri7DkP7USu44FnSoTgTSqdGegCBtjRf3DfblSEL0,1311
 compressed_tensors/config/sparse_bitmask.py,sha256=y8fmQaOoGjIiI4FR6BJjfIqisAcqNQ_zjKyjT75bXwY,1284
 compressed_tensors/quantization/__init__.py,sha256=83J5bPB7PavN2TfCoW7_vEDhfYpm4TDrqYO9vdSQ5bk,760
-compressed_tensors/quantization/quant_args.py,sha256=dxSrq0_88ORQXcyIMYqoMZJvYEjnqdYl37f7lgZQqhw,2742
+compressed_tensors/quantization/quant_args.py,sha256=cZhe5hRmvU_HnnUDw1kmqzMAGFb0r5t0IL2cobBNw28,3371
 compressed_tensors/quantization/quant_config.py,sha256=DWx8ae3gDlw99zAn3MUN9I4qeksbbmITmOXHRynqPB8,6650
 compressed_tensors/quantization/quant_scheme.py,sha256=X3oqmZPiIKtX5tEKKUj-0N6hB68NeiU2b1GcQEQPadQ,1480
 compressed_tensors/quantization/lifecycle/__init__.py,sha256=fM9XBtPgJX6z54PTm3Sd0SpK5od95ibwaSf2FFR8DqE,772
 compressed_tensors/quantization/lifecycle/apply.py,sha256=WXUL3q1g0s244k0wuqGYZPXTXiscdyrp7RScN2j_KGA,6651
 compressed_tensors/quantization/lifecycle/calibration.py,sha256=mLns4jlaWmBwOW8Jtlm5bMX-JET1AiZYUBO7qa-XuxI,1776
-compressed_tensors/quantization/lifecycle/forward.py,sha256=hnjk7pocZLDhLdMx237FKayYdvsdKbYSjTmSN5xbQO8,4599
-compressed_tensors/quantization/lifecycle/frozen.py,sha256=NHNmlDIaxurifqeI_qZC8xa4BstQsBNdOCXJjRzAfNU,1596
-compressed_tensors/quantization/lifecycle/initialize.py,sha256=8pifqZQSgVqWYI_Qtv6QfBICPbCTFHy48OWPeQsxEHQ,3578
+compressed_tensors/quantization/lifecycle/forward.py,sha256=COTlfH92JkwM9Vd6m07tK_dhSiC77SFS0-MHU_DbQko,4941
+compressed_tensors/quantization/lifecycle/frozen.py,sha256=FF7BleuOUX46Egk7F1ZE5r4fjWt9jG5-tO8BjXU1r78,1606
+compressed_tensors/quantization/lifecycle/initialize.py,sha256=U6g9qifSF6pagQZQZEwd-rwWC6uQ_dZXn1wg6nr1Abg,3697
 compressed_tensors/quantization/observers/__init__.py,sha256=DNH31NQYrIBBcmHsMyFA6whh4pbRsLwuNa6L8AeXaGc,745
 compressed_tensors/quantization/observers/base.py,sha256=O76dAxkin7bB602e9kjmxc84p71-PxBtjIq5L69xplI,2786
 compressed_tensors/quantization/observers/helpers.py,sha256=SxvOf9zwZ9NDRC3E4Xm7z3RqHcbcPtCABLKX9GnGGHM,2109
-compressed_tensors/quantization/observers/memoryless.py,sha256=3f6bUlcf5mzOHPkTRhoQ7Zd8xu_pUmj8e3Y85fGysSU,1848
+compressed_tensors/quantization/observers/memoryless.py,sha256=ZHTPh4aURE8LvHBFaP--HIC2JanMX5-VRdIkE2JHthw,1859
 compressed_tensors/quantization/observers/min_max.py,sha256=uAcZd5aY6WKM-KumTb2ybX28s8iKGVy6Nrje5Sddqew,2439
 compressed_tensors/quantization/utils/__init__.py,sha256=VdtEmP0bvuND_IGQnyqUPc5lnFp-1_yD7StKSX4x80w,656
 compressed_tensors/quantization/utils/helpers.py,sha256=N_wYfrPcFr__Q1mn6mHoNUTclwpTW8P5PDHkR7GvXWo,3694
 compressed_tensors/registry/__init__.py,sha256=FwLSNYqfIrb5JD_6OK_MT4_svvKTN_nEhpgQlQvGbjI,658
 compressed_tensors/registry/registry.py,sha256=fxjOjh2wklCvJhQxwofdy-zV8q7MkQ85SLG77nml2iA,11890
 compressed_tensors/utils/__init__.py,sha256=5DrYjoZbaEvSkJcC-GRSbM_RBHVF4tG9gMd3zsJnjLw,665
-compressed_tensors/utils/helpers.py,sha256=wLgiPrk7Vn29AijOGQGk3UnXItRd1jpROS6FxHoC4VQ,5530
+compressed_tensors/utils/helpers.py,sha256=b2zQimHNn3emCgUGsVYMpaWQJH_tR9Uso819bU5r78Y,5909
 compressed_tensors/utils/safetensors_load.py,sha256=wo9UirGrGlenBqZeqotvpCT7D5MEdjCo2J3HeRaIFoU,8502
-compressed_tensors-0.3.1.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-compressed_tensors-0.3.1.dist-info/METADATA,sha256=vpGRbjHWdPUTl9HFoDxIkwAKQJNpff75P4pKC3nJE4A,3850
-compressed_tensors-0.3.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-compressed_tensors-0.3.1.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
-compressed_tensors-0.3.1.dist-info/RECORD,,
+compressed_tensors-0.3.2.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+compressed_tensors-0.3.2.dist-info/METADATA,sha256=658VPFfv3kqbIbEjY0tJNOkRoTdeZApVUbqcC1vtMLM,4060
+compressed_tensors-0.3.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+compressed_tensors-0.3.2.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
+compressed_tensors-0.3.2.dist-info/RECORD,,