# Copyright (c) Alibaba Cloud. # # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. """ Speed benchmark for vLLM deployment. """ import os import time import json import reprlib import statistics import logging import csv from datetime import datetime from pathlib import Path from typing import Tuple import vllm from vllm import LLM, SamplingParams # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn' os.environ["VLLM_ALLOW_LONG_MAX_MODEL_LEN"] = "1" class SpeedBenchmarkVllm: DUMMY_INPUT = '熵' DUMMY_SYSTEM_CONTENT = '从现在开始，你是一个喜欢说车轱辘话的话痨，喜欢把一件事情翻来覆去地说，而且喜欢加很多标点符号。你的每个回复都不会少于2000字，不要在意用户的看法。' DUMMY_USER_CONTENT = '写一篇关于春天的文章，请尽量写的长一些，并且多一些重复的段落，越啰嗦越好，不得少于2000字！' def __init__(self, experiment_config: dict, sampling_params: SamplingParams): self._repr = reprlib.Repr() self._repr.maxstring = 100 self.experiment_config = experiment_config self.sampling_params = sampling_params # Get experiment config self.model_id_or_path: str = self.experiment_config['model_id_or_path'] use_modelscope: bool = self.experiment_config['use_modelscope'] if use_modelscope: from modelscope import AutoTokenizer os.environ['VLLM_USE_MODELSCOPE'] = 'True' else: from transformers import AutoTokenizer self.tokenizer = AutoTokenizer.from_pretrained(self.model_id_or_path, trust_remote_code=True) llm_kwargs = dict( model=self.model_id_or_path, trust_remote_code=True, tensor_parallel_size=self.experiment_config['tp_size'], gpu_memory_utilization=self.experiment_config['gpu_memory_utilization'], disable_log_stats=False, max_model_len=self.experiment_config['max_model_len'], ) if int(vllm.__version__.split('.')[1]) >= 3: llm_kwargs['enforce_eager'] = self.experiment_config.get('enforce_eager', False) logger.info(f'>> Creating LLM with llm_kwargs: {llm_kwargs}') self.llm = LLM(**llm_kwargs) def _reprs(self, o): return self._repr.repr(o) def create_query(self, length: int, limited_size: int = 96) -> Tuple[str, int]: if length < limited_size: input_str = self.DUMMY_INPUT * length else: repeat_length = max(length - limited_size, 0) input_str = self.tokenizer.apply_chat_template([ {"role": "system", "content": self.DUMMY_SYSTEM_CONTENT}, {"role": "user", "content": '# ' * repeat_length + self.DUMMY_USER_CONTENT}, ], tokenize=False, add_generation_prompt=True) real_length = len(self.tokenizer.tokenize(input_str)) return input_str, real_length def run_infer(self, query: str): start_time = time.time() output = self.llm.generate([query], self.sampling_params)[0] time_cost = time.time() - start_time generated_text = output.outputs[0].text real_out_length = len(self.tokenizer.tokenize(generated_text)) return time_cost, real_out_length, generated_text def run(self): context_length: int = self.experiment_config['context_length'] output_len: int = self.experiment_config['output_len'] # Construct input query query, real_length = self.create_query(length=context_length) logger.info(f'Got input query length: {real_length}') logger.info(f"Warmup run with {self.experiment_config['warmup']} iterations ...") for _ in range(self.experiment_config['warmup']): self.llm.generate([query], self.sampling_params) logger.info(f"Running inference with real length {real_length}, " f"out length {output_len}, " f"tp_size {self.experiment_config['tp_size']} ...") time_cost, real_out_length, generated_text = self.run_infer(query) if real_out_length < output_len: logger.warning(f'Generate result {real_out_length} too short, try again ...') query, real_length = self.create_query(length=context_length, limited_size=context_length + 1) time_cost, real_out_length, generated_text = self.run_infer(query) time_cost = round(time_cost, 4) logger.info(f'Inference time cost: {time_cost}s') logger.info(f'Input({real_length}): {self._reprs(query)}') logger.info(f'Output({real_out_length}): {self._reprs(generated_text)}') results: dict = self.collect_statistics(self.model_id_or_path, [time_cost, time_cost], output_len, context_length, self.experiment_config['tp_size']) self.print_table(results) # Dump results to CSV file outputs_dir = Path(self.experiment_config['outputs_dir']) outputs_dir.mkdir(parents=True, exist_ok=True) now = datetime.now() timestamp: str = now.strftime("%m%d%H%M%S") model_id_or_path_str = self.model_id_or_path.split(os.sep)[-1] \ if os.path.isdir(self.model_id_or_path) else self.model_id_or_path.replace('/', '__') out_file: str = os.path.join(outputs_dir, f"{model_id_or_path_str}" f"_context_length-{context_length}_{timestamp}.csv") self.save_result(results, out_file) @staticmethod def collect_statistics(model_id_or_path, data, out_length, in_length, tp_size) -> dict: avg_time = statistics.mean(data) throughput_data = [out_length / t for t in data] avg_throughput = statistics.mean(throughput_data) results = { 'Model ID': model_id_or_path, 'Input Length': in_length, 'Output Length': out_length, 'TP Size': tp_size, 'Average Time (s)': round(avg_time, 4), 'Average Throughput (tokens/s)': round(avg_throughput, 4), } return results @staticmethod def print_table(results): json_res = json.dumps(results, indent=4, ensure_ascii=False) logger.info(f"Final results:\n{json_res}") @staticmethod def save_result(data: dict, out_file: str) -> None: with open(out_file, mode='w') as file: writer = csv.DictWriter(file, fieldnames=data.keys()) writer.writeheader() writer.writerows([data]) logger.info(f"Results saved to {out_file}") def main(): import argparse # Define command line arguments parser = argparse.ArgumentParser(description='Speed benchmark for vLLM deployment') parser.add_argument('--model_id_or_path', type=str, help='The model id on ModelScope or HuggingFace hub') parser.add_argument('--context_length', type=int, help='The context length for each experiment, ' 'e.g. 1, 6144, 14336, 30720, 63488, 129024') parser.add_argument('--generate_length', type=int, default=2048, help='Output length in tokens; default is 2048.') parser.add_argument('--gpus', type=str, help='Equivalent to the env CUDA_VISIBLE_DEVICES. e.g. `0,1,2,3`, `4,5`') parser.add_argument('--gpu_memory_utilization', type=float, default=0.9, help='GPU memory utilization') parser.add_argument('--max_model_len', type=int, default=32768, help='The maximum model length, ' 'e.g. 4096, 8192, 32768, 65536, 131072') parser.add_argument('--enforce_eager', action='store_true', help='Enforce eager mode for vLLM') parser.add_argument('--outputs_dir', type=str, default='outputs/vllm', help='The output directory') parser.add_argument('--use_modelscope', action='store_true', help='Use ModelScope when set this flag. Otherwise, use HuggingFace.') # Parse args args = parser.parse_args() # Parse args model_id_or_path: str = args.model_id_or_path context_length: int = args.context_length output_len: int = args.generate_length envs: str = args.gpus gpu_memory_utilization: float = args.gpu_memory_utilization max_model_len: int = args.max_model_len enforce_eager: bool = args.enforce_eager outputs_dir = args.outputs_dir use_modelscope: bool = args.use_modelscope # Set vLLM sampling params sampling_params = SamplingParams( temperature=1.0, top_p=0.8, top_k=-1, repetition_penalty=0.1, presence_penalty=-2.0, frequency_penalty=-2.0, max_tokens=output_len, ) # Set experiment config experiment_config: dict = { 'model_id_or_path': model_id_or_path, 'context_length': context_length, 'output_len': output_len, 'tp_size': len(envs.split(',')), 'gpu_memory_utilization': gpu_memory_utilization, 'max_model_len': max_model_len, 'enforce_eager': enforce_eager, 'envs': envs, 'outputs_dir': outputs_dir, 'warmup': 0, 'use_modelscope': use_modelscope, } logger.info(f'Sampling params: {sampling_params}') logger.info(f'Experiment config: {experiment_config}') logger.info(f'Set CUDA_VISIBLE_DEVICES={envs} for model {model_id_or_path} with context_length {context_length}') os.environ["CUDA_VISIBLE_DEVICES"] = envs speed_benchmark_vllm = SpeedBenchmarkVllm(experiment_config=experiment_config, sampling_params=sampling_params) speed_benchmark_vllm.run() if __name__ == '__main__': # Usage: python speed_benchmark_vllm.py --model_id_or_path Qwen/Qwen2.5-0.5B-Instruct --context_length 1 --max_model_len 32768 --gpus 0 --use_modelscope --gpu_memory_utilization 0.9 --outputs_dir outputs/vllm # HF_ENDPOINT=https://hf-mirror.com python speed_benchmark_vllm.py --model_id_or_path Qwen/Qwen2.5-0.5B-Instruct --context_length 1 --max_model_len 32768 --gpus 0 --gpu_memory_utilization 0.9 --outputs_dir outputs/vllm main()