Made some adjustments for tensor parallel vllm to work

deepvk · Feb 17, 2024 · e40ccab · e40ccab
1 parent b2d237c
commit e40ccab
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/lm_eval/models/vllm_causallms.py b/lm_eval/models/vllm_causallms.py
@@ -61,6 +61,9 @@ def __init__(
         device: str = "cuda",
         data_parallel_size: int = 1,
     ):
+        import ray
+        ray.init(ignore_reinit_error=True, num_cpus=8, num_gpus=tensor_parallel_size)
+
         super().__init__()
 
         if not find_spec("vllm"):