Don't use cupy when `enforce_eager=True` (#3037)

c1c0d00b · Roy · GitHub · d9f726c4 · c1c0d00b
Unverified Commit c1c0d00b authored Feb 27, 2024 by Roy Committed by GitHub Feb 26, 2024
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

vllm/engine/llm_engine.py vllm/engine/llm_engine.py +4 -1

No files found.
--- a/vllm/engine/llm_engine.py
+++ b/vllm/engine/llm_engine.py
@@ -284,7 +284,10 @@ class LLMEngine:
            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.