diff --git a/Dockerfile b/Dockerfile index 011fdde..e3ba4c0 100644 --- a/Dockerfile +++ b/Dockerfile @@ -48,6 +48,8 @@ ARG TOKENIZER= ENV TOKENIZER=$TOKENIZER ARG STREAMING= ENV STREAMING=$STREAMING +ARG QUANTIZATION= +ENV QUANTIZATION=$QUANTIZATION ENV HF_DATASETS_CACHE="/runpod-volume/huggingface-cache/datasets" ENV HUGGINGFACE_HUB_CACHE="/runpod-volume/huggingface-cache/hub" @@ -66,4 +68,4 @@ ENV MODEL_NAME=$MODEL_NAME \ RUN python -u /download_model.py # Start the handler -CMD STREAMING=$STREAMING MODEL_NAME=$MODEL_NAME MODEL_BASE_PATH=$MODEL_BASE_PATH TOKENIZER=$TOKENIZER python -u /handler.py +CMD STREAMING=$STREAMING MODEL_NAME=$MODEL_NAME MODEL_BASE_PATH=$MODEL_BASE_PATH TOKENIZER=$TOKENIZER QUANTIZATION=$QUANTIZATION python -u /handler.py diff --git a/README.md b/README.md index af5bd02..711c690 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,7 @@ 2. `MODEL_NAME`: The Hugging Face model to use. Please ensure that the chosen model is supported by vLLM. Refer to the list of supported models for compatibility. 3. `TOKENIZER`: (Optional) The specified tokenizer to use. If you want to use the default tokenizer for the model, do not provide this docker argument at all. 4. `STREAMING`: Whether to use HTTP Streaming or not. Specify True if you want to enable HTTP Streaming; otherwise, omit this argument. +5. `USE_QUANTIZATION`: (Optional) Use AWQ Quantization. Base model must be in AWQ format. #### llama2 7B Chat: `docker build . --platform linux/amd64 --build-arg HUGGING_FACE_HUB_TOKEN=your_hugging_face_token_here --build-arg MODEL_NAME=meta-llama/Llama-2-7b-chat-hf --build-arg TOKENIZER=hf-internal-testing/llama-tokenizer --build-arg STREAMING=True` diff --git a/builder/requirements.txt b/builder/requirements.txt index cf7ece2..70aa770 100644 --- a/builder/requirements.txt +++ b/builder/requirements.txt @@ -1,6 +1,6 @@ # Required Python packages get listed here, one per line. # Recomended to lock the version number to avoid unexpected changes. fastapi==0.99.1 -vllm==0.1.7 -huggingface-hub==0.16.4 +vllm==0.2.1.post1 +huggingface-hub==0.17.3 runpod @ git+https://github.com/runpod/runpod-python@main#egg=runpod diff --git a/src/handler.py b/src/handler.py index 5ab4882..682f8cb 100644 --- a/src/handler.py +++ b/src/handler.py @@ -16,6 +16,7 @@ MODEL_BASE_PATH = os.environ.get('MODEL_BASE_PATH', '/runpod-volume/') STREAMING = os.environ.get('STREAMING', False) == 'True' TOKENIZER = os.environ.get('TOKENIZER', None) USE_FULL_METRICS = os.environ.get('USE_FULL_METRICS', True) +QUANTIZATION = os.environ.get('QUANTIZATION', None) if not MODEL_NAME: print("Error: The model has not been provided.") @@ -27,17 +28,21 @@ except ValueError: print("Error: NUM_GPU_SHARD should be an integer. Using default value of 1.") NUM_GPU_SHARD = 1 +# Setup quantization parameter +if type(QUANTIZATION) is str and QUANTIZATION.lower() != "awq": + QUANTIZATION = None + print("Invalid quantization parameter. Using default value of None.") + # 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=NUM_GPU_SHARD, - dtype="auto", + dtype="auto" if QUANTIZATION is None else "half", seed=0, - max_num_batched_tokens=8192, disable_log_stats=False, - # max_num_seqs=256, + quantization=QUANTIZATION, ) # Create the vLLM asynchronous engine @@ -47,6 +52,7 @@ llm = AsyncLLMEngine.from_engine_args(engine_args) llm.engine._log_system_stats = lambda x, y: vllm_log_system_stats( llm.engine, x, y) + def concurrency_controller() -> bool: # Calculate pending sequences total_pending_sequences = len(llm.engine.scheduler.waiting) + len(llm.engine.scheduler.swapped) @@ -55,6 +61,7 @@ def concurrency_controller() -> bool: # Enable auto-scaling if pending sequences exist return total_pending_sequences > 30 + def prepare_metrics() -> dict: # The vLLM metrics are updated every 5 seconds, see metrics.py for the _LOGGING_INTERVAL_SEC field. if hasattr(llm.engine, 'metrics'): @@ -62,6 +69,7 @@ def prepare_metrics() -> dict: else: return {} + # Validation def validate_sampling_params(sampling_params): def validate_int(value, default): @@ -162,7 +170,7 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None] self.stream_index = 0 def inc_stream_idx(self): - self.stream_index +=1 + self.stream_index += 1 tracker = Tracker() @@ -243,9 +251,9 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None] # Include metrics for the job. runpod.serverless.modules.rp_metrics.metrics_collector.push_metrics_internal( - job_id=job['id'], + job_id=job['id'], metrics=runpod_metrics - ) + ) # Keep track of the final output final_output = request_output @@ -269,10 +277,10 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None] "input_tokens": input_tokens, "output_tokens": output_tokens, } - + # Update the aggregate transformation function runpod.serverless.modules.rp_metrics.metrics_collector.update_stream_aggregate( - job_id=job['id'], + job_id=job['id'], aggregate_function=aggregate_function ) @@ -336,7 +344,7 @@ async def handler(job: dict) -> dict[str, list]: # Include metrics for the job. runpod.serverless.modules.rp_metrics.metrics_collector.push_metrics_internal( - job_id=job['id'], + job_id=job['id'], metrics=runpod_metrics ) @@ -347,18 +355,19 @@ async def handler(job: dict) -> dict[str, list]: } return ret + # Start the serverless worker with appropriate settings if STREAMING: print("Starting the vLLM serverless worker with streaming enabled.") runpod.serverless.start({ - "handler": handler_streaming, - "concurrency_controller": concurrency_controller, + "handler": handler_streaming, + "concurrency_controller": concurrency_controller, "return_aggregate_stream": True }) else: print("Starting the vLLM serverless worker with streaming disabled.") runpod.serverless.start({ - "handler": handler, - "concurrency_controller": - concurrency_controller + "handler": handler, + "concurrency_controller": + concurrency_controller })