support more than one quant compressor

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[str] = None,
+        quantization_format: Optional[List[str]] = None,
     ) -> Optional["ModelCompressor"]:
         """
         Given a pytorch model and optional sparsity and/or quantization configs,
@@ -267,9 +267,18 @@ def __init__(
                 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 isinstance(quantization_config.format, list):
+                self.quantization_compressor = {}
+                for format in quantization_config.format:
+                    self.quantization_compressor[
+                        format
+                    ] = BaseCompressor.load_from_registry(
+                        format, config=quantization_config
+                    )
+            else:
+                self.quantization_compressor = BaseCompressor.load_from_registry(
+                    quantization_config.format, config=quantization_config
+                )
 
     # ----- used by hf quantizer ----- #
 
@@ -407,12 +416,23 @@ def compress_model(self, model: Module):
 
                 # quantization first
                 if prefix in module_to_scheme:
-                    state_dict = self.quantization_compressor.compress(
-                        state_dict,
-                        names_to_scheme=module_to_scheme,
-                        show_progress=False,
-                        compression_device=exec_device,
-                    )
+                    if isinstance(self.quantization_compressor, dict):
+                        quant_compressor = self.quantization_compressor.get(
+                            module.quantization_scheme.format
+                        )
+                        state_dict = quant_compressor.compress(
+                            state_dict,
+                            names_to_scheme=module_to_scheme,
+                            show_progress=False,
+                            compression_device=exec_device,
+                        )
+                    else:
+                        state_dict = self.quantization_compressor.compress(
+                            state_dict,
+                            names_to_scheme=module_to_scheme,
+                            show_progress=False,
+                            compression_device=exec_device,
+                        )
 
                 # sparsity second
                 if prefix in sparse_compression_targets:

src/compressed_tensors/quantization/quant_config.py

@@ -138,7 +138,7 @@ class QuantizationConfig(BaseModel):
     config_groups: Dict[str, Union[QuantizationScheme, List[str]]]
     quant_method: str = DEFAULT_QUANTIZATION_METHOD
     kv_cache_scheme: Optional[QuantizationArgs] = None
-    format: str = DEFAULT_QUANTIZATION_FORMAT
+    format: Union[List[str], str] = DEFAULT_QUANTIZATION_FORMAT
     quantization_status: QuantizationStatus = QuantizationStatus.INITIALIZED
     global_compression_ratio: Optional[float] = None
     ignore: Optional[List[str]] = Field(default_factory=list)
@@ -162,7 +162,7 @@ def to_dict(self):
 
     @staticmethod
     def from_pretrained(
-        model: Module, format: Optional[str] = None
+        model: Module, format: Optional[Union[List[str], str]] = None
     ) -> Optional["QuantizationConfig"]:
         """
         Converts a model into its associated QuantizationConfig based on the
@@ -228,7 +228,7 @@ def from_pretrained(
 
         if format is None:
             if quantization_status == QuantizationStatus.COMPRESSED:
-                format = CompressionFormat.int_quantized.value
+                format = CompressionFormat.int_quantized.value  # why?!
             else:
                 format = CompressionFormat.dense.value
 

src/compressed_tensors/quantization/quant_scheme.py

@@ -16,6 +16,7 @@
 from copy import deepcopy
 from typing import Any, Dict, List, Optional
 
+from compressed_tensors.config import CompressionFormat
 from compressed_tensors.quantization.quant_args import (
     DynamicType,
     QuantizationArgs,
@@ -48,6 +49,7 @@ class QuantizationScheme(BaseModel):
     weights: Optional[QuantizationArgs] = None
     input_activations: Optional[QuantizationArgs] = None
     output_activations: Optional[QuantizationArgs] = None
+    format: Optional[CompressionFormat] = None
 
     @model_validator(mode="after")
     def validate_model_after(model: "QuantizationScheme") -> "QuantizationScheme":

src/compressed_tensors/transform/factory/base.py

@@ -14,7 +14,7 @@
 
 from abc import ABC, abstractmethod
 from collections import defaultdict
-from typing import List, Optional, Tuple, Set
+from typing import List, Optional, Set, Tuple
 
 import torch
 import torch.nn.utils.parametrize as P