@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 276 lines 10.1 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Benchmark the efficiency of prefix caching.
4
5 This script allows you to benchmark the performance of
6 a model with and without prefix caching using either fixed prompts
7 or prompts sampled from the ShareGPT dataset.
8
9 Fixed example usage:
10 python benchmark_prefix_caching.py \
11 --model meta-llama/Llama-2-7b-chat-hf \
12 --enable-prefix-caching \
13 --num-prompts 1 \
14 --repeat-count 100 \
15 --input-length-range 128:256
16
17 ShareGPT example usage:
18 # This command samples 20 prompts with input lengths
19 # between 128 and 256 tokens from the ShareGPT dataset,
20 # then replicates each prompt 5 times.
21 python benchmark_prefix_caching.py \
22 --model meta-llama/Llama-2-7b-chat-hf \
23 --dataset-path /path/to/ShareGPT_V3_unfiltered_cleaned_split.json \
24 --enable-prefix-caching \
25 --num-prompts 20 \
26 --repeat-count 5 \
27 --input-length-range 128:256
28 """
29
30 import dataclasses
31 import json
32 import random
33 import time
34
35 from transformers import PreTrainedTokenizerBase
36
37 from vllm import LLM, SamplingParams
38 from vllm.engine.arg_utils import EngineArgs
39 from vllm.utils.argparse_utils import FlexibleArgumentParser
40
41 try:
42 from vllm.tokenizers import get_tokenizer
43 except ImportError:
44 from backend_request_func import get_tokenizer
45
46 PROMPT = "You are a helpful assistant in recognizes the content of tables in markdown format. Here is a table as fellows. You need to answer my question about the table.\n# Table\n|Opening|Opening|Sl. No.|Film|Cast|Director|Music Director|Notes|\n|----|----|----|----|----|----|----|----|\n|J A N|9|1|Agni Pushpam|Jayabharathi, Kamalahasan|Jeassy|M. K. Arjunan||\n|J A N|16|2|Priyamvada|Mohan Sharma, Lakshmi, KPAC Lalitha|K. S. Sethumadhavan|V. Dakshinamoorthy||\n|J A N|23|3|Yakshagaanam|Madhu, Sheela|Sheela|M. S. Viswanathan||\n|J A N|30|4|Paalkkadal|Sheela, Sharada|T. K. Prasad|A. T. Ummer||\n|F E B|5|5|Amma|Madhu, Srividya|M. Krishnan Nair|M. K. Arjunan||\n|F E B|13|6|Appooppan|Thikkurissi Sukumaran Nair, Kamal Haasan|P. Bhaskaran|M. S. Baburaj||\n|F E B|20|7|Srishti|Chowalloor Krishnankutty, Ravi Alummoodu|K. T. Muhammad|M. S. Baburaj||\n|F E B|20|8|Vanadevatha|Prem Nazir, Madhubala|Yusufali Kechery|G. Devarajan||\n|F E B|27|9|Samasya|Madhu, Kamalahaasan|K. Thankappan|Shyam||\n|F E B|27|10|Yudhabhoomi|K. P. Ummer, Vidhubala|Crossbelt Mani|R. K. Shekhar||\n|M A R|5|11|Seemantha Puthran|Prem Nazir, Jayabharathi|A. B. Raj|M. K. Arjunan||\n|M A R|12|12|Swapnadanam|Rani Chandra, Dr. Mohandas|K. G. George|Bhaskar Chandavarkar||\n|M A R|19|13|Thulavarsham|Prem Nazir, sreedevi, Sudheer|N. Sankaran Nair|V. Dakshinamoorthy||\n|M A R|20|14|Aruthu|Kaviyoor Ponnamma, Kamalahasan|Ravi|G. Devarajan||\n|M A R|26|15|Swimming Pool|Kamal Haasan, M. G. Soman|J. Sasikumar|M. K. Arjunan||\n\n# Question\nWhat' s the content in the (1,1) cells\n" # noqa: E501
47
48
49 def test_prefix(llm=None, sampling_params=None, prompts=None):
50 start_time = time.time()
51
52 llm.generate(prompts, sampling_params=sampling_params)
53
54 end_time = time.time()
55 print(f"cost time {end_time - start_time}")
56
57
58 @dataclasses.dataclass
59 class Request:
60 prompt: str
61 prompt_len: int
62 output_len: int
63
64
65 def sample_tokens(tokenizer: PreTrainedTokenizerBase, length: int) -> list[int]:
66 vocab = tokenizer.get_vocab()
67 all_special_ids = set(tokenizer.all_special_ids)
68
69 # Remove the special tokens.
70 return random.choices(
71 [v for v in vocab.values() if v not in all_special_ids],
72 k=length,
73 )
74
75
76 def sample_requests_from_dataset(
77 dataset_path: str,
78 num_requests: int,
79 tokenizer: PreTrainedTokenizerBase,
80 input_length_range: tuple[int, int],
81 fixed_output_len: int | None,
82 ) -> list[Request]:
83 if fixed_output_len is not None and fixed_output_len < 4:
84 raise ValueError("output_len too small")
85
86 # Load the dataset.
87 with open(dataset_path) as f:
88 dataset = json.load(f)
89 # Filter out the conversations with less than 2 turns.
90 dataset = [data for data in dataset if len(data["conversations"]) >= 2]
91 # Only keep the first two turns of each conversation.
92 dataset = [
93 (data["conversations"][0]["value"], data["conversations"][1]["value"])
94 for data in dataset
95 ]
96
97 # Shuffle the dataset.
98 random.shuffle(dataset)
99
100 min_len, max_len = input_length_range
101 assert min_len >= 0 and max_len >= min_len, "input_length_range too small"
102
103 # Filter out sequences that are too long or too short
104 filtered_requests: list[Request] = []
105
106 for i in range(len(dataset)):
107 if len(filtered_requests) == num_requests:
108 break
109
110 # Tokenize the prompts and completions.
111 prompt_token_ids = tokenizer(dataset[i][0]).input_ids
112 prompt = tokenizer.decode(prompt_token_ids)
113 completion = dataset[i][1]
114 completion_token_ids = tokenizer(completion).input_ids
115 prompt_len = len(prompt_token_ids)
116 output_len = (
117 len(completion_token_ids) if fixed_output_len is None else fixed_output_len
118 )
119 if min_len <= prompt_len <= max_len:
120 filtered_requests.append(Request(prompt, prompt_len, output_len))
121
122 return filtered_requests
123
124
125 def sample_requests_from_random(
126 num_requests: int,
127 tokenizer: PreTrainedTokenizerBase,
128 input_length_range: tuple[int, int],
129 fixed_output_len: int | None,
130 prefix_len: int,
131 ) -> list[Request]:
132 requests = []
133 prefix_token_ids = sample_tokens(tokenizer, prefix_len)
134 min_len, max_len = input_length_range
135
136 for i in range(num_requests):
137 unique_part_token_ids = sample_tokens(
138 tokenizer, random.randint(min_len - prefix_len, max_len - prefix_len)
139 )
140 prompt_token_ids = prefix_token_ids + unique_part_token_ids
141 prompt = tokenizer.decode(prompt_token_ids)
142 prompt_len = len(prompt_token_ids)
143 assert min_len <= prompt_len <= max_len, (
144 f"prompt_len {prompt_len} out of range {min_len}:{max_len}"
145 )
146 requests.append(Request(prompt, prompt_len, fixed_output_len))
147 return requests
148
149
150 def repeat_and_sort_requests(
151 requests: list[Request], repeat_count: int, sort: bool = False
152 ) -> list[str]:
153 repeated_requests = requests * repeat_count
154 if sort:
155 repeated_requests.sort(key=lambda x: x[1])
156 else:
157 random.shuffle(repeated_requests)
158 return [req.prompt for req in repeated_requests]
159
160
161 def main(args):
162 tokenizer = get_tokenizer(args.model, trust_remote_code=True)
163 input_length_range = tuple(map(int, args.input_length_range.split(":")))
164 random.seed(args.seed)
165 if args.dataset_path is not None:
166 if args.prefix_len > 0:
167 raise ValueError(
168 "prefix-len is not supported when dataset-path is provided."
169 )
170 print(f"Start to sample {args.num_prompts} prompts from {args.dataset_path}")
171 filtered_requests = sample_requests_from_dataset(
172 dataset_path=args.dataset_path,
173 num_requests=args.num_prompts,
174 tokenizer=tokenizer,
175 input_length_range=input_length_range,
176 fixed_output_len=args.output_len,
177 )
178 else:
179 print(f"Start to sample {args.num_prompts} prompts from random")
180 filtered_requests = sample_requests_from_random(
181 num_requests=args.num_prompts,
182 tokenizer=tokenizer,
183 input_length_range=input_length_range,
184 fixed_output_len=args.output_len,
185 prefix_len=args.prefix_len,
186 )
187
188 # Print some helpful stats of the requests.
189 print(f"Sampled {len(filtered_requests)} requests.")
190 prompt_lens = [req.prompt_len for req in filtered_requests]
191 print(f"Average input length: {sum(prompt_lens) / len(prompt_lens)}")
192 print(f"P50 input length: {sorted(prompt_lens)[len(prompt_lens) // 2]}")
193 print(f"Min Prompt Length: {min(prompt_lens)}")
194 print(f"Max Prompt Length: {max(prompt_lens)}")
195
196 engine_args = EngineArgs.from_cli_args(args)
197
198 llm = LLM.from_engine_args(engine_args)
199
200 sampling_params = SamplingParams(
201 temperature=0,
202 max_tokens=args.output_len,
203 detokenize=not args.disable_detokenize,
204 )
205
206 print("Testing filtered requests")
207 prompts = repeat_and_sort_requests(
208 filtered_requests, repeat_count=args.repeat_count, sort=args.sort
209 )
210
211 print("------start generating------")
212 test_prefix(
213 llm=llm,
214 prompts=prompts,
215 sampling_params=sampling_params,
216 )
217
218
219 def create_argument_parser():
220 parser = FlexibleArgumentParser(
221 description="Benchmark the performance with or without "
222 "automatic prefix caching."
223 )
224 parser.add_argument(
225 "--dataset-path", type=str, default=None, help="Path to the dataset."
226 )
227 parser.add_argument("--output-len", type=int, default=10)
228 parser.add_argument(
229 "--num-prompts",
230 type=int,
231 required=True,
232 help="Number of the prompts sampled from dataset",
233 )
234 parser.add_argument(
235 "--repeat-count",
236 type=int,
237 default=1,
238 help="Number of times to repeat each prompt",
239 )
240 parser.add_argument(
241 "--sort", action="store_true", help="Sort prompts by input length"
242 )
243 parser.add_argument(
244 "--input-length-range",
245 type=str,
246 required=True,
247 help="Range of input lengths for sampling prompts,"
248 'specified as "min:max" (e.g., "128:256").',
249 )
250 parser.add_argument(
251 "--prefix-len",
252 type=int,
253 default=0,
254 help="Specifies the length of a common prefix to be "
255 "added to the input prompt. The input-length-range will "
256 "subtract this length when filtering prompts. Only used "
257 "when dataset-path is not provided.",
258 )
259 parser.add_argument(
260 "--disable-detokenize",
261 action="store_true",
262 help=(
263 "Do not detokenize responses (i.e. do not include "
264 "detokenization time in the latency measurement)"
265 ),
266 )
267
268 parser = EngineArgs.add_cli_args(parser)
269
270 return parser
271
272
273 if __name__ == "__main__":
274 parser = create_argument_parser()
275 args = parser.parse_args()
276 main(args)