From 3ab0770b1b69d27b2449a204788016a43d41efbc Mon Sep 17 00:00:00 2001
From: Jorg Doku <jd728@cornell.edu>
Date: Wed, 9 Aug 2023 10:37:43 -0500
Subject: [PATCH] including env variable for tensor parallelism / multi-gpu

---
 src/handler.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/handler.py b/src/handler.py
index 142bb2d..5a3b5e1 100644
--- a/src/handler.py
+++ b/src/handler.py
@@ -17,12 +17,15 @@ TOKENIZER = os.environ.get('TOKENIZER', None)
 if not MODEL_NAME:
     print("Error: The model has not been provided.")
 
+# Tensor parallelism
+TENSOR_PARALLEL_SIZE = os.environ.get('TENSOR_PARALLEL_SIZE', 1)
+
 # Prepare the engine's arguments
 engine_args = AsyncEngineArgs(
     model=f"{MODEL_BASE_PATH}{MODEL_NAME.split('/')[1]}",
     tokenizer=TOKENIZER,
     tokenizer_mode="auto",
-    tensor_parallel_size=1,
+    tensor_parallel_size=TENSOR_PARALLEL_SIZE,
     dtype="auto",
     seed=0,
     worker_use_ray=False,