From 12d6f0778ea96418bd13460f597802631b94022f Mon Sep 17 00:00:00 2001 From: alpayariyak Date: Wed, 31 Jan 2024 03:13:21 +0000 Subject: [PATCH] Fixed Model bake-in, added Custom Chat Templates, Custom Tokenizer --- Dockerfile | 13 ++++++------- README.md | 10 ++++++---- src/download_model.py | 37 ++++++++++++++++++++----------------- src/engine.py | 30 +++++++++++++++++++++--------- 4 files changed, 53 insertions(+), 37 deletions(-) diff --git a/Dockerfile b/Dockerfile index 3919085..552471c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -15,15 +15,14 @@ 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 \ 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" @@ -33,7 +32,7 @@ RUN --mount=type=secret,id=HF_TOKEN,required=false \ 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 # Start the handler diff --git a/README.md b/README.md index 24cf9fc..f0933a2 100644 --- a/README.md +++ b/README.md @@ -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`) @@ -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 diff --git a/src/download_model.py b/src/download_model.py index 999f903..086a2a0 100644 --- a/src/download_model.py +++ b/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) \ No newline at end of file diff --git a/src/engine.py b/src/engine.py index 3f69139..2c9140c 100644 --- a/src/engine.py +++ b/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.environ.get("TOKENIZER_NAME", os.environ.get("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,13 +165,15 @@ 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))), @@ -187,10 +192,17 @@ 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_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