owen/worker-vllm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fixed Model bake-in, added Custom Chat Templates, Custom Tokenizer

Browse Source
This commit is contained in:
alpayariyak
2024-01-31 03:13:21 +00:00
parent fa5556434c
commit 12d6f0778e
4 changed files with 53 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Dockerfile
+6 -7
View File
@@ -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
README.md
+6 -4
View File
@@ -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
src/download_model.py
+18 -15
View File
@@ -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")
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,
)
args = parser.parse_args()
if not args.model or not args.download_dir:
raise ValueError("Must specify model and download_dir")
logging.info(f"Finished downloading model {model} to {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,
)
# Wrie hf_folder to file
with open("/local_model_path.txt", "w") as f:
f.write(hf_folder)
src/engine.py
+20 -8
View File
@@ -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