add tests

Author: dsikka

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -189,6 +189,7 @@ def from_pretrained_model(
             if isinstance(quantization_format, (str, CompressionFormat)):
                 quantization_format = [quantization_format]
 
+            compression_formats = quantization_format
             # assume multiple compression formats means mixed-precision
             # as we currently only support one compressor per precision type and scheme
             if len(quantization_format) > 1:
@@ -216,6 +217,7 @@ def from_pretrained_model(
             sparsity_config=sparsity_config,
             quantization_config=quantization_config,
             transform_config=transform_config,
+            compression_formats=compression_formats,
         )
 
     @staticmethod
@@ -296,10 +298,12 @@ def __init__(
         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[
@@ -313,9 +317,11 @@ def __init__(
             )
 
         if quantization_config is not None:
-            quantization_formats = self._fetch_unique_quantization_formats()
+            if not self.compression_formats:
+                self.compression_formats = self._fetch_unique_quantization_formats()
+
             self.quantization_compressor = {}
-            for format in quantization_formats:
+            for format in self.compression_formats:
                 self.quantization_compressor[
                     format
                 ] = BaseCompressor.load_from_registry(

src/compressed_tensors/quantization/quant_scheme.py

@@ -43,6 +43,7 @@ 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]

tests/test_compressors/model_compressors/test_model_compressor.py

@@ -20,8 +20,12 @@
 import torch
 import torch.nn as nn
 from compressed_tensors.compressors import ModelCompressor
-from compressed_tensors.config import SparsityCompressionConfig
-from compressed_tensors.quantization import QuantizationConfig
+from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig
+from compressed_tensors.quantization import (
+    QuantizationArgs,
+    QuantizationConfig,
+    QuantizationScheme,
+)
 from safetensors.torch import save_file
 from tests.testing_utils import induce_sparsity, requires_hf_quantizer
 from transformers import AutoModelForCausalLM
@@ -469,6 +473,43 @@ def test_compress_model_meta(model_stub, q_format, s_config):
         assert compressed[key].dtype == dtype, f"{key} has incorrect dtype"
 
 
+def test_multiple_quant_compressors():
+    model = torch.nn.Sequential(torch.nn.Linear(1, 2), torch.nn.Linear(2, 3))
+    input_activations = QuantizationArgs(num_bits=8, type="float")
+    weights = QuantizationArgs(num_bits=8, type="float")
+
+    scheme_fp8 = QuantizationScheme(
+        targets=["Linear"],
+        weights=weights,
+        input_activations=input_activations,
+        format=CompressionFormat.float_quantized.value,
+    )
+
+    input_activations = QuantizationArgs(num_bits=4, type="float")
+    weights = QuantizationArgs(num_bits=4, type="float")
+
+    scheme_nvfp4 = QuantizationScheme(
+        targets=["Linear"],
+        weights=weights,
+        input_activations=input_activations,
+        format=CompressionFormat.nvfp4_pack_quantized.value,
+    )
+
+    model[0].quantization_scheme = scheme_fp8
+    model[0].quantization_status = "frozen"
+    model[1].quantization_scheme = scheme_nvfp4
+    model[1].quantization_status = "frozen"
+
+    formats = [scheme_fp8.format, scheme_nvfp4.format]
+
+    compressor = ModelCompressor.from_pretrained_model(model, None, formats)
+    assert isinstance(compressor.quantization_compressor, dict)
+    assert (
+        compressor.quantization_config.format == CompressionFormat.mixed_precision.value
+    )
+    assert all(format in compressor.quantization_compressor for format in formats)
+
+
 @pytest.mark.parametrize(
     "model_stub,comp_stub",
     [