handle mixed-precision case

Author: dsikka

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -164,7 +164,7 @@ def from_pretrained_model(
         cls,
         model: Module,
         sparsity_config: Union[SparsityCompressionConfig, str, None] = None,
-        quantization_format: Optional[List[str]] = None,
+        quantization_format: Optional[Union[str, List[str]]] = None,
     ) -> Optional["ModelCompressor"]:
         """
         Given a pytorch model and optional sparsity and/or quantization configs,
@@ -179,10 +179,14 @@ def from_pretrained_model(
         """
         # assume multiple compression formats means mixed-precision
         # as we currently only support one compressor per precision type and scheme
-        if len(quantization_format) > 1:
-            quantization_format = CompressionFormat.mixed_precision.value
-        else:
-            quantization_format = quantization_format[0]
+        if quantization_format is not None:
+            if isinstance(quantization_format, str):
+                quantization_format = [quantization_format]
+
+            if len(quantization_format) > 1:
+                quantization_format = CompressionFormat.mixed_precision.value
+            else:
+                quantization_format = quantization_format[0]
 
         quantization_config = QuantizationConfig.from_pretrained(
             model, format=quantization_format
@@ -393,12 +397,13 @@ def get_unexpected_file_keys(self, model: Module) -> List[str]:
                     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)
 
@@ -436,9 +441,24 @@ def compress_model(self, model: Module):
 
                 # quantization first
                 if prefix in module_to_scheme:
-                    quant_compressor = self.quantization_compressor.get(
-                        module.quantization_scheme.format
-                    )
+                    if (
+                        not hasattr(module.quantization_scheme, "format")
+                        or module.quantization_scheme.format is None
+                    ):
+                        if (
+                            self.quantization_config.format
+                            == CompressionFormat.mixed_precision.value
+                        ):
+                            raise ValueError(
+                                "Compressing mixed-precision models without defining "
+                                "per module quantization_scheme.format is currently "
+                                "not supported"
+                            )
+                        format = self.quantization_config.format
+                    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,
@@ -510,9 +530,24 @@ def decompress_model(self, model: Module):
 
                 # quantization second
                 if prefix in module_to_scheme:
-                    quant_compressor = self.quantization_compressor.get(
-                        module.quantization_scheme.format
-                    )
+
+                    if (
+                        not hasattr(module.quantization_scheme, "format")
+                        or module.quantization_scheme.format is None
+                    ):
+                        if (
+                            self.quantization_config.format
+                            == CompressionFormat.mixed_precision.value
+                        ):
+                            raise ValueError(
+                                "Decompressing mixed-precision models without defining "
+                                "per module quantization_scheme.format is currently not "
+                                "supported"
+                            )
+                        format = self.quantization_config.format
+                    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,
@@ -606,15 +641,19 @@ def decompress(self, model_path: str, model: Module):
         """
         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
-            quant_compressor = next(iter(self.quantization_compressor.values()))
             params_to_ignore = None
-            if self.quantization_compressor is not None:
+            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

tests/test_compressors/model_compressors/test_model_compressor.py

@@ -443,7 +443,7 @@ def test_compress_model_meta(model_stub, q_format, s_config):
         model_stub, torch_dtype=torch.float32
     )
     reference_compressor = ModelCompressor.from_pretrained_model(
-        cpu_model, s_config, q_format
+        cpu_model, s_config, [q_format]
     )
     # Only stores dtype because meta model does not store values
     expected = {k: v.dtype for k, v in reference_compressor.compress(cpu_model).items()}
@@ -459,7 +459,7 @@ def test_compress_model_meta(model_stub, q_format, s_config):
             module.to_empty(device="meta")
 
     # Compress in-place on meta model
-    compressor = ModelCompressor.from_pretrained_model(meta_model, s_config, q_format)
+    compressor = ModelCompressor.from_pretrained_model(meta_model, s_config, [q_format])
     compressor.compress_model(meta_model)
 
     # Compare keys and dtypes