Enable sparse24bytemask compressor

vllm-project · Dec 12, 2024 · 8699a24 · 8699a24
1 parent 938a526
commit 8699a24
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/llmcompressor/transformers/compression/sparsity_config.py b/src/llmcompressor/transformers/compression/sparsity_config.py
@@ -102,7 +102,11 @@ def from_pretrained(
             # compression
             format = CompressionFormat.dense.value
         if compress:
-            format = CompressionFormat.sparse_bitmask.value
+            format = (
+                CompressionFormat.sparse_24_bytemask.value
+                if sparsity_structure == SparsityStructure.TWO_FOUR.value
+                else CompressionFormat.sparse_bitmask.value
+            )
 
         else:
             format = CompressionFormat.dense.value