| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Benchmark offline prioritization.""" |
| 4 | |
| 5 | import argparse |
| 6 | import json |
| 7 | import random |
| 8 | import time |
| 9 | |
| 10 | from transformers import AutoTokenizer, PreTrainedTokenizerBase |
| 11 | |
| 12 | from vllm.engine.arg_utils import EngineArgs |
| 13 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 14 | |
| 15 | |
| 16 | # Select a equi-probable random priority |
| 17 | def get_random_flag(): |
| 18 | return 0 if random.random() < 0.5 else 1 |
| 19 | |
| 20 | |
| 21 | def sample_requests( |
| 22 | dataset_path: str, |
| 23 | num_requests: int, |
| 24 | tokenizer: PreTrainedTokenizerBase, |
| 25 | fixed_output_len: int | None, |
| 26 | ) -> list[tuple[str, int, int, int]]: |
| 27 | if fixed_output_len is not None and fixed_output_len < 4: |
| 28 | raise ValueError("output_len too small") |
| 29 | |
| 30 | # Load the dataset. |
| 31 | with open(dataset_path) as f: |
| 32 | dataset = json.load(f) |
| 33 | # Filter out the conversations with less than 2 turns. |
| 34 | dataset = [data for data in dataset if len(data["conversations"]) >= 2] |
| 35 | # Only keep the first two turns of each conversation. |
| 36 | dataset = [ |
| 37 | (data["conversations"][0]["value"], data["conversations"][1]["value"]) |
| 38 | for data in dataset |
| 39 | ] |
| 40 | |
| 41 | # Shuffle the dataset. |
| 42 | random.shuffle(dataset) |
| 43 | |
| 44 | # Filter out sequences that are too long or too short |
| 45 | filtered_dataset: list[tuple[str, int, int]] = [] |
| 46 | for i in range(len(dataset)): |
| 47 | if len(filtered_dataset) == num_requests: |
| 48 | break |
| 49 | |
| 50 | # Tokenize the prompts and completions. |
| 51 | prompt = dataset[i][0] |
| 52 | prompt_token_ids = tokenizer(prompt).input_ids |
| 53 | completion = dataset[i][1] |
| 54 | completion_token_ids = tokenizer(completion).input_ids |
| 55 | prompt_len = len(prompt_token_ids) |
| 56 | output_len = ( |
| 57 | len(completion_token_ids) if fixed_output_len is None else fixed_output_len |
| 58 | ) |
| 59 | if prompt_len < 4 or output_len < 4: |
| 60 | # Prune too short sequences. |
| 61 | continue |
| 62 | if prompt_len > 1024 or prompt_len + output_len > 2048: |
| 63 | # Prune too long sequences. |
| 64 | continue |
| 65 | |
| 66 | priority = get_random_flag() |
| 67 | |
| 68 | filtered_dataset.append((prompt, prompt_len, output_len, priority)) |
| 69 | |
| 70 | return filtered_dataset |
| 71 | |
| 72 | |
| 73 | def run_vllm( |
| 74 | requests: list[tuple[str, int, int]], |
| 75 | n: int, |
| 76 | engine_args: EngineArgs, |
| 77 | disable_detokenize: bool = False, |
| 78 | ) -> float: |
| 79 | from vllm import LLM, SamplingParams |
| 80 | |
| 81 | llm = LLM.from_engine_args(engine_args) |
| 82 | |
| 83 | assert all( |
| 84 | llm.llm_engine.model_config.max_model_len >= (request[1] + request[2]) |
| 85 | for request in requests |
| 86 | ), ( |
| 87 | "Please ensure that max_model_len is greater than the sum of" |
| 88 | " input_len and output_len for all requests." |
| 89 | ) |
| 90 | |
| 91 | # Add the requests to the engine. |
| 92 | prompts = [] |
| 93 | sampling_params = [] |
| 94 | priority = [] |
| 95 | for prompt, _, output_len, _priority in requests: |
| 96 | prompts.append(prompt) |
| 97 | priority.append(_priority) |
| 98 | sampling_params.append( |
| 99 | SamplingParams( |
| 100 | n=n, |
| 101 | temperature=1.0, |
| 102 | top_p=1.0, |
| 103 | ignore_eos=True, |
| 104 | max_tokens=output_len, |
| 105 | detokenize=not disable_detokenize, |
| 106 | ) |
| 107 | ) |
| 108 | |
| 109 | start = time.perf_counter() |
| 110 | llm.generate(prompts, sampling_params, priority=priority, use_tqdm=True) |
| 111 | end = time.perf_counter() |
| 112 | return end - start |
| 113 | |
| 114 | |
| 115 | def main(args: argparse.Namespace): |
| 116 | print(args) |
| 117 | random.seed(args.seed) |
| 118 | |
| 119 | # Sample the requests. |
| 120 | tokenizer = AutoTokenizer.from_pretrained( |
| 121 | args.tokenizer, trust_remote_code=args.trust_remote_code |
| 122 | ) |
| 123 | if args.dataset is None: |
| 124 | # Synthesize a prompt with the given input length. |
| 125 | prompt = "hi" * (args.input_len - 1) |
| 126 | requests = [ |
| 127 | (prompt, args.input_len, args.output_len, get_random_flag()) |
| 128 | for _ in range(args.num_prompts) |
| 129 | ] |
| 130 | else: |
| 131 | requests = sample_requests( |
| 132 | args.dataset, args.num_prompts, tokenizer, args.output_len |
| 133 | ) |
| 134 | |
| 135 | if args.backend == "vllm": |
| 136 | elapsed_time = run_vllm( |
| 137 | requests, args.n, EngineArgs.from_cli_args(args), args.disable_detokenize |
| 138 | ) |
| 139 | else: |
| 140 | raise ValueError(f"Unknown backend: {args.backend}") |
| 141 | total_num_tokens = sum( |
| 142 | prompt_len + output_len for _, prompt_len, output_len, priority in requests |
| 143 | ) |
| 144 | print( |
| 145 | f"Throughput: {len(requests) / elapsed_time:.2f} requests/s, " |
| 146 | f"{total_num_tokens / elapsed_time:.2f} tokens/s" |
| 147 | ) |
| 148 | |
| 149 | # Output JSON results if specified |
| 150 | if args.output_json: |
| 151 | results = { |
| 152 | "elapsed_time": elapsed_time, |
| 153 | "num_requests": len(requests), |
| 154 | "total_num_tokens": total_num_tokens, |
| 155 | "requests_per_second": len(requests) / elapsed_time, |
| 156 | "tokens_per_second": total_num_tokens / elapsed_time, |
| 157 | } |
| 158 | with open(args.output_json, "w") as f: |
| 159 | json.dump(results, f, indent=4) |
| 160 | |
| 161 | |
| 162 | def create_argument_parser(): |
| 163 | parser = FlexibleArgumentParser(description="Benchmark the throughput.") |
| 164 | parser.add_argument( |
| 165 | "--backend", type=str, choices=["vllm", "hf", "mii"], default="vllm" |
| 166 | ) |
| 167 | parser.add_argument( |
| 168 | "--dataset", type=str, default=None, help="Path to the dataset." |
| 169 | ) |
| 170 | parser.add_argument( |
| 171 | "--input-len", |
| 172 | type=int, |
| 173 | default=None, |
| 174 | help="Input prompt length for each request", |
| 175 | ) |
| 176 | parser.add_argument( |
| 177 | "--output-len", |
| 178 | type=int, |
| 179 | default=None, |
| 180 | help="Output length for each request. Overrides the " |
| 181 | "output length from the dataset.", |
| 182 | ) |
| 183 | parser.add_argument( |
| 184 | "--n", type=int, default=1, help="Number of generated sequences per prompt." |
| 185 | ) |
| 186 | parser.add_argument( |
| 187 | "--num-prompts", type=int, default=200, help="Number of prompts to process." |
| 188 | ) |
| 189 | parser.add_argument( |
| 190 | "--output-json", |
| 191 | type=str, |
| 192 | default=None, |
| 193 | help="Path to save the throughput results in JSON format.", |
| 194 | ) |
| 195 | parser.add_argument( |
| 196 | "--disable-detokenize", |
| 197 | action="store_true", |
| 198 | help=( |
| 199 | "Do not detokenize responses (i.e. do not include " |
| 200 | "detokenization time in the latency measurement)" |
| 201 | ), |
| 202 | ) |
| 203 | |
| 204 | parser = EngineArgs.add_cli_args(parser) |
| 205 | |
| 206 | return parser |
| 207 | |
| 208 | |
| 209 | if __name__ == "__main__": |
| 210 | parser = create_argument_parser() |
| 211 | args = parser.parse_args() |
| 212 | if args.tokenizer is None: |
| 213 | args.tokenizer = args.model |
| 214 | if args.dataset is None: |
| 215 | assert args.input_len is not None |
| 216 | assert args.output_len is not None |
| 217 | else: |
| 218 | assert args.input_len is None |
| 219 | |
| 220 | main(args) |