Don't use cupy when enforce_eager=True (vllm-project#3037)

xjpang · Mar 4, 2024 · d55c43c · d55c43c
1 parent 88ffe97
commit d55c43c
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py
@@ -284,7 +284,10 @@ def _init_workers_ray(self, placement_group: "PlacementGroup",
             is_driver_worker=True,
         )
 
-        self._run_workers("init_model", cupy_port=get_open_port())
+        # don't use cupy for eager mode
+        self._run_workers("init_model",
+                          cupy_port=get_open_port()
+                          if not model_config.enforce_eager else None)
         self._run_workers(
             "load_model",
             max_concurrent_workers=self.parallel_config.