add support for pipeline-parallel-size in vLLM example (ray-project#2370

) Signed-off-by: Andrew Sy Kim <[email protected]>
kevin85421 · Sep 10, 2024 · d6fbdd5 · d6fbdd5
1 parent 3e68606
commit d6fbdd5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/ray-operator/config/samples/vllm/ray-service.vllm.yaml b/ray-operator/config/samples/vllm/ray-service.vllm.yaml
@@ -20,6 +20,7 @@ spec:
         env_vars:
           MODEL_ID: "meta-llama/Meta-Llama-3-8B-Instruct"
           TENSOR_PARALLELISM: "2"
+          PIPELINE_PARALLELISM: "1"
   rayClusterConfig:
     headGroupSpec:
       rayStartParams:

diff --git a/ray-operator/config/samples/vllm/serve.py b/ray-operator/config/samples/vllm/serve.py
@@ -122,4 +122,4 @@ def build_app(cli_args: Dict[str, str]) -> serve.Application:
 
 
 model = build_app(
-    {"model": os.environ['MODEL_ID'], "tensor-parallel-size": os.environ['TENSOR_PARALLELISM']})
+    {"model": os.environ['MODEL_ID'], "tensor-parallel-size": os.environ['TENSOR_PARALLELISM'], "pipeline-parallel-size": os.environ['PIPELINE_PARALLELISM']})