vllm 0.2.1.post1 - speed boost, quantization, mistral support

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 

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`

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 

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
     })