compressed-tensors 0.11.1a20250929__py3-none-any.whl → 0.12.1__py3-none-any.whl

compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -344,6 +344,55 @@ class ModelCompressor:
                     format, config=quantization_config
                 )
 
+    def get_missing_module_keys(self, model: Module) -> List[str]:
+        """
+        Identifies the expected missing weight keys in the compressed state_dict.
+
+        When a model undergoes sparsity or quantization compression, certain
+        weight tensors may be absent from the checkpoint by virtue of compression.
+        This function determines which weight keys are missing based on the
+        applied compression techniques.
+
+        :param model: The PyTorch model to check for missing keys.
+        :return: A list of missing keys expected in the compressed state_dict.
+        """
+        missing_keys = set()
+
+        # Determine missing keys due to sparsity compression
+        if (
+            self.sparsity_compressor
+            and self.sparsity_config.format != CompressionFormat.dense.value
+        ):
+            sparse_targets = match_named_modules(
+                model=model,
+                targets=self.sparsity_config.targets,
+                ignore=self.sparsity_config.ignore,
+            )
+
+            missing_keys.update(
+                merge_names(target_name, "weight")
+                for target_name, _module in sparse_targets
+            )
+
+        # Determine missing keys due to pack quantization
+        if (
+            self.quantization_compressor
+            and self.quantization_config.format
+            == CompressionFormat.pack_quantized.value
+        ):
+            for scheme in self.quantization_config.config_groups.values():
+                quant_targets = match_named_modules(
+                    model=model,
+                    targets=scheme.targets,
+                    ignore=self.quantization_config.ignore,
+                )
+                missing_keys.update(
+                    merge_names(target_name, "weight")
+                    for target_name, _module in quant_targets
+                )
+
+        return list(missing_keys)
+
     def get_unexpected_file_keys(self, model: Module) -> List[str]:
         """
         Identifies extra keys introduced by the compression process in the

compressed_tensors/transform/factory/hadamard.py

@@ -52,7 +52,7 @@ class HadamardFactory(TransformFactory):
         :param args: defines how the transform will be applied to the module
         """
         assert hasattr(module, "weight")
-        size = get_transform_size(module, args.location, self.scheme.block_size)
+        size = get_transform_size(module, args.location, self.scheme.head_dim)
         exec_device = get_execution_device(module)
         device = get_offloaded_device(module)
         precision = self.scheme.precision if args.is_online() else torch.float64

compressed_tensors/transform/factory/matrix_multiply.py

@@ -51,7 +51,7 @@ class RandomMatrixFactory(TransformFactory):
         :param args: defines how the transform will be applied to the module
         """
         assert hasattr(module, "weight")
-        size = get_transform_size(module, args.location, self.scheme.block_size)
+        size = get_transform_size(module, args.location, self.scheme.head_dim)
         device = get_offloaded_device(module)
         precision = self.scheme.precision if args.is_online() else torch.float64
 

compressed_tensors/transform/transform_scheme.py

@@ -17,7 +17,7 @@ from typing import List, Optional
 import torch
 from compressed_tensors.transform import TransformArgs
 from compressed_tensors.utils import TorchDtype
-from pydantic import BaseModel, ConfigDict, Field, model_validator
+from pydantic import BaseModel, ConfigDict, Field
 
 
 __all__ = ["TransformScheme"]
@@ -36,8 +36,11 @@ class TransformScheme(BaseModel):
     :param randomize: True if uniquely randomized transform weights should be used,
         otherwise use identical transform weights where applicable
     :param requires_grad: True if weights include gradients for training
-    :param block_size: If set, the transform matrix will be block diagonal, with each
-        block being a square matrix of this size.
+    :param head_dim: If set, the transform matrix will be block diagonal with each
+        block being a square matrix of this size. The name head_dim was chosen because
+        some rotations need to be block-diagonal with block size equal to the head_dim,
+        but research has shown value in applying some rotations with smaller block size,
+        irrespective of head_dim.
     :param precision: Precision at which this transform should be applied during online
         rotations. Fused (offline) rotations are always performed in float64
     """
@@ -46,21 +49,7 @@ class TransformScheme(BaseModel):
     apply: List[TransformArgs] = Field(default_factory=list)
     randomize: bool = Field(default=False)
     requires_grad: bool = Field(default=False)
-    block_size: Optional[int] = Field(default=None)
-    head_dim: Optional[int] = Field(
-        default=None, deprecated="head_dim is deprecated, use block_size instead"
-    )
+    head_dim: Optional[int] = Field(default=None)
     precision: TorchDtype = Field(default=torch.float32)
 
-    @model_validator(mode="after")
-    def validate_model_after(model: "TransformScheme") -> "TransformScheme":
-        """
-        If head_dim is used instead of block_size, set block_size to head_dim
-        and remove head_dim
-        """
-        if model.block_size is None and model.head_dim is not None:
-            model.block_size = model.head_dim
-            model.head_dim = None
-        return model
-
     model_config = ConfigDict(extra="forbid")

compressed_tensors/version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.11.1.a20250929'
-__version_tuple__ = version_tuple = (0, 11, 1)
+__version__ = version = '0.12.1'
+__version_tuple__ = version_tuple = (0, 12, 1)

{compressed_tensors-0.11.1a20250929.dist-info → compressed_tensors-0.12.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.11.1a20250929
+Version: 0.12.1
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.

{compressed_tensors-0.11.1a20250929.dist-info → compressed_tensors-0.12.1.dist-info}/RECORD

@@ -1,12 +1,12 @@
 compressed_tensors/__init__.py,sha256=SRqNYFVvxAaLa4SImhoiIBKfoOSj7EUdx0CxXjGC2PA,884
 compressed_tensors/base.py,sha256=-gxWvDF4LCkyeDP8YlGzvBBKxo4Dk9h4NINPD61drFU,921
 compressed_tensors/logger.py,sha256=sTm1Od1cV0aDxBm3YN-PPvsOATxY_2tBV62TQE4HiPw,4032
-compressed_tensors/version.py,sha256=hBkazgCi53CqyaCNPq_qa6m80E1jS-UDDX5P6Om8KaA,523
+compressed_tensors/version.py,sha256=DzB47l1iPNTY8yuA60VuZvneNqNdRBd1svEjKxC37o4,513
 compressed_tensors/compressors/__init__.py,sha256=smSygTSfcfuujRrAXDc6uZm4L_ccV1tWZewqVnOb4lM,825
 compressed_tensors/compressors/base.py,sha256=nvWsv4xEw1Tkxkxth6TmHplDYXfBeP22xWxOsZERyDY,7204
 compressed_tensors/compressors/helpers.py,sha256=OK6qxX9j3bHwF9JfIYSGMgBJe2PWjlTA3byXKCJaTIQ,5431
 compressed_tensors/compressors/model_compressors/__init__.py,sha256=5RGGPFu4YqEt_aOdFSQYFYFDjcZFJN0CsMqRtDZz3Js,666
-compressed_tensors/compressors/model_compressors/model_compressor.py,sha256=xoQzV660vjjqwo_kek_xHJYIeqtt2hsD-QNAlqLHDvo,36144
+compressed_tensors/compressors/model_compressors/model_compressor.py,sha256=zPe3T0hyHuvIzSXHBWHPiqJ3sQcdVY3tgwF1aBeG7oo,38044
 compressed_tensors/compressors/quantized_compressors/__init__.py,sha256=KvaFBL_Q84LxRGJOV035M8OBoCkAx8kOkfphswgkKWk,745
 compressed_tensors/compressors/quantized_compressors/base.py,sha256=6GXhc2E5qFHnkcE3H2mx5CD2YtwfPXKUsvmsakzkkfA,10088
 compressed_tensors/compressors/quantized_compressors/naive_quantized.py,sha256=0ANDcuD8aXPqTYNPY6GnX9iS6eXJw6P0TzNV_rYS2l8,5369
@@ -46,11 +46,11 @@ compressed_tensors/transform/__init__.py,sha256=v2wfl4CMfA6KbD7Hxx_MbRev63y_6QLD
 compressed_tensors/transform/apply.py,sha256=YldYh3DGY09wJEKlFRSbg7-kRdCv7uFEuGkr5ZWCEXs,2984
 compressed_tensors/transform/transform_args.py,sha256=rVgReFp7wMXcYugkfd325e2tTFh8pGV3FnYTGCEv5jY,3429
 compressed_tensors/transform/transform_config.py,sha256=3YdtGcau3qkcapX9GMUiLuhQHFQZKFYT3eLgJGj1L6s,1204
-compressed_tensors/transform/transform_scheme.py,sha256=FJjkJcMHA6p6zh6yMunIgpzlgZgjkc39CqkeL65n_dQ,2786
+compressed_tensors/transform/transform_scheme.py,sha256=YR7Ri5hXWYQ8MSgMcpYg_BMyZtvXVPShz6iUf4giI8g,2475
 compressed_tensors/transform/factory/__init__.py,sha256=fH6rjBYAxuwrTzBTlTjTgCYNyh6TCvCqajCz4Im4YrA,617
 compressed_tensors/transform/factory/base.py,sha256=Menoz_nU64wLEPPP_FdGg_2HhdkN8mKpOCGGOkgoQIY,6532
-compressed_tensors/transform/factory/hadamard.py,sha256=DpDcBwqdccxn6pJ10CSGanTU7mnuFyUBiOX_DvtI1w8,4622
-compressed_tensors/transform/factory/matrix_multiply.py,sha256=poCvuscNHgxBoQLHogy9rNezL3zMm47AxhooT2BOlZ8,4458
+compressed_tensors/transform/factory/hadamard.py,sha256=uLemSQC7xuDb4k0D5INoPi-uZtRNxtN0Dj-i8ladHfE,4620
+compressed_tensors/transform/factory/matrix_multiply.py,sha256=u-7V04EvEe9G3VEF--YwoVV-h5kmh6hXq8stY_EWmLY,4456
 compressed_tensors/transform/factory/random_hadamard.py,sha256=ck-LF7sl7i9NW4fxLypgHgkw91lc_TpwHO8bXX-0fPU,1577
 compressed_tensors/transform/utils/__init__.py,sha256=fH6rjBYAxuwrTzBTlTjTgCYNyh6TCvCqajCz4Im4YrA,617
 compressed_tensors/transform/utils/hadamard.py,sha256=9JMJPtlYI2HwJ2Y2I6uyVCBhKju8yXuQ9DhK4ZpEpvY,5753
@@ -66,8 +66,8 @@ compressed_tensors/utils/permute.py,sha256=wB8LCvQxmI7xnm67S8pbHzM5CDFgG8t4D8ema
 compressed_tensors/utils/safetensors_load.py,sha256=Vql34aCTDHwmTZXJHzCyBISJo7iA7EQ78LdTlMjdpZo,12023
 compressed_tensors/utils/semi_structured_conversions.py,sha256=XKNffPum54kPASgqKzgKvyeqWPAkair2XEQXjkp7ho8,13489
 compressed_tensors/utils/type.py,sha256=bNwoo_FWlvLuDpYAGGzZJITRg0JA_Ngk9LGPo-kvjeU,2554
-compressed_tensors-0.11.1a20250929.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-compressed_tensors-0.11.1a20250929.dist-info/METADATA,sha256=xn1t6NC9__fTUjuX0ctMoF-vXQVWpRalSaJMrmS8A5w,7027
-compressed_tensors-0.11.1a20250929.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-compressed_tensors-0.11.1a20250929.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
-compressed_tensors-0.11.1a20250929.dist-info/RECORD,,
+compressed_tensors-0.12.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+compressed_tensors-0.12.1.dist-info/METADATA,sha256=rGHOzy2FZFRUNCuiBKb99tnWtlks5mAIBrR76QltFkM,7018
+compressed_tensors-0.12.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+compressed_tensors-0.12.1.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
+compressed_tensors-0.12.1.dist-info/RECORD,,