[Model Compressor] Move infer call to from_pretrained_model method (#470)

* move infer call to model

* global format support

* fix

* more clena-up

* fix type hint

* update

* docstring

* sqap

* update

* update

* update

Author: dsikka

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -33,6 +33,9 @@
 from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.compressors.sparse_compressors import DenseCompressor
 from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig
+from compressed_tensors.config.format import (
+    infer_and_set_per_module_quantization_format,
+)
 from compressed_tensors.quantization import (
     DEFAULT_QUANTIZATION_METHOD,
     QuantizationConfig,
@@ -58,6 +61,7 @@
     is_compressed_tensors_config,
 )
 from compressed_tensors.utils.match import match_named_modules
+from loguru import logger
 from torch import Tensor
 from torch.nn import Module
 from tqdm import tqdm
@@ -166,29 +170,50 @@ def from_compression_config(
     def from_pretrained_model(
         cls,
         model: Module,
+        sparsity_config_or_format: Union[SparsityCompressionConfig, str, None] = None,
+        quantization_format: Optional[str] = None,
         sparsity_config: Union[SparsityCompressionConfig, str, None] = None,
-        quantization_format: Optional[Union[str, List[str]]] = None,
     ) -> Optional["ModelCompressor"]:
         """
         Given a pytorch model and optional sparsity and/or quantization configs,
         load the appropriate compressors
 
         :param model: pytorch model to target for compression
         :param sparsity_config: a filled in sparsity config or string corresponding
-            to a sparsity compression algorithm
-        :param quantization_format: string corresponding to a quantization compression
-            algorithm
+            to a sparsity format
+        :param quantization_format: string corresponding to a quantization
+            format that should be applied to the entire model
         :return: compressor for the configs, or None if model is not compressed
         """
-        quantization_config = QuantizationConfig.from_pretrained(
-            model, format=quantization_format
-        )
+        if sparsity_config:
+            logger.warning(
+                "sparsity_config is deprecated, use sparsity_config_or_format"
+            )
+            sparsity_config_or_format = sparsity_config
 
-        # use config passed as argument
-        if isinstance(sparsity_config, str):  # we passed in a sparsity format
+        if sparsity_config_or_format and isinstance(
+            sparsity_config_or_format, str
+        ):  # we passed in a sparsity format
             sparsity_config = SparsityCompressionConfig.load_from_registry(
-                sparsity_config
+                sparsity_config_or_format
             )
+        else:
+            # otherwise, config or None
+            sparsity_config = sparsity_config_or_format
+
+        quantization_format = infer_and_set_per_module_quantization_format(
+            model=model,
+            sparsity_structure=(
+                sparsity_config.sparsity_structure
+                if sparsity_config is not None
+                else None
+            ),
+            quantization_format=quantization_format,
+        )
+
+        quantization_config = QuantizationConfig.from_pretrained(
+            model, format=quantization_format
+        )
 
         # use config attached to model
         transform_config = getattr(model, TRANSFORM_CONFIG_NAME, None)
@@ -200,11 +225,7 @@ def from_pretrained_model(
             sparsity_config=sparsity_config,
             quantization_config=quantization_config,
             transform_config=transform_config,
-            compression_formats=(
-                [quantization_format]
-                if isinstance(quantization_format, str)
-                else quantization_format
-            ),
+            compression_formats=quantization_format,
         )
 
     @staticmethod
@@ -620,6 +641,7 @@ def decompress(self, model_path: str, model: Module):
                     # compressor in a follow-up including initialization
                     load_weight_qparams=load_weight_qparams,
                 )
+
             model_path_or_state_dict = (
                 model.state_dict() if sparse_decompressed else model_path
             )

src/compressed_tensors/config/format.py

@@ -82,14 +82,18 @@ def _get_quant_compression_format(
 
 
 def set_per_module_format(
-    module: torch.nn.Module, sparsity_structure: Optional[str] = None
+    module: torch.nn.Module,
+    sparsity_structure: Optional[str] = None,
+    quantization_format: Optional[str] = None,
 ):
     """
     Determine and set the per module quantization format given quantization args
     and sparsity structure.
 
     :param module: module which has its quantization inferred
     :param sparsity_structure: optional sparsity applied to the module
+    :param quantization_format: optional global format to override
+        the per module formats
 
     """
     weight_scheme = module.quantization_scheme.weights
@@ -100,41 +104,56 @@ def set_per_module_format(
         input_scheme, weight_scheme, sparsity_structure
     )
 
-    # If set, we check if it matches our inferred one
-    if module.quantization_scheme.format is not None:
+    # Check if a global format was provided first
+    # This will override any per module format
+    if quantization_format is not None:
+        if quantization_format != compression_format.value:
+            logger.warning(
+                "The provided format for the module does not match the "
+                "inferred format. Compression may fail "
+            )
+        module.quantization_scheme.format = quantization_format
+    # If a per module format is not provided, we check if it matches our inferred one
+    elif module.quantization_scheme.format is not None:
         # If it does not, warn the user
         if module.quantization_scheme.format != compression_format.value:
             logger.warning(
                 "The provided format for the module does not match the "
                 "inferred format. Compression may fail "
             )
+    # If neither provided, set ours
     else:
-        # If not set, we set ours
         module.quantization_scheme.format = compression_format.value
 
 
 def infer_and_set_per_module_quantization_format(
     model: torch.nn.Module,
     sparsity_structure: Optional[str] = None,
+    quantization_format: Optional[str] = None,
 ) -> List[str]:
     """
     Infers the quantization format for a model based on its state and provided
     compression arguments. Updates thhe quantization_scheme.format value
-    based on the inferred format. Returns the unique list of formats in the model
-    or None if empty list
+    based on the inferred format. Returns the unique list of formats in the model.
+    All None formats are mapped to CompressionFormat.dense.value
 
     For a summary of the formats, see `docs/guides/compression_formats.md`.
 
     :param model: model to check for quantization
     :param sparsity_structure: optional sparsity applied to the module
-    :return compression format appropriate for model
+    :param quantization_format: optional global format to override
+        the per module formats
+    :return compression format appropriate for the model
     """
     unique_formats = []
     for submodule in model.modules():
         if is_module_quantized(submodule):
             assert hasattr(submodule, "quantization_scheme")
-            set_per_module_format(submodule, sparsity_structure)
-            if submodule.quantization_scheme.format not in unique_formats:
+            set_per_module_format(submodule, sparsity_structure, quantization_format)
+            if (
+                submodule.quantization_scheme.format
+                and submodule.quantization_scheme.format not in unique_formats
+            ):
                 unique_formats.append(submodule.quantization_scheme.format)
 
     if len(unique_formats) > 0:

src/compressed_tensors/quantization/quant_config.py

@@ -165,7 +165,7 @@ def to_dict(self):
 
     @staticmethod
     def from_pretrained(
-        model: Module, format: Optional[str] = None
+        model: Module, format: Optional[Union[str, list]] = None
     ) -> Optional["QuantizationConfig"]:
         """
         Converts a model into its associated QuantizationConfig based on the

tests/test_compressors/model_compressors/test_model_compressor.py

@@ -342,7 +342,7 @@ def _get_combined_config(s_config, q_config):
 )
 def test_compress_model(model_stub, q_format, s_config, tmpdir):
     model = AutoModelForCausalLM.from_pretrained(model_stub, torch_dtype=torch.float32)
-    compressor = ModelCompressor.from_pretrained_model(model, s_config, [q_format])
+    compressor = ModelCompressor.from_pretrained_model(model, s_config, q_format)
 
     # compress model by eagerly compressing state dict
     true_compressed = dict(compressor.compress(model))
@@ -388,7 +388,7 @@ def test_compress_model_meta(model_stub, q_format, s_config):
     # Load model on CPU to get expected compressed state_dict
     cpu_model = AutoModelForCausalLM.from_pretrained(model_stub)
     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()}
@@ -403,7 +403,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
@@ -442,7 +442,7 @@ def test_multiple_quant_compressors():
 
     formats = [scheme_fp8.format, scheme_nvfp4.format]
 
-    compressor = ModelCompressor.from_pretrained_model(model, None, formats)
+    compressor = ModelCompressor.from_pretrained_model(model, None)
     assert isinstance(compressor.quantization_compressor, dict)
     assert (
         compressor.quantization_config.format == CompressionFormat.mixed_precision.value