[Quantization] Support mixed-precision compression

src/llmcompressor/transformers/compression/quantization_format.py

@@ -5,7 +5,7 @@
 from compressed_tensors.quantization import QuantizationStrategy, QuantizationType
 from compressed_tensors.quantization.utils import is_module_quantized
 
-__all__ = ["infer_quantization_format"]
+__all__ = ["infer_quantization_format", "infer_per_module_quantization"]
 
 
 def infer_quantization_format(
@@ -97,6 +97,60 @@ def infer_quantization_format(
         return None
 
 
+def _get_quant_method(input_args, weight_args, sparsity_structure):
+    is_24_structure = (
+        SparsityStructure(sparsity_structure) == SparsityStructure.TWO_FOUR
+    )
+    is_weight_only = weight_args is not None and input_args is None
+
+    if weight_args.num_bits == 4 and weight_args.type == QuantizationType.FLOAT.value:
+        return CompressionFormat.nvfp4_pack_quantized
+
+    if is_weight_only:  # w4a16 and w8a16
+        is_valid_pack = (
+            weight_args.num_bits in [4, 8]
+            and weight_args.type == QuantizationType.INT.value
+        )
+        if not is_valid_pack:  # packing only valid for int4 and int 8
+            return CompressionFormat.naive_quantized
+        if is_24_structure:
+            if (
+                weight_args.strategy is not QuantizationStrategy.CHANNEL.value
+                and weight_args.strategy is not QuantizationStrategy.GROUP.value
+            ):
+                # marlin24 kernel only applicable for channel/group quantization
+                return CompressionFormat.pack_quantized
+            return CompressionFormat.marlin_24
+        return CompressionFormat.pack_quantized
+
+    else:  # w8a8 float and int
+        if (
+            weight_args.type == QuantizationType.FLOAT.value
+            and weight_args.num_bits == 8
+        ):
+            return CompressionFormat.float_quantized
+        if weight_args.type == QuantizationType.INT.value:
+            return CompressionFormat.int_quantized
+
+        return CompressionFormat.naive_quantized
+
+
+def infer_per_module_quantization(model, sparsity_structure):
+    unique_formats = []
+    for submodule in model.modules():
+        if is_module_quantized(submodule):
+            weight_scheme = submodule.quantization_scheme.weights
+            input_scheme = submodule.quantization_scheme.input_activations
+            compression_format = _get_quant_method(
+                input_scheme, weight_scheme, sparsity_structure
+            )
+            print(compression_format)
+            submodule.quantization_scheme.format = compression_format.value
+            if compression_format not in unique_formats:
+                unique_formats.append(compression_format)
+    return unique_formats
+
+
 def _get_unique_quant_args(model):
     """
     Gets a list of all the unique quantization settings present in model

src/llmcompressor/transformers/sparsification/compressed_tensors_utils.py

@@ -19,6 +19,7 @@
 from llmcompressor.core import active_session
 from llmcompressor.pytorch.model_load.helpers import copy_python_files_from_model_cache
 from llmcompressor.transformers.compression.quantization_format import (
+    infer_per_module_quantization,
     infer_quantization_format,
 )
 from llmcompressor.transformers.compression.sparsity_metadata_config import (
@@ -228,6 +229,7 @@ def get_model_compressor(
                 SparsityConfigMetadata.infer_sparsity_structure(model)
             )
 
+    """
     quantization_format: Optional[CompressionFormat] = infer_quantization_format(
         model=model,
         quantization_format=quantization_format,
@@ -236,6 +238,14 @@ def get_model_compressor(
         if sparsity_config is None
         else sparsity_config.sparsity_structure,
     )
+    """
+
+    quantization_format = infer_per_module_quantization(
+        model,
+        sparsity_structure=None
+        if sparsity_config is None
+        else sparsity_config.sparsity_structure,
+    )
 
     return ModelCompressor.from_pretrained_model(
         model,