More robust load device autodetection

HabanaAI · Sep 24, 2024 · 7adace5 · 7adace5
1 parent 838aedc
commit 7adace5
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
@@ -844,11 +844,13 @@ def create_model_config(self) -> ModelConfig:
             mm_processor_kwargs=self.mm_processor_kwargs,
         )
 
-    def create_load_config(self) -> LoadConfig:
+    def create_load_config(self, load_device=None) -> LoadConfig:
+        if load_device is None:
+            load_device = DeviceConfig(device=self.device).device
         return LoadConfig(
             load_format=self.load_format,
             download_dir=self.download_dir,
-            device=self.load_device,
+            device=load_device,
             model_loader_extra_config=self.model_loader_extra_config,
             ignore_patterns=self.ignore_patterns,
         )
@@ -1032,9 +1034,9 @@ def create_engine_config(self) -> EngineConfig:
             self.model_loader_extra_config[
                 "qlora_adapter_name_or_path"] = self.qlora_adapter_name_or_path
 
-        self.load_device = device_config.device if self.weights_load_device is \
+        load_device = device_config.device if self.weights_load_device is \
             None else self.weights_load_device
-        load_config = self.create_load_config()
+        load_config = self.create_load_config(load_device)
 
         prompt_adapter_config = PromptAdapterConfig(
             max_prompt_adapters=self.max_prompt_adapters,