Set VLLM_WORKER_MULTIPROC_METHOD=spawn when tensor_parallel_size>1

ProbablyFaiz · Jul 21, 2024 · 1d57936 · 1d57936
1 parent 6d3dc31
commit 1d57936
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/rl/llm/engines.py b/rl/llm/engines.py
@@ -530,6 +530,16 @@ def _get_vllm_engine(
         )
 
     engine_args_kwargs = _get_vllm_kwargs(llm_config)
+    if (
+        engine_args_kwargs["tensor_parallel_size"] > 1
+        and "VLLM_WORKER_MULTIPROC_METHOD" not in os.environ
+    ):
+        LOGGER.warning(
+            "Setting VLLM_WORKER_MULTIPROC_METHOD to 'spawn' to avoid issues with "
+            "CUDA re-initialization."
+        )
+        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
     engine_cls = AsyncLLMEngine if use_async else LLMEngine
     engine_args_cls = AsyncEngineArgs if use_async else EngineArgs
     engine_args = engine_args_cls(**engine_args_kwargs)  # type: ignore