Fix Tensor Parallel

Dockerfile

@@ -1,5 +1,5 @@
 ARG WORKER_CUDA_VERSION=11.8.0
-FROM runpod/worker-vllm:base-0.2.0-cuda${WORKER_CUDA_VERSION} AS vllm-base
+FROM runpod/worker-vllm:base-0.2.2-cuda${WORKER_CUDA_VERSION} AS vllm-base
 
 RUN apt-get update -y \
     && apt-get install -y python3-pip
@@ -15,26 +15,26 @@ COPY src /src
 
 # Setup for Option 2: Building the Image with the Model included
 ARG MODEL_NAME=""
-ARG MODEL_BASE_PATH="/runpod-volume"
+ARG BASE_PATH="/runpod-volume"
 ARG QUANTIZATION=""
 
-ENV MODEL_BASE_PATH=$MODEL_BASE_PATH \
-    MODEL_NAME=$MODEL_NAME \
+ENV MODEL_NAME=$MODEL_NAME \
+    BASE_PATH=$BASE_PATH \
     QUANTIZATION=$QUANTIZATION \
-    HF_DATASETS_CACHE="${MODEL_BASE_PATH}/huggingface-cache/datasets" \
-    HUGGINGFACE_HUB_CACHE="${MODEL_BASE_PATH}/huggingface-cache/hub" \
-    HF_HOME="${MODEL_BASE_PATH}/huggingface-cache/hub" \
+    HF_DATASETS_CACHE="${BASE_PATH}/huggingface-cache/datasets" \
+    HUGGINGFACE_HUB_CACHE="${BASE_PATH}/huggingface-cache/hub" \
+    HF_HOME="${BASE_PATH}/huggingface-cache/hub" \
     HF_TRANSFER=1 
+
+ENV PYTHONPATH="/:/vllm-installation"
     
 RUN --mount=type=secret,id=HF_TOKEN,required=false \
     if [ -f /run/secrets/HF_TOKEN ]; then \
         export HF_TOKEN=$(cat /run/secrets/HF_TOKEN); \
     fi && \
     if [ -n "$MODEL_NAME" ]; then \
-        python3 /src/download_model.py --model $MODEL_NAME; \
+        python3 /src/download_model.py; \
     fi
 
-ENV PYTHONPATH="/:/vllm-installation"
-
 # Start the handler
-CMD ["python3", "/src/handler.py"]
+CMD ["python3", "/src/handler.py"]

README.md

@@ -42,7 +42,7 @@ We now offer a pre-built Docker Image for the vLLM Worker that you can configure
 
 <div align="center">
 
-Stable Image: ```runpod/worker-vllm:0.2.0```
+Stable Image: ```runpod/worker-vllm:0.2.1```
 
 Development Image: ```runpod/worker-vllm:dev```
 
@@ -57,9 +57,11 @@ Development Image: ```runpod/worker-vllm:dev```
    - `MODEL_NAME`: Hugging Face Model Repository (e.g., `openchat/openchat-3.5-1210`).
 
 **Optional**:
-- Model Settings:
+- LLM Settings:
+  - `TOKENIZER_NAME`: Tokenizer repository if you would like to use a different tokenizer than the one that comes with the model. (default: `None`)
+  - `CUSTOM_CHAT_TEMPLATE`: Custom chat jinja template, read more about Hugging Face chat templates [here](https://huggingface.co/docs/transformers/chat_templating). (default: `None`) 
   - `MAX_MODEL_LENGTH`: Maximum number of tokens for the engine to be able to handle. (default: maximum supported by the model)
-  - `MODEL_BASE_PATH`: Model storage directory (default: `/runpod-volume`).
+  - `BASE_PATH`: Storage directory where huggingface cache and model will be located. (default: `/runpod-volume`, which will utilize network storage if you attach it or create a local directory within the image if you don't)
   - `LOAD_FORMAT`: Format to load model in (default: `auto`).
   - `HF_TOKEN`: Hugging Face token for private and gated models (e.g., Llama, Falcon).
   - `QUANTIZATION`: AWQ (`awq`), SqueezeLLM (`squeezellm`) or GPTQ (`gptq`) Quantization. The specified Model Repo must be of a quantized model. (default: `None`)
@@ -68,12 +70,12 @@ Development Image: ```runpod/worker-vllm:dev```
 - Tensor Parallelism:
 
   Note that the more GPUs you split a model's weights accross, the slower it will be due to inter-GPU communication overhead. If you can fit the model on a single GPU, it is recommended to do so. 
-  - `USE_TENSOR_PARALLEL`: Enable (`1`) or disable (`0`) Tensor Parallelism. (default: `0`)
   - `TENSOR_PARALLEL_SIZE`: Number of GPUs to shard the model across (default: `1`).
+  - If you are having issues loading your model with Tensor Parallelism, try decreasing `VLLM_CPU_FRACTION` (default: `1`).
   
 - System Settings:
   - `GPU_MEMORY_UTILIZATION`: GPU VRAM utilization (default: `0.98`).
-  - `MAX_PARALLEL_LOADING_WORKERS`: Maximum number of parallel workers for loading models (default: `number of available CPU cores`).
+  - `MAX_PARALLEL_LOADING_WORKERS`: Maximum number of parallel workers for loading models (default: `number of available CPU cores` if `TENSOR_PARALLEL_SIZE` is `1`, otherwise `None`).
 
 
 - Serverless Settings:
@@ -94,7 +96,7 @@ To build an image with the model baked in, you must specify the following docker
 - **Required**
   - `MODEL_NAME`
 - **Optional**
-  - `MODEL_BASE_PATH`: Defaults to `/runpod-volume` for network storage. Use `/models` or for local container storage.
+  - `BASE_PATH`: Storage directory where huggingface cache and model will be located. (default: `/runpod-volume`, which will utilize network storage if you attach it or create a local directory within the image if you don't. If your intention is to bake the model into the image, you should set this to something like `/models` to make sure there are no issues if you were to accidentally attach network storage.)
   - `QUANTIZATION`
   - `WORKER_CUDA_VERSION`: `11.8.0` or `12.1.0` (default: `11.8.0` due to a small amount of workers not having CUDA 12.1 support yet. `12.1.0` is recommended for optimal performance).
 
@@ -102,7 +104,7 @@ For the remaining settings, you may apply them as environment variables when run
 
 #### Example: Building an image with OpenChat-3.5
 ```bash
-sudo docker build -t username/image:tag --build-arg MODEL_NAME="openchat/openchat_3.5" --build-arg MODEL_BASE_PATH="/models" .
+sudo docker build -t username/image:tag --build-arg MODEL_NAME="openchat/openchat_3.5" --build-arg BASE_PATH="/models" .
 ```
 
 ##### (Optional) Including Huggingface Token

src/download_model.py

@@ -1,22 +1,25 @@
-import argparse
 import os
+import logging
 from vllm.model_executor.weight_utils import prepare_hf_model_weights
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model", type=str)
-    parser.add_argument(
-        "--download_dir", type=str, default=os.environ.get("MODEL_BASE_PATH")
-    )
-
-    args = parser.parse_args()
-    if not args.model or not args.download_dir:
-        raise ValueError("Must specify model and download_dir")
-
-    if not os.path.exists(args.download_dir):
-        os.makedirs(args.download_dir)
-
-    prepare_hf_model_weights(
-        model_name_or_path=args.model,
-        cache_dir=args.download_dir,
+    model = os.getenv("MODEL_NAME")
+    download_dir = os.getenv("HF_HOME")
+    if not model or not download_dir:
+        raise ValueError(f"Must specify model and download_dir. Model: {model}, download_dir: {download_dir}")
+
+    if not os.path.exists(download_dir):
+        os.makedirs(download_dir)
+    
+    logging.info(f"Downloading model {model} to {download_dir}")
+        
+    hf_folder, hf_weights_files, use_safetensors = prepare_hf_model_weights(
+        model_name_or_path=model,
+        cache_dir=download_dir,
     )
+    
+    logging.info(f"Finished downloading model {model} to {download_dir}")
+    
+    # Wrie hf_folder to file
+    with open("/local_model_path.txt", "w") as f:
+        f.write(hf_folder) 

src/engine.py

@@ -13,9 +13,12 @@ from dotenv import load_dotenv
 
 
 class Tokenizer:
-    def __init__(self, model_name: str):
+    def __init__(self, model_name):
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.has_chat_template = bool(self.tokenizer.chat_template)
+        self.custom_chat_template = os.getenv("CUSTOM_CHAT_TEMPLATE")
+        self.has_chat_template = bool(self.tokenizer.chat_template) or bool(self.custom_chat_template)
+        if self.custom_chat_template and isinstance(self.custom_chat_template, str):
+            self.tokenizer.chat_template = self.custom_chat_template
 
     def apply_chat_template(self, input: Union[str, list[dict[str, str]]]) -> str:
         if isinstance(input, list):
@@ -38,7 +41,7 @@ class vLLMEngine:
         load_dotenv() # For local development
         self.config = self._initialize_config()
         logging.info("vLLM config: %s", self.config)
-        self.tokenizer = Tokenizer(self.config["model"])
+        self.tokenizer = Tokenizer(os.getenv("TOKENIZER_NAME", os.getenv("MODEL_NAME")))
         self.llm = self._initialize_llm() if engine is None else engine
         self.openai_engine = self._initialize_openai()
         self.max_concurrency = int(os.getenv("MAX_CONCURRENCY", DEFAULT_MAX_CONCURRENCY))
@@ -162,18 +165,20 @@ class vLLMEngine:
     
     def _initialize_config(self):
         quantization = self._get_quantization()
-        dtype = "half" if quantization else "auto"
+        model, download_dir = self._get_model_name_and_path()
+        
         return {
-            "model": os.getenv("MODEL_NAME"),
-            "download_dir": os.getenv("MODEL_BASE_PATH", "/runpod-volume/"),
+            "model": model,
+            "download_dir": download_dir,
             "quantization": quantization,
             "load_format": os.getenv("LOAD_FORMAT", "auto"),
-            "dtype": dtype,
+            "dtype": "half" if quantization else "auto",
+            "tokenizer": os.getenv("TOKENIZER_NAME"),
             "disable_log_stats": bool(int(os.getenv("DISABLE_LOG_STATS", 1))),
             "disable_log_requests": bool(int(os.getenv("DISABLE_LOG_REQUESTS", 1))),
             "trust_remote_code": bool(int(os.getenv("TRUST_REMOTE_CODE", 0))),
             "gpu_memory_utilization": float(os.getenv("GPU_MEMORY_UTILIZATION", 0.95)),
-            "max_parallel_loading_workers": int(os.getenv("MAX_PARALLEL_LOADING_WORKERS", count_physical_cores())),
+            "max_parallel_loading_workers": self._get_max_parallel_loading_workers(),
             "max_model_len": self._get_max_model_len(),
             "tensor_parallel_size": self._get_num_gpu_shard(),
         }
@@ -187,22 +192,34 @@ class vLLMEngine:
     
     def _initialize_openai(self):
         if bool(int(os.getenv("ALLOW_OPENAI_FORMAT", 1))) and self.tokenizer.has_chat_template:
-            return OpenAIServingChat(self.llm, self.config["model"], "assistant")
+            return OpenAIServingChat(self.llm, self.config["model"], "assistant", self.tokenizer.tokenizer.chat_template)
         else: 
             return None
-            
+        
+    def _get_max_parallel_loading_workers(self):
+        if int(os.getenv("TENSOR_PARALLEL_SIZE", 1)) > 1:
+            return None
+        else:
+            return int(os.getenv("MAX_PARALLEL_LOADING_WORKERS", count_physical_cores()))
+        
+    def _get_model_name_and_path(self):
+        if os.path.exists("/local_model_path.txt"):
+            model, download_dir = open("/local_model_path.txt", "r").read().strip(), None
+            logging.info("Using local model at %s", model)
+        else:
+            model, download_dir = os.getenv("MODEL_NAME"), os.getenv("HF_HOME")  
+        return model, download_dir
         
     def _get_num_gpu_shard(self):
-        final_num_gpu_shard = 1
-        if bool(int(os.getenv("USE_TENSOR_PARALLEL", 0))):
-            env_num_gpu_shard = int(os.getenv("TENSOR_PARALLEL_SIZE", 1))
+        num_gpu_shard = int(os.getenv("TENSOR_PARALLEL_SIZE", 1))
+        if num_gpu_shard > 1:
             num_gpu_available = device_count()
-            final_num_gpu_shard = min(env_num_gpu_shard, num_gpu_available)
-            logging.info("Using %s GPU shards", final_num_gpu_shard)
-        return final_num_gpu_shard
+            num_gpu_shard = min(num_gpu_shard, num_gpu_available)
+            logging.info("Using %s GPU shards", num_gpu_shard)
+        return num_gpu_shard
     
     def _get_max_model_len(self):
-        max_model_len = os.getenv("MAX_MODEL_LEN")
+        max_model_len = os.getenv("MAX_MODEL_LENGTH")
         return int(max_model_len) if max_model_len is not None else None
     
     def _get_n_current_jobs(self):