compressed-tensors 0.10.3a20250811__py3-none-any.whl → 0.10.3a20250814__py3-none-any.whl

compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -169,7 +169,7 @@ class ModelCompressor:
         cls,
         model: Module,
         sparsity_config: Union[SparsityCompressionConfig, str, None] = None,
-        quantization_format: Optional[str] = None,
+        quantization_format: Optional[Union[str, List[str]]] = None,
     ) -> Optional["ModelCompressor"]:
         """
         Given a pytorch model and optional sparsity and/or quantization configs,
@@ -182,7 +182,6 @@ class ModelCompressor:
             algorithm
         :return: compressor for the configs, or None if model is not compressed
         """
-        # reconstruct config from schemes attached to modules
         quantization_config = QuantizationConfig.from_pretrained(
             model, format=quantization_format
         )
@@ -203,11 +202,14 @@ class ModelCompressor:
             sparsity_config=sparsity_config,
             quantization_config=quantization_config,
             transform_config=transform_config,
+            compression_formats=[quantization_format]
+            if isinstance(quantization_format, str)
+            else quantization_format,
         )
 
     @staticmethod
     def parse_sparsity_config(
-        compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"]
+        compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"],
     ) -> Union[Dict[str, Any], None]:
         """
         Parse sparsity config from quantization/compression config. Sparsity
@@ -227,7 +229,7 @@ class ModelCompressor:
 
     @staticmethod
     def parse_quantization_config(
-        compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"]
+        compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"],
     ) -> Union[Dict[str, Any], None]:
         """
         Parse quantization config from quantization/compression config. The
@@ -246,6 +248,7 @@ class ModelCompressor:
 
         quantization_config = deepcopy(compression_config)
         quantization_config.pop(SPARSITY_CONFIG_NAME, None)
+        quantization_config.pop(TRANSFORM_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
@@ -262,19 +265,39 @@ class ModelCompressor:
 
         return quantization_config
 
+    def _fetch_unique_quantization_formats(self) -> List[str]:
+        """
+        Get all unique compression formats present in a model.
+        :return: list of quantization formats
+        """
+        quantization_formats = []
+        for _, scheme in self.quantization_config.config_groups.items():
+            if scheme.format is not None and scheme.format not in quantization_formats:
+                quantization_formats.append(scheme.format)
+
+        if (
+            len(quantization_formats) == 0
+            and self.quantization_config.format
+            != CompressionFormat.mixed_precision.value
+        ):
+            quantization_formats.append(self.quantization_config.format)
+        return quantization_formats
+
     def __init__(
         self,
         sparsity_config: Optional[SparsityCompressionConfig] = None,
         quantization_config: Optional[QuantizationConfig] = None,
         transform_config: Optional[TransformConfig] = None,
+        compression_formats: Optional[List[str]] = None,
     ):
         self.sparsity_config = sparsity_config
         self.quantization_config = quantization_config
         self.transform_config = transform_config
+        self.compression_formats = compression_formats
 
         self.sparsity_compressor = None
         self.quantization_compressor: Optional[
-            Union[BaseQuantizationCompressor, DenseCompressor]
+            Dict[str, Union[BaseQuantizationCompressor, DenseCompressor]]
         ] = None
         # no transform compressor is required
 
@@ -282,10 +305,21 @@ class ModelCompressor:
             self.sparsity_compressor = BaseCompressor.load_from_registry(
                 sparsity_config.format, config=sparsity_config
             )
+
         if quantization_config is not None:
-            self.quantization_compressor = BaseCompressor.load_from_registry(
-                quantization_config.format, config=quantization_config
-            )
+            # If a list of compression_format is not provided, we resolve the
+            # relevant quantization formats using the config groups from the config
+            # and if those are not defined, we fall-back to the global quantization format
+            if not self.compression_formats:
+                self.compression_formats = self._fetch_unique_quantization_formats()
+
+            self.quantization_compressor = {}
+            for format in self.compression_formats:
+                self.quantization_compressor[
+                    format
+                ] = BaseCompressor.load_from_registry(
+                    format, config=quantization_config
+                )
 
     # ----- used by hf quantizer ----- #
 
@@ -380,12 +414,13 @@ class ModelCompressor:
                     targets=scheme.targets,
                     ignore=self.quantization_config.ignore,
                 )
-                unexpected_keys.update(
-                    merge_names(target, param)
-                    for target in quant_targets
-                    for param in self.quantization_compressor.compression_param_names
-                    if param != "weight"
-                )
+                for quant_compressor in self.quantization_compressor.values():
+                    unexpected_keys.update(
+                        merge_names(target, param)
+                        for target in quant_targets
+                        for param in quant_compressor.compression_param_names
+                        if param != "weight"
+                    )
 
         return list(unexpected_keys)
 
@@ -423,7 +458,21 @@ class ModelCompressor:
 
                 # quantization first
                 if prefix in module_to_scheme:
-                    state_dict = self.quantization_compressor.compress(
+                    if (
+                        not hasattr(module.quantization_scheme, "format")
+                        or module.quantization_scheme.format is None
+                    ):
+                        if len(self.compression_formats) > 1:
+                            raise ValueError(
+                                "Applying multiple compressors without defining "
+                                "per module formats is not supported "
+                            )
+                        format = self.compression_formats[0]
+                    else:
+                        format = module.quantization_scheme.format
+
+                    quant_compressor = self.quantization_compressor.get(format)
+                    state_dict = quant_compressor.compress(
                         state_dict,
                         names_to_scheme=module_to_scheme,
                         show_progress=False,
@@ -494,12 +543,24 @@ class ModelCompressor:
 
                 # quantization second
                 if prefix in module_to_scheme:
-                    state_dict = (
-                        self.quantization_compressor.decompress_module_from_state_dict(
-                            prefix,
-                            state_dict,
-                            scheme=module_to_scheme[prefix],
-                        )
+
+                    if (
+                        not hasattr(module.quantization_scheme, "format")
+                        or module.quantization_scheme.format is None
+                    ):
+                        if len(self.compression_formats) > 1:
+                            raise ValueError(
+                                "Applying multiple compressors without defining "
+                                "per module formats is not supported "
+                            )
+                        format = self.compression_formats[0]
+                    else:
+                        format = module.quantization_scheme.format
+                    quant_compressor = self.quantization_compressor.get(format)
+                    state_dict = quant_compressor.decompress_module_from_state_dict(
+                        prefix,
+                        state_dict,
+                        scheme=module_to_scheme[prefix],
                     )
 
                 # remove any existing parameters
@@ -538,7 +599,9 @@ class ModelCompressor:
 
         if self.quantization_compressor is not None:
             module_to_scheme = map_module_to_scheme(model)
-            state_dict = self.quantization_compressor.compress(
+            # Note - compress only supports one compression format atm
+            quant_compressor = next(iter(self.quantization_compressor.values()))
+            state_dict = quant_compressor.compress(
                 state_dict,
                 names_to_scheme=module_to_scheme,
                 show_progress=show_progress,
@@ -587,14 +650,20 @@ class ModelCompressor:
         """
         model_path = get_safetensors_folder(model_path)
         sparse_decompressed = False
+        quant_compressor = (
+            next(iter(self.quantization_compressor.values()))
+            if self.quantization_compressor is not None
+            else None
+        )
 
         if (
             self.sparsity_compressor is not None
             and self.sparsity_config.format != CompressionFormat.dense.value
         ):
+            # note - decompress only supports one compressor atm
             params_to_ignore = None
-            if self.quantization_compressor is not None:
-                params_to_ignore = self.quantization_compressor.compression_param_names
+            if quant_compressor is not None:
+                params_to_ignore = quant_compressor.compression_param_names
             # Sparse decompression is applied on the model_path
             # The compressor will try and load any quantization parameters as well
             # params_to_skip_load will skip over quantization params from being loaded
@@ -605,7 +674,7 @@ class ModelCompressor:
             setattr(model, SPARSITY_CONFIG_NAME, self.sparsity_compressor.config)
             sparse_decompressed = True
 
-        if self.quantization_compressor is not None:
+        if quant_compressor is not None:
             # Temporarily set quantization status to FROZEN to prevent
             # quantization during apply_quantization_config. This ensures
             # that the dtypes of the weights are not unintentionally updated.
@@ -628,7 +697,7 @@ class ModelCompressor:
                     # including initialization
                     load_weight_quantization=(
                         sparse_decompressed
-                        or isinstance(self.quantization_compressor, DenseCompressor)
+                        or isinstance(quant_compressor, DenseCompressor)
                     ),
                 )
 
@@ -636,7 +705,7 @@ class ModelCompressor:
                 model.state_dict() if sparse_decompressed else model_path
             )
 
-            dense_gen = self.quantization_compressor.decompress(
+            dense_gen = quant_compressor.decompress(
                 model_path_or_state_dict, names_to_scheme=names_to_scheme
             )
             # TODO: all weight quantization params will be moved to the compressor
@@ -674,7 +743,7 @@ class ModelCompressor:
 
         # serialize configs into json
         qconfig_data = (
-            self.quantization_config.model_dump(exclude=["quant_method", "format"])
+            self.quantization_config.model_dump(exclude=["quant_method"])
             if self.quantization_config is not None
             else {}
         )

compressed_tensors/config/base.py

@@ -32,6 +32,7 @@ class CompressionFormat(Enum):
     naive_quantized = "naive-quantized"
     pack_quantized = "pack-quantized"
     marlin_24 = "marlin-24"
+    mixed_precision = "mixed-precision"
     nvfp4_pack_quantized = "nvfp4-pack-quantized"
 
 

compressed_tensors/quantization/quant_config.py

@@ -234,6 +234,12 @@ class QuantizationConfig(BaseModel):
                 format = CompressionFormat.int_quantized.value
             else:
                 format = CompressionFormat.dense.value
+        elif isinstance(format, list):
+            format = (
+                CompressionFormat.mixed_precision.value
+                if len(format) > 1
+                else format[0]
+            )
 
         return QuantizationConfig(
             config_groups=config_groups,

compressed_tensors/quantization/quant_scheme.py

@@ -16,6 +16,7 @@ import warnings
 from copy import deepcopy
 from typing import List, Optional
 
+from compressed_tensors.config import CompressionFormat
 from compressed_tensors.quantization.quant_args import (
     DynamicType,
     QuantizationArgs,
@@ -42,12 +43,14 @@ class QuantizationScheme(BaseModel):
     :param weights: quantization config for layer weights
     :param input_activations: quantization config for layer inputs
     :param output_activations: quantization config for layer outputs
+    :param format: CompressionFormat for the layer
     """
 
     targets: List[str]
     weights: Optional[QuantizationArgs] = None
     input_activations: Optional[QuantizationArgs] = None
     output_activations: Optional[QuantizationArgs] = None
+    format: Optional[str] = None
 
     @model_validator(mode="after")
     def validate_model_after(model: "QuantizationScheme") -> "QuantizationScheme":

compressed_tensors/utils/offload.py

@@ -86,6 +86,7 @@ __all__ = [
     "offloaded_dispatch",
     "disable_offloading",
     "remove_dispatch",
+    "cast_to_device",
 ]
 
 
@@ -169,6 +170,19 @@ def update_parameter_data(
 """ Candidates for Upstreaming """
 
 
+def cast_to_device(device_spec: Union[int, torch.device]) -> torch.device:
+    """
+    Convert an integer device index or torch.device into a torch.device object.
+
+    :param device_spec: Device index (int) or torch.device object.
+                        Negative integers map to CPU.
+    :return: torch.device corresponding to the given device specification.
+    """
+    if isinstance(device_spec, int):
+        return torch.device(f"cuda:{device_spec}" if device_spec >= 0 else "cpu")
+    return device_spec
+
+
 def get_execution_device(module: torch.nn.Module) -> torch.device:
     """
     Get the device which inputs should be moved to before module execution.
@@ -179,7 +193,7 @@ def get_execution_device(module: torch.nn.Module) -> torch.device:
     """
     for submodule in module.modules():
         if has_offloaded_params(submodule):
-            return submodule._hf_hook.execution_device
+            return cast_to_device(submodule._hf_hook.execution_device)
 
         param = next(submodule.parameters(recurse=False), None)
         if param is not None:

compressed_tensors/version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.10.3.a20250811'
+__version__ = version = '0.10.3.a20250814'
 __version_tuple__ = version_tuple = (0, 10, 3)

{compressed_tensors-0.10.3a20250811.dist-info → compressed_tensors-0.10.3a20250814.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.10.3a20250811
+Version: 0.10.3a20250814
 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.10.3a20250811.dist-info → compressed_tensors-0.10.3a20250814.dist-info}/RECORD

@@ -1,11 +1,11 @@
 compressed_tensors/__init__.py,sha256=UtKmifNeBCSE2TZSAfduVNNzHY-3V7bLjZ7n7RuXLOE,812
 compressed_tensors/base.py,sha256=-gxWvDF4LCkyeDP8YlGzvBBKxo4Dk9h4NINPD61drFU,921
-compressed_tensors/version.py,sha256=9NgEdMzgL7r039RghUuU-BhjHVuGK1utk2z8Au9OlWA,523
+compressed_tensors/version.py,sha256=fAUC53w9XJ-gbZ3V6UhPrss4y7OIGxsBJ0rFa1T--zA,523
 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=0WULLKpgWuTQLjKsCstiTssT778wp9TWMkQjHbYO4Zo,33989
+compressed_tensors/compressors/model_compressors/model_compressor.py,sha256=FuPS3LYSJk0ATu6caW_GQsFi31EqFTnQtR6mIe6fDAU,37278
 compressed_tensors/compressors/quantized_compressors/__init__.py,sha256=KvaFBL_Q84LxRGJOV035M8OBoCkAx8kOkfphswgkKWk,745
 compressed_tensors/compressors/quantized_compressors/base.py,sha256=YGUMzbxekj_36ChgQnVZN6T8uDjXtGG1zfMIBGBLWco,10354
 compressed_tensors/compressors/quantized_compressors/naive_quantized.py,sha256=0ANDcuD8aXPqTYNPY6GnX9iS6eXJw6P0TzNV_rYS2l8,5369
@@ -19,7 +19,7 @@ compressed_tensors/compressors/sparse_compressors/sparse_bitmask.py,sha256=S8vW0
 compressed_tensors/compressors/sparse_quantized_compressors/__init__.py,sha256=4f_cwcKXB1nVVMoiKgTFAc8jAPjPLElo-Df_EDm1_xw,675
 compressed_tensors/compressors/sparse_quantized_compressors/marlin_24.py,sha256=7F9J6wgkecitK5hHuqjetZ18HExHIF4QIw1wgm2Y6U8,10099
 compressed_tensors/config/__init__.py,sha256=8sOoZ6xvYSC79mBvEtO8l6xk4PC80d29AnnJiGMrY2M,737
-compressed_tensors/config/base.py,sha256=p3glQHvC2fjodf_SvlelVrTWSIjGXgGC86t8oVOlMng,3529
+compressed_tensors/config/base.py,sha256=FaImUwb5G93en2BHUKDs76L_tO8NFpdxlfwAgQL7mNM,3569
 compressed_tensors/config/dense.py,sha256=NgSxnFCnckU9-iunxEaqiFwqgdO7YYxlWKR74jNbjks,1317
 compressed_tensors/config/sparse_24_bitmask.py,sha256=Lhj39zT2V1hxftprvxvneyhv45ShlXOKd75DBbDTyTE,1401
 compressed_tensors/config/sparse_bitmask.py,sha256=pZUboRNZTu6NajGOQEFExoPknak5ynVAUeiiYpS1Gt8,1308
@@ -27,8 +27,8 @@ compressed_tensors/linear/__init__.py,sha256=fH6rjBYAxuwrTzBTlTjTgCYNyh6TCvCqajC
 compressed_tensors/linear/compressed_linear.py,sha256=1yo9RyjA0aQ--iuIknFfcSorJn43Mn4CoV-q4JlTJ_o,4052
 compressed_tensors/quantization/__init__.py,sha256=83J5bPB7PavN2TfCoW7_vEDhfYpm4TDrqYO9vdSQ5bk,760
 compressed_tensors/quantization/quant_args.py,sha256=PMoaa6hpyJLGGSeCWefGmzGVxbOtxAdDunHJi_L5gNs,12894
-compressed_tensors/quantization/quant_config.py,sha256=StEpCvc70JasE1srLaHqI-TJlasLWGtHU2o0E_gDJhQ,10400
-compressed_tensors/quantization/quant_scheme.py,sha256=3EUGCw5_e7nnmvYPK_UlQKaaskOLIAo30dHYn0z7HmQ,8521
+compressed_tensors/quantization/quant_config.py,sha256=2NgDwKuQn0f-ojiHC8c6tXtYX_zQlk26Rj-bU71QKvA,10598
+compressed_tensors/quantization/quant_scheme.py,sha256=k25Cdx7BZCvLlRlENu4BVoFxquqcErP58P3Y_1HsKB4,8661
 compressed_tensors/quantization/lifecycle/__init__.py,sha256=_uItzFWusyV74Zco_pHLOTdE9a83cL-R-ZdyQrBkIyw,772
 compressed_tensors/quantization/lifecycle/apply.py,sha256=wM8mVcbKvZjBo18pSXMp28i30YWwUXJPSS7_HCakH9U,17892
 compressed_tensors/quantization/lifecycle/compressed.py,sha256=Fj9n66IN0EWsOAkBHg3O0GlOQpxstqjCcs0ttzMXrJ0,2296
@@ -57,14 +57,14 @@ compressed_tensors/utils/__init__.py,sha256=spzbjUO4-hZ2jXGST27r3MIt2yzIXsjdbEaY
 compressed_tensors/utils/helpers.py,sha256=Q3iRAa2XSdmmn4vSpUplnvKOmWwn4Clao9ZkPBHXtpI,12604
 compressed_tensors/utils/internal.py,sha256=7SSWgDoNFRnlfadwkoFhLW-T2jOc7Po_WzWv5h32Sa8,982
 compressed_tensors/utils/match.py,sha256=9x-yZIlq7ndSLf2aQwNT7IpBQDe-8H6utiJkji8wPrQ,9397
-compressed_tensors/utils/offload.py,sha256=3XiBuWbUkBAt8v1t5i57qDcbB3VJQs_FDeayi-JzIWg,23896
+compressed_tensors/utils/offload.py,sha256=gFoEDaissHsLM5-JDbgPxh5hiE9VFN4HFxvszYvReos,24446
 compressed_tensors/utils/permutations_24.py,sha256=kx6fsfDHebx94zsSzhXGyCyuC9sVyah6BUUir_StT28,2530
 compressed_tensors/utils/permute.py,sha256=V6tJLKo3Syccj-viv4F7ZKZgJeCB-hl-dK8RKI_kBwI,2355
 compressed_tensors/utils/safetensors_load.py,sha256=DMfZBuUbA6qp_BG_zIWT3ckiEE33K9ob34s-OgzReO4,12057
 compressed_tensors/utils/semi_structured_conversions.py,sha256=XKNffPum54kPASgqKzgKvyeqWPAkair2XEQXjkp7ho8,13489
 compressed_tensors/utils/type.py,sha256=bNwoo_FWlvLuDpYAGGzZJITRg0JA_Ngk9LGPo-kvjeU,2554
-compressed_tensors-0.10.3a20250811.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-compressed_tensors-0.10.3a20250811.dist-info/METADATA,sha256=mYMXLEK9r53lXrMbZBRmkimI3aW-X1x4n-8DUThb0K8,7031
-compressed_tensors-0.10.3a20250811.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-compressed_tensors-0.10.3a20250811.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
-compressed_tensors-0.10.3a20250811.dist-info/RECORD,,
+compressed_tensors-0.10.3a20250814.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+compressed_tensors-0.10.3a20250814.dist-info/METADATA,sha256=-7voWXyJPB13WkMJADa57hDE4euKxrCjnQfYdHROjKg,7031
+compressed_tensors-0.10.3a20250814.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+compressed_tensors-0.10.3a20250814.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
+compressed_tensors-0.10.3a20250814.dist-info/RECORD,,