Enable sparse24bytemask compressor

vllm-project · Dec 12, 2024 · 2a5a71a · 2a5a71a
1 parent 938a526
commit 2a5a71a
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/src/llmcompressor/transformers/compression/sparsity_config.py b/src/llmcompressor/transformers/compression/sparsity_config.py
@@ -94,15 +94,19 @@ def from_pretrained(
         if global_sparsity < 0.05:
             return None
 
-        sparsity_structure = SparsityConfigMetadata.infer_sparsity_structure(
-            model=model
+        sparsity_structure = SparsityStructure(
+            SparsityConfigMetadata.infer_sparsity_structure(model=model)
         )
         if is_marlin:
             # sparse compressor should be dense for marlin
             # compression
             format = CompressionFormat.dense.value
         if compress:
-            format = CompressionFormat.sparse_bitmask.value
+            format = (
+                CompressionFormat.sparse_bitmask.value
+                if sparsity_structure == SparsityStructure.TWO_FOUR
+                else CompressionFormat.sparse_24_bytemask.value
+            )
 
         else:
             format = CompressionFormat.dense.value