Enable Sparse compression

src/llmcompressor/transformers/compression/quantization_format.py

@@ -1,7 +1,7 @@
 from typing import Optional
 
 from compressed_tensors import CompressionFormat
-from compressed_tensors.config import SparsityCompressionConfig
+from compressed_tensors.config import SparsityStructure
 from compressed_tensors.quantization import QuantizationStrategy, QuantizationType
 from compressed_tensors.quantization.utils import (
     is_model_quantized,
@@ -16,7 +16,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
@@ -37,7 +37,8 @@ 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"
+            SparsityStructure(sparsity_structure).value
+            == SparsityStructure.TWO_FOUR.value
         )
         is_weight_only = len(input_args) == 0 and len(weight_args) > 0
 

src/llmcompressor/transformers/compression/sparsity_config.py

@@ -1,7 +1,7 @@
 from typing import Dict, Optional
 
 from compressed_tensors import CompressionFormat, SparsityCompressionConfig
-from compressed_tensors.quantization.utils import is_model_quantized
+from compressed_tensors.config import SparsityStructure
 from torch import Tensor
 from torch.nn import Module
 
@@ -20,7 +20,7 @@ class SparsityConfigMetadata:
     metadata from the model
     """
 
-    SPARSITY_THRESHOLD: float = 0.4
+    SPARSITY_THRESHOLD: float = 0.5
 
     @staticmethod
     def infer_global_sparsity(
@@ -67,13 +67,14 @@ def infer_sparsity_structure(model: Optional[Module] = None) -> str:
         if model and sparsity_structure is None:
             sparsity_structure = infer_sparsity_structure_from_model(model)
 
-        return sparsity_structure or "unstructured"
+        return SparsityStructure(sparsity_structure).value
 
     @staticmethod
     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 +83,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,11 +97,15 @@ 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:
-            format = CompressionFormat.sparse_bitmask.value
+        if compress:
+            if sparsity_structure == SparsityStructure.TWO_FOUR.value:
+                format = CompressionFormat.sparse_24.value
+            else:
+                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 safetensors.torch import storage_ptr
 from transformers import PreTrainedModel
@@ -78,15 +82,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(
@@ -96,15 +107,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