support for sparse only ct models

neuralmagic · Dec 8, 2024 · 4820ebe · 4820ebe
1 parent 6dfc5c9
commit 4820ebe
Showing 1 changed file with 17 additions and 8 deletions.
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
@@ -367,14 +367,23 @@ def get_scheme(
         # so we do not have to re-write these functions
         # need to make accelerate optional in ct to do this
 
-        matched_target = find_matched_target(
-            layer_name=layer_name,
-            module=layer,
-            targets=self.target_scheme_map.keys())
-
-        scheme_dict = self.target_scheme_map[matched_target]
-        weight_quant = scheme_dict.get("weights")
-        input_quant = scheme_dict.get("input_activations")
+        # Will be empty for models with only sparsity
+        if self.target_scheme_map:
+            matched_target = find_matched_target(
+                layer_name=layer_name,
+                module=layer,
+                targets=self.target_scheme_map.keys())
+
+            scheme_dict = self.target_scheme_map[matched_target]
+            weight_quant = scheme_dict.get("weights")
+            input_quant = scheme_dict.get("input_activations")
+        elif self.sparsity_scheme_map:
+            matched_target = find_matched_target(
+                layer_name=layer_name,
+                module=layer,
+                targets=self.sparsity_scheme_map.keys())
+            weight_quant = None
+            input_quant = None
 
         sparsity_scheme: Optional[
             SparsityCompressionConfig] = self.sparsity_scheme_map.get(