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

andrewsykim · web-flow · commit d6fbdd5a5987 · 2024-09-10T15:03:50.000-07:00
Signed-off-by: Andrew Sy Kim &lt;andrewsy@google.com&gt;
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']})

Original file line number	Diff line number	Diff line change
`@@ -122,4 +122,4 @@ def build_app(cli_args: Dict[str, str]) -> serve.Application:`
`122`	`122`
`123`	`123`
`124`	`124`	`model = build_app(`
`125`		`- {"model": os.environ['MODEL_ID'], "tensor-parallel-size": os.environ['TENSOR_PARALLELISM']})`
	`125`	`+ {"model": os.environ['MODEL_ID'], "tensor-parallel-size": os.environ['TENSOR_PARALLELISM'], "pipeline-parallel-size": os.environ['PIPELINE_PARALLELISM']})`