Make compressors stackable

Special condition for marlin_24 compressor
Update tests

src/llmcompressor/transformers/compression/quantization_format.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
 from compressed_tensors import CompressionFormat
-from compressed_tensors.config import SparsityCompressionConfig
 from compressed_tensors.quantization import QuantizationStrategy, QuantizationType
 from compressed_tensors.quantization.utils import (
     is_model_quantized,
@@ -16,7 +15,7 @@ def infer_quantization_format(
     model,
     quantization_format: Optional[str] = None,
     save_compressed: bool = False,
-    sparsity_config: Optional[SparsityCompressionConfig] = None,
+    sparsity_structure: Optional[str] = None,
 ) -> str:
     """
     Infers a quantization format based on model state and compression args
@@ -36,9 +35,7 @@ def infer_quantization_format(
 
     if save_compressed:
         weight_args, input_args = _get_unique_quant_args(model)
-        is_24_structure = (
-            sparsity_config and sparsity_config.sparsity_structure == "2:4"
-        )
+        is_24_structure = sparsity_structure is not None and sparsity_structure == "2:4"
         is_weight_only = len(input_args) == 0 and len(weight_args) > 0
 
         if is_weight_only:  # w4a16 and w8a16

src/llmcompressor/transformers/compression/sparsity_config.py

@@ -1,7 +1,6 @@
 from typing import Dict, Optional
 
 from compressed_tensors import CompressionFormat, SparsityCompressionConfig
-from compressed_tensors.quantization.utils import is_model_quantized
 from torch import Tensor
 from torch.nn import Module
 
@@ -74,6 +73,7 @@ def from_pretrained(
         model: Module,
         state_dict: Optional[Dict[str, Tensor]] = None,
         compress: bool = False,
+        is_marlin: bool = False,
     ) -> Optional["SparsityCompressionConfig"]:
         """
         Determines compression type and informational parameters for a given model
@@ -82,6 +82,7 @@ def from_pretrained(
         :param state_dict: optional state_dict to replace that in model, used for
         gathering global FSDP model info
         :param compress: whether or not to compress the model on disk
+        :param is_marlin: whether or not marlin compression is being used
         :return: compression config inferred from the model
         """
 
@@ -95,10 +96,11 @@ def from_pretrained(
         sparsity_structure = SparsityConfigMetadata.infer_sparsity_structure(
             model=model
         )
-        if is_model_quantized(model):
-            # compressing a sparse quantized model is not supported yet
+        if is_marlin:
+            # sparse compressor should be dense for marlin
+            # compression
             format = CompressionFormat.dense.value
-        elif compress:
+        if compress:
             format = CompressionFormat.sparse_bitmask.value
         else:
             format = CompressionFormat.dense.value

src/llmcompressor/transformers/sparsification/compressed_tensors_utils.py

@@ -6,7 +6,11 @@
 import torch
 import transformers
 from accelerate.accelerator import get_state_dict_offloaded_model
-from compressed_tensors import ModelCompressor, SparsityCompressionConfig
+from compressed_tensors import (
+    CompressionFormat,
+    ModelCompressor,
+    SparsityCompressionConfig,
+)
 from loguru import logger
 from transformers import PreTrainedModel
 
@@ -77,15 +81,22 @@ def save_pretrained_wrapper(
             if state_dict is None:
                 state_dict = get_state_dict_offloaded_model(model)
 
+            sparsity_stucture = SparsityConfigMetadata.infer_sparsity_structure(model)
+            quantization_format = infer_quantization_format(
+                model=model,
+                quantization_format=quantization_format,
+                save_compressed=save_compressed,
+                sparsity_structure=sparsity_stucture,
+            )
+            is_marlin = quantization_format == CompressionFormat.marlin_24.value
+
             if sparsity_config is not None:
                 sparsity_config.global_sparsity = (
                     SparsityConfigMetadata.infer_global_sparsity(
                         model, state_dict=state_dict
                     )
                 )
-                sparsity_config.sparsity_structure = (
-                    SparsityConfigMetadata.infer_sparsity_structure()
-                )
+                sparsity_config.sparsity_structure = sparsity_stucture
             elif not skip_compression_stats:
                 # try to infer a sparsity config from the model if none is provided
                 logger.info(
@@ -95,15 +106,12 @@ def save_pretrained_wrapper(
                     "skip_compression_stats=True"
                 )
                 sparsity_config = SparsityConfigMetadata.from_pretrained(
-                    model, state_dict=state_dict, compress=save_compressed
+                    model,
+                    state_dict=state_dict,
+                    compress=save_compressed,
+                    is_marlin=is_marlin,
                 )
 
-            quantization_format = infer_quantization_format(
-                model=model,
-                quantization_format=quantization_format,
-                save_compressed=save_compressed,
-                sparsity_config=sparsity_config,
-            )
             compressor = ModelCompressor.from_pretrained_model(
                 model,
                 sparsity_config=sparsity_config,

tests/llmcompressor/transformers/compression/test_infer_quant_format.py

@@ -30,6 +30,8 @@ def test_infer_quant_format(preset, sparsity_structure, expected_format):
         module.quantization_scheme = quant_scheme
 
     inferred_format = infer_quantization_format(
-        dummy_model, save_compressed=True, sparsity_config=sparsity_config
+        dummy_model,
+        save_compressed=True,
+        sparsity_structure=sparsity_config.sparsity_structure,
     )
     assert inferred_format.value == expected_format