Use correct model_source for loading GPTQ models with FMS (#61)

IBM · Jan 17, 2025 · 7e068b5 · 7e068b5
2 parents 590042e + d19ff3d
commit 7e068b5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/vllm/core/scheduler.py b/vllm/core/scheduler.py
@@ -919,8 +919,9 @@ def _schedule_prefills(
         ignored_seq_groups: List[SequenceGroup] = []
         seq_groups: List[ScheduledSequenceGroup] = []
 
-        applicable_spyre_warmup_shapes = list(
-            self.scheduler_config.spyre_warmup_shapes)
+        if self.scheduler_config.spyre_scheduling_enabled:
+            applicable_spyre_warmup_shapes = list(
+                self.scheduler_config.spyre_warmup_shapes)
 
         waiting_queue = self.waiting
 

diff --git a/vllm/model_executor/model_loader/spyre.py b/vllm/model_executor/model_loader/spyre.py
@@ -143,7 +143,7 @@ def load_weights(self, model_config: ModelConfig, max_prompt_length: int,
                 "desc_act": quant_cfg['desc_act'],
             }
             data_type = None
-            model_source = "llama_gptq_hf_unfused_aiu"
+            model_source = "hf_gptq_aiu"
         else:
             linear_config = {"linear_type": "torch_linear"}
             data_type = self.dtype