@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1253b15774
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 19:45 UTC

[Feature] Consolidate performance benchmark datasets (#14036)

Signed-off-by: Jennifer Zhao <7443418+JenZhao@users.noreply.github.com> Signed-off-by: Roger Wang <ywang@roblox.com> Co-authored-by: Jennifer Zhao <7443418+JenZhao@users.noreply.github.com> Co-authored-by: Roger Wang <ywang@roblox.com>

Jennifer Zhao committed Mar 10, 2025 at 00:23 UTC 1253b1577408f7981d11495b1fda71cbcbe48dc4
3 files changed +825 -579
benchmarks/benchmark_dataset.py new
+667
@@ -0,0 +1,667 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""
3 +This module defines a framework for sampling benchmark requests from various
4 +datasets. Each dataset subclass of BenchmarkDataset must implement sample
5 +generation. Supported dataset types include:
6 + - ShareGPT
7 + - Random (synthetic)
8 + - Sonnet
9 + - BurstGPT
10 + - HuggingFace
11 + - VisionArena
12 +
13 +TODO: Implement CustomDataset to parse a JSON file and convert its contents into
14 +SampleRequest instances, similar to the approach used in ShareGPT.
15 +"""
16 +
17 +import base64
18 +import io
19 +import json
20 +import random
21 +from abc import ABC, abstractmethod
22 +from collections.abc import Mapping
23 +from dataclasses import dataclass
24 +from functools import cache
25 +from typing import Any, Optional, Union
26 +
27 +import numpy as np
28 +import pandas as pd
29 +from datasets import load_dataset
30 +from PIL import Image
31 +from transformers import PreTrainedTokenizerBase
32 +
33 +from vllm.lora.request import LoRARequest
34 +from vllm.lora.utils import get_adapter_absolute_path
35 +from vllm.multimodal import MultiModalDataDict
36 +from vllm.transformers_utils.tokenizer import AnyTokenizer, get_lora_tokenizer
37 +
38 +# -----------------------------------------------------------------------------
39 +# Data Classes
40 +# -----------------------------------------------------------------------------
41 +
42 +
43 +@dataclass
44 +class SampleRequest:
45 + """
46 + Represents a single inference request for benchmarking.
47 + """
48 +
49 + prompt: str
50 + prompt_len: int
51 + expected_output_len: int
52 + multi_modal_data: Optional[Union[MultiModalDataDict, dict]] = None
53 + lora_request: Optional[LoRARequest] = None
54 +
55 +
56 +# -----------------------------------------------------------------------------
57 +# Benchmark Dataset Base Class
58 +# -----------------------------------------------------------------------------
59 +
60 +
61 +class BenchmarkDataset(ABC):
62 + DEFAULT_SEED = 0
63 +
64 + # num_requests has default 1000 in both the benchmark_serving.py and
65 + # benchmark_throughput.py
66 +
67 + def __init__(
68 + self,
69 + dataset_path: Optional[str] = None,
70 + random_seed: int = DEFAULT_SEED,
71 + ) -> None:
72 + """
73 + Initialize the BenchmarkDataset with an optional dataset path and random
74 + seed. Args:
75 + dataset_path (Optional[str]): Path to the dataset. If None, it
76 + indicates that a default or random dataset might be used.
77 + random_seed (int): Seed value for reproducible shuffling or
78 + sampling. Defaults to DEFAULT_SEED.
79 + """
80 + self.dataset_path = dataset_path
81 + # Set the random seed, ensuring that a None value is replaced with the
82 + # default seed.
83 + self.random_seed = (random_seed
84 + if random_seed is not None else self.DEFAULT_SEED)
85 + self.data = None
86 +
87 + def load_data(self) -> None:
88 + """
89 + Load data from the dataset path into self.data.
90 +
91 + This method must be overridden by subclasses since the method to load
92 + data will vary depending on the dataset format and source.
93 +
94 + Raises:
95 + NotImplementedError: If a subclass does not implement this method.
96 + """
97 + # TODO (jenniferzhao): add support for downloading data
98 + raise NotImplementedError(
99 + "load_data must be implemented in subclasses.")
100 +
101 + def get_random_lora_request(
102 + self,
103 + tokenizer: PreTrainedTokenizerBase,
104 + max_loras: Optional[int] = None,
105 + lora_path: Optional[str] = None,
106 + ) -> tuple[Optional[LoRARequest], AnyTokenizer]:
107 + """
108 + Optionally select a random LoRA request and return its associated
109 + tokenizer.
110 +
111 + This method is used when LoRA parameters are provided. It randomly
112 + selects a LoRA based on max_loras and retrieves a cached tokenizer for
113 + that LoRA if available. Otherwise, it returns the base tokenizer.
114 +
115 + Args:
116 + tokenizer (PreTrainedTokenizerBase): The base tokenizer to use if no
117 + LoRA is selected. max_loras (Optional[int]): The maximum number of
118 + LoRAs available. If None, LoRA is not used. lora_path
119 + (Optional[str]): Path to the LoRA parameters on disk. If None, LoRA
120 + is not used.
121 +
122 + Returns:
123 + tuple[Optional[LoRARequest], AnyTokenizer]: A tuple where the first
124 + element is a LoRARequest (or None if not applicable) and the second
125 + element is the tokenizer associated with the LoRA request (or the
126 + base tokenizer).
127 + """
128 + if max_loras is None or lora_path is None:
129 + return None, tokenizer
130 +
131 + # Generate a random LoRA ID in the range [1, max_loras].
132 + lora_id = random.randint(1, max_loras)
133 + lora_request = LoRARequest(
134 + lora_name=str(lora_id),
135 + lora_int_id=lora_id,
136 + lora_path=lora_path_on_disk(lora_path),
137 + )
138 + if lora_id not in lora_tokenizer_cache:
139 + lora_tokenizer_cache[lora_id] = get_lora_tokenizer(lora_request)
140 + # Return lora_request and the cached tokenizer if available; otherwise,
141 + # return the base tokenizer
142 + return lora_request, lora_tokenizer_cache[lora_id] or tokenizer
143 +
144 + @abstractmethod
145 + def sample(self, tokenizer: PreTrainedTokenizerBase,
146 + num_requests: int) -> list[SampleRequest]:
147 + """
148 + Abstract method to generate sample requests from the dataset.
149 +
150 + Subclasses must override this method to implement dataset-specific logic
151 + for generating a list of SampleRequest objects.
152 +
153 + Args:
154 + tokenizer (PreTrainedTokenizerBase): The tokenizer to be used
155 + for processing the dataset's text.
156 + num_requests (int): The number of sample requests to generate.
157 +
158 + Returns:
159 + list[SampleRequest]: A list of sample requests generated from the
160 + dataset.
161 + """
162 + raise NotImplementedError("sample must be implemented in subclasses.")
163 +
164 +
165 +# -----------------------------------------------------------------------------
166 +# Utility Functions and Global Caches
167 +# -----------------------------------------------------------------------------
168 +
169 +
170 +def is_valid_sequence(
171 + prompt_len: int,
172 + output_len: int,
173 + min_len: int = 4,
174 + max_prompt_len: int = 1024,
175 + max_total_len: int = 2048,
176 + skip_min_output_len_check: bool = False,
177 +) -> bool:
178 + """
179 + Validate a sequence based on prompt and output lengths.
180 +
181 + Default pruning criteria are copied from the original `sample_hf_requests`
182 + and `sample_sharegpt_requests` functions in benchmark_serving.py, as well as
183 + from `sample_requests` in benchmark_throughput.py.
184 + """
185 + # Check for invalid conditions
186 + prompt_too_short = prompt_len < min_len
187 + output_too_short = (not skip_min_output_len_check) and (output_len
188 + < min_len)
189 + prompt_too_long = prompt_len > max_prompt_len
190 + combined_too_long = (prompt_len + output_len) > max_total_len
191 +
192 + # Return True if none of the invalid conditions are met
193 + return not (prompt_too_short or output_too_short or prompt_too_long
194 + or combined_too_long)
195 +
196 +
197 +@cache
198 +def lora_path_on_disk(lora_path: str) -> str:
199 + return get_adapter_absolute_path(lora_path)
200 +
201 +
202 +# Global cache for LoRA tokenizers.
203 +lora_tokenizer_cache: dict[int, AnyTokenizer] = {}
204 +
205 +
206 +def process_image(image: Any) -> Mapping[str, Any]:
207 + """
208 + Process a single image input and return a multimedia content dictionary.
209 +
210 + For a PIL.Image.Image input:
211 + - Converts the image to RGB.
212 + - Saves the image as a JPEG in-memory.
213 + - Encodes the JPEG data as a base64 string.
214 + - Returns a dictionary with the image as a base64 data URL.
215 +
216 + For a string input:
217 + - Treats the string as a URL or file path.
218 + - Prepends "file://" if the string doesn't start with "http://" or
219 + "file://".
220 + - Returns a dictionary with the image URL.
221 +
222 + Raises:
223 + ValueError: If the input is neither a PIL.Image.Image nor a string.
224 + """
225 + if isinstance(image, Image.Image):
226 + image = image.convert("RGB")
227 + with io.BytesIO() as image_data:
228 + image.save(image_data, format="JPEG")
229 + image_base64 = base64.b64encode(
230 + image_data.getvalue()).decode("utf-8")
231 + return {
232 + "type": "image_url",
233 + "image_url": {
234 + "url": f"data:image/jpeg;base64,{image_base64}"
235 + },
236 + }
237 +
238 + if isinstance(image, str):
239 + image_url = (image if image.startswith(
240 + ("http://", "file://")) else f"file://{image}")
241 + return {"type": "image_url", "image_url": {"url": image_url}}
242 +
243 + raise ValueError(
244 + f"Invalid image input {image}. Must be a PIL.Image.Image or str.")
245 +
246 +
247 +# -----------------------------------------------------------------------------
248 +# Random Dataset Implementation (Synthetic Data)
249 +# -----------------------------------------------------------------------------
250 +
251 +
252 +class RandomDataset(BenchmarkDataset):
253 + # Default values copied from benchmark_serving.py for the random dataset.
254 + DEFAULT_PREFIX_LEN = 0
255 + DEFAULT_RANGE_RATIO = 1.0
256 + DEFAULT_INPUT_LEN = 1024
257 + DEFAULT_OUTPUT_LEN = 128
258 +
259 + def __init__(
260 + self,
261 + **kwargs,
262 + ) -> None:
263 + super().__init__(**kwargs)
264 +
265 + def sample(self,
266 + tokenizer: PreTrainedTokenizerBase,
267 + num_requests: int,
268 + prefix_len: int = DEFAULT_PREFIX_LEN,
269 + range_ratio: float = DEFAULT_RANGE_RATIO,
270 + input_len: int = DEFAULT_INPUT_LEN,
271 + output_len: int = DEFAULT_OUTPUT_LEN,
272 + **kwargs) -> list[SampleRequest]:
273 +
274 + vocab_size = tokenizer.vocab_size
275 +
276 + prefix_token_ids = (np.random.randint(
277 + 0, vocab_size, size=prefix_len).tolist() if prefix_len > 0 else [])
278 +
279 + input_low = int(input_len * range_ratio)
280 + output_low = int(output_len * range_ratio)
281 +
282 + input_lens = np.random.randint(input_low,
283 + input_len + 1,
284 + size=num_requests)
285 + output_lens = np.random.randint(output_low,
286 + output_len + 1,
287 + size=num_requests)
288 + offsets = np.random.randint(0, vocab_size, size=num_requests)
289 +
290 + requests = []
291 + for i in range(num_requests):
292 + inner_seq = ((offsets[i] + i + np.arange(input_lens[i])) %
293 + vocab_size).tolist()
294 + token_sequence = prefix_token_ids + inner_seq
295 + prompt = tokenizer.decode(token_sequence)
296 + total_input_len = prefix_len + int(input_lens[i])
297 + requests.append(
298 + SampleRequest(
299 + prompt=prompt,
300 + prompt_len=total_input_len,
301 + expected_output_len=int(output_lens[i]),
302 + ))
303 + return requests
304 +
305 +
306 +# -----------------------------------------------------------------------------
307 +# ShareGPT Dataset Implementation
308 +# -----------------------------------------------------------------------------
309 +
310 +
311 +class ShareGPTDataset(BenchmarkDataset):
312 + """
313 + Implements the ShareGPT dataset. Loads data from a JSON file and generates
314 + sample requests based on conversation turns.
315 + """
316 +
317 + def __init__(self, **kwargs) -> None:
318 + super().__init__(**kwargs)
319 + self.load_data()
320 +
321 + def load_data(self) -> None:
322 + if self.dataset_path is None:
323 + raise ValueError("dataset_path must be provided for loading data.")
324 +
325 + with open(self.dataset_path, encoding="utf-8") as f:
326 + self.data = json.load(f)
327 + # Filter entries with at least two conversation turns.
328 + self.data = [
329 + entry for entry in self.data
330 + if "conversations" in entry and len(entry["conversations"]) >= 2
331 + ]
332 + random.seed(self.random_seed)
333 + random.shuffle(self.data)
334 +
335 + def sample(self,
336 + tokenizer: PreTrainedTokenizerBase,
337 + num_requests: int,
338 + lora_path: Optional[str] = None,
339 + max_loras: Optional[int] = None,
340 + output_len: Optional[int] = None,
341 + **kwargs) -> list:
342 + samples: list = []
343 + for entry in self.data:
344 + if len(samples) >= num_requests:
345 + break
346 + prompt, completion = entry["conversations"][0]["value"],\
347 + entry["conversations"][1]["value"]
348 +
349 + lora_request, tokenizer = self.get_random_lora_request(
350 + tokenizer=tokenizer, max_loras=max_loras, lora_path=lora_path)
351 + prompt_ids = tokenizer(prompt).input_ids
352 + completion_ids = tokenizer(completion).input_ids
353 + prompt_len = len(prompt_ids)
354 + new_output_len = (len(completion_ids)
355 + if output_len is None else output_len)
356 + if not is_valid_sequence(prompt_len,
357 + new_output_len,
358 + skip_min_output_len_check=output_len
359 + is not None):
360 + continue
361 + samples.append(
362 + SampleRequest(
363 + prompt=prompt,
364 + prompt_len=prompt_len,
365 + expected_output_len=new_output_len,
366 + lora_request=lora_request,
367 + ))
368 + return samples
369 +
370 +
371 +# -----------------------------------------------------------------------------
372 +# Sonnet Dataset Implementation
373 +# -----------------------------------------------------------------------------
374 +
375 +
376 +class SonnetDataset(BenchmarkDataset):
377 + """
378 + Simplified implementation of the Sonnet dataset. Loads poem lines from a
379 + text file and generates sample requests. Default values here copied from
380 + `benchmark_serving.py` for the sonnet dataset.
381 + """
382 +
383 + DEFAULT_PREFIX_LEN = 200
384 + DEFAULT_INPUT_LEN = 550
385 + DEFAULT_OUTPUT_LEN = 150
386 +
387 + def __init__(
388 + self,
389 + **kwargs,
390 + ) -> None:
391 + super().__init__(**kwargs)
392 + self.load_data()
393 +
394 + def load_data(self) -> None:
395 + if not self.dataset_path:
396 + raise ValueError("dataset_path must be provided.")
397 + with open(self.dataset_path, encoding="utf-8") as f:
398 + self.data = f.readlines()
399 +
400 + def sample(self,
401 + tokenizer,
402 + num_requests: int,
403 + prefix_len: int = DEFAULT_PREFIX_LEN,
404 + input_len: int = DEFAULT_INPUT_LEN,
405 + output_len: int = DEFAULT_OUTPUT_LEN,
406 + return_prompt_formatted: bool = False,
407 + **kwargs) -> list:
408 + # Calculate average token length for a poem line.
409 + tokenized_lines = [tokenizer(line).input_ids for line in self.data]
410 + avg_len = sum(len(tokens)
411 + for tokens in \
412 + tokenized_lines) / len(tokenized_lines)
413 +
414 + # Build the base prompt.
415 + base_prompt = "Pick as many lines as you can from these poem lines:\n"
416 + base_msg = [{"role": "user", "content": base_prompt}]
417 + base_fmt = tokenizer.apply_chat_template(base_msg,
418 + add_generation_prompt=True,
419 + tokenize=False)
420 + base_offset = len(tokenizer(base_fmt).input_ids)
421 + if input_len <= base_offset:
422 + raise ValueError(
423 + f"'input_len' must be higher than the base prompt length "
424 + f"({base_offset}).")
425 +
426 + # Determine how many poem lines to use.
427 + num_input_lines = round((input_len - base_offset) / avg_len)
428 + num_prefix_lines = round((prefix_len - base_offset) / avg_len)
429 + prefix_lines = self.data[:num_prefix_lines]
430 +
431 + samples = []
432 + for _ in range(num_requests):
433 + extra_lines = random.choices(self.data,
434 + k=num_input_lines - num_prefix_lines)
435 + prompt = f"{base_prompt}{''.join(prefix_lines + extra_lines)}"
436 + msg = [{"role": "user", "content": prompt}]
437 + prompt_formatted = tokenizer.apply_chat_template(
438 + msg, add_generation_prompt=True, tokenize=False)
439 + prompt_len = len(tokenizer(prompt_formatted).input_ids)
440 + samples.append(
441 + SampleRequest(
442 + prompt=prompt_formatted
443 + if return_prompt_formatted else prompt,
444 + prompt_len=prompt_len,
445 + expected_output_len=output_len,
446 + ))
447 + return samples
448 +
449 +
450 +# -----------------------------------------------------------------------------
451 +# BurstGPT Dataset Implementation
452 +# -----------------------------------------------------------------------------
453 +
454 +
455 +class BurstGPTDataset(BenchmarkDataset):
456 + """
457 + Implements the BurstGPT dataset. Loads data from a CSV file and generates
458 + sample requests based on synthetic prompt generation. Only rows with Model
459 + "GPT-4" and positive response tokens are used.
460 + """
461 +
462 + def __init__(self, **kwargs) -> None:
463 + super().__init__(**kwargs)
464 + self.load_data()
465 +
466 + def load_data(self, ):
467 + if self.dataset_path is None:
468 + raise ValueError("dataset_path must be provided for loading data.")
469 +
470 + df = pd.read_csv(self.dataset_path)
471 + # Filter to keep only GPT-4 rows.
472 + gpt4_df = df[df["Model"] == "GPT-4"]
473 + # Remove failed requests (where Response tokens is 0 or less).
474 + gpt4_df = gpt4_df[gpt4_df["Response tokens"] > 0]
475 + # Sample the desired number of rows.
476 + self.data = gpt4_df
477 +
478 + def _sample_loaded_data(self, num_requests: int) -> list:
479 + if num_requests <= len(self.data):
480 + data = self.data.sample(n=num_requests,
481 + random_state=self.random_seed)
482 + else:
483 + data = self.data.sample(
484 + n=num_requests,
485 + random_state=self.random_seed,
486 + replace=True,
487 + )
488 + # Convert the dataframe to a list of lists.
489 + return data.values.tolist()
490 +
491 + def sample(self,
492 + tokenizer: PreTrainedTokenizerBase,
493 + num_requests: int,
494 + max_loras: Optional[int] = None,
495 + lora_path: Optional[str] = None,
496 + **kwargs) -> list[SampleRequest]:
497 + samples = []
498 + data = self._sample_loaded_data(num_requests=num_requests)
499 + for i in range(num_requests):
500 + input_len = int(data[i][2])
501 + output_len = int(data[i][3])
502 + lora_req, tokenizer = self.get_random_lora_request(
503 + tokenizer=tokenizer, max_loras=max_loras, lora_path=lora_path)
504 + vocab_size = tokenizer.vocab_size
505 + # Generate a synthetic prompt: a list of token IDs computed as (i +
506 + # j) modulo vocab_size.
507 + token_ids = [(i + j) % vocab_size for j in range(input_len)]
508 + prompt = tokenizer.decode(token_ids)
509 + samples.append(
510 + SampleRequest(
511 + prompt=prompt,
512 + prompt_len=input_len,
513 + expected_output_len=output_len,
514 + lora_request=lora_req,
515 + ))
516 + return samples
517 +
518 +
519 +# -----------------------------------------------------------------------------
520 +# HuggingFace Dataset Implementation
521 +# -----------------------------------------------------------------------------
522 +
523 +
524 +class HuggingFaceDataset(BenchmarkDataset):
525 + """
526 + Dataset class for processing a HuggingFace dataset with conversation data
527 + and optional images.
528 + """
529 + DEFAULT_NUM_REQUESTS = 1000
530 +
531 + def __init__(
532 + self,
533 + dataset_split: str,
534 + dataset_subset: Optional[str] = None,
535 + **kwargs,
536 + ) -> None:
537 + super().__init__(**kwargs)
538 + self.dataset_split = dataset_split
539 + self.dataset_subset = dataset_subset
540 +
541 + self.load_data()
542 +
543 + def load_data(self) -> None:
544 + if not self.dataset_path:
545 + raise ValueError("dataset_path must be provided for loading data.")
546 +
547 + self.data = load_dataset(
548 + self.dataset_path,
549 + name=self.dataset_subset,
550 + split=self.dataset_split,
551 + streaming=True,
552 + )
553 +
554 + if "conversations" not in self.data.features:
555 + raise ValueError("HF Dataset must have a 'conversations' column.")
556 +
557 + # Shuffle and filter examples with at least 2 conversations.
558 + self.data = self.data.shuffle(seed=self.random_seed).filter(
559 + lambda x: len(x["conversations"]) >= 2)
560 +
561 + def sample(self,
562 + tokenizer: PreTrainedTokenizerBase,
563 + num_requests: int,
564 + lora_path: Optional[str] = None,
565 + max_loras: Optional[int] = None,
566 + output_len: Optional[int] = None,
567 + **kwargs) -> list:
568 + sampled_requests = []
569 + dynamic_output = output_len is None
570 +
571 + for item in self.data:
572 + if len(sampled_requests) >= num_requests:
573 + break
574 +
575 + conv = item["conversations"]
576 + prompt, completion = conv[0]["value"], conv[1]["value"]
577 +
578 + lora_request, tokenizer = self.get_random_lora_request(
579 + tokenizer, lora_path=lora_path, max_loras=max_loras)
580 +
581 + prompt_ids = tokenizer(prompt).input_ids
582 + completion_ids = tokenizer(completion).input_ids
583 + prompt_len = len(prompt_ids)
584 + completion_len = len(completion_ids)
585 + output_len = completion_len if dynamic_output else output_len
586 + assert isinstance(output_len, int) and output_len > 0
587 + if dynamic_output and not is_valid_sequence(
588 + prompt_len, completion_len):
589 + continue
590 +
591 + mm_content = process_image(
592 + item["image"]) if "image" in item else None
593 + sampled_requests.append(
594 + SampleRequest(
595 + prompt=prompt,
596 + prompt_len=prompt_len,
597 + expected_output_len=output_len,
598 + multi_modal_data=mm_content,
599 + lora_request=lora_request,
600 + ))
601 + return sampled_requests
602 +
603 +
604 +# -----------------------------------------------------------------------------
605 +# Vision Arena Dataset Implementation
606 +# -----------------------------------------------------------------------------
607 +
608 +
609 +class VisionArenaDataset(BenchmarkDataset):
610 + """
611 + Vision Arena Dataset.
612 + """
613 +
614 + DEFAULT_OUTPUT_LEN = 128
615 + DEFAULT_NUM_REQUESTS = 1000
616 + VISION_ARENA_DATASET_PATH = "lmarena-ai/vision-arena-bench-v0.1"
617 +
618 + def __init__(
619 + self,
620 + dataset_split: str,
621 + dataset_subset: Optional[str] = None,
622 + **kwargs,
623 + ) -> None:
624 + super().__init__(**kwargs)
625 + self.dataset_split = dataset_split
626 + self.dataset_subset = dataset_subset
627 +
628 + if self.dataset_path != self.VISION_ARENA_DATASET_PATH:
629 + raise ValueError(f"Only support Vision Arena dataset.\
630 + This data path {self.dataset_path} is not valid.")
631 + if self.dataset_subset is None and self.dataset_split != "train":
632 + raise ValueError("Dataset split must be 'train'.")
633 +
634 + self.load_data()
635 +
636 + def load_data(self) -> None:
637 + dataset = load_dataset(
638 + self.dataset_path,
639 + name=self.dataset_subset,
640 + split=self.dataset_split,
641 + streaming=True,
642 + )
643 + self.data = dataset.shuffle(seed=self.random_seed)
644 +
645 + def sample(self,
646 + tokenizer: PreTrainedTokenizerBase,
647 + num_requests: int,
648 + output_len: int = DEFAULT_OUTPUT_LEN,
649 + **kwargs) -> list:
650 + # TODO (jenniferzhao): Add support for offline benchmark sampling
651 + output_len = (output_len
652 + if output_len is not None else self.DEFAULT_OUTPUT_LEN)
653 + sampled_requests = []
654 + for item in self.data:
655 + if len(sampled_requests) >= num_requests:
656 + break
657 + prompt = item["turns"][0][0]["content"]
658 + prompt_len = len(tokenizer(prompt).input_ids)
659 + mm_content = process_image(item["images"][0])
660 + sampled_requests.append(
661 + SampleRequest(
662 + prompt=prompt,
663 + prompt_len=prompt_len,
664 + expected_output_len=output_len,
665 + multi_modal_data=mm_content,
666 + ))
667 + return sampled_requests
benchmarks/benchmark_serving.py
+77 -382
@@ -25,25 +25,20 @@ On the client side, run:
25 """
26 import argparse
27 import asyncio
28 -import base64
28 import gc
30 -import io
29 import json
30 import os
31 import random
32 import time
33 import warnings
36 -from collections.abc import AsyncGenerator, Collection
34 +from collections.abc import AsyncGenerator, Iterable
35 from dataclasses import dataclass
36 from datetime import datetime
37 from typing import Any, Optional
38
39 import numpy as np
42 -import pandas as pd
40 from backend_request_func import (ASYNC_REQUEST_FUNCS, RequestFuncInput,
41 RequestFuncOutput)
45 -from datasets import load_dataset
46 -from PIL.Image import Image
42 from tqdm.asyncio import tqdm
43 from transformers import PreTrainedTokenizerBase
44
@@ -57,6 +52,9 @@ try:
52 except ImportError:
53 from argparse import ArgumentParser as FlexibleArgumentParser
54
55 +from benchmark_dataset import (BurstGPTDataset, HuggingFaceDataset,
56 + RandomDataset, SampleRequest, ShareGPTDataset,
57 + SonnetDataset, VisionArenaDataset)
58 from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json
59
60 MILLISECONDS_TO_SECONDS_CONVERSION = 1000
@@ -92,325 +90,18 @@ class BenchmarkMetrics:
90 percentiles_e2el_ms: list[tuple[float, float]]
91
92
95 -def sample_sharegpt_requests(
96 - dataset_path: str,
97 - num_requests: int,
98 - tokenizer: PreTrainedTokenizerBase,
99 - fixed_output_len: Optional[int] = None,
100 -) -> list[tuple[str, int, int, None]]:
101 - # Load the dataset.
102 - with open(dataset_path, encoding='utf-8') as f:
103 - dataset = json.load(f)
104 - # Filter out the conversations with less than 2 turns.
105 - dataset = [data for data in dataset if len(data["conversations"]) >= 2]
106 - # Only keep the first two turns of each conversation.
107 - dataset = [(data["conversations"][0]["value"],
108 - data["conversations"][1]["value"]) for data in dataset]
109 -
110 - # Shuffle the dataset.
111 - random.shuffle(dataset)
112 -
113 - # Filter out sequences that are too long or too short
114 - filtered_dataset: list[tuple[str, int, int]] = []
115 - for i in range(len(dataset)):
116 - if len(filtered_dataset) == num_requests:
117 - break
118 -
119 - # Tokenize the prompts and completions.
120 - prompt = dataset[i][0]
121 - prompt_token_ids = tokenizer(prompt).input_ids
122 - completion = dataset[i][1]
123 - completion_token_ids = tokenizer(completion).input_ids
124 - prompt_len = len(prompt_token_ids)
125 - output_len = len(completion_token_ids
126 - ) if fixed_output_len is None else fixed_output_len
127 - if prompt_len < 4 or (fixed_output_len is None and output_len < 4):
128 - # Prune too short sequences.
129 - continue
130 - if prompt_len > 1024 or prompt_len + output_len > 2048:
131 - # Prune too long sequences.
132 - continue
133 - filtered_dataset.append((prompt, prompt_len, output_len, None))
134 -
135 - return filtered_dataset
136 -
137 -
138 -def sample_burstgpt_requests(
139 - dataset_path: str,
140 - num_requests: int,
141 - random_seed: int,
142 - tokenizer: PreTrainedTokenizerBase,
143 -) -> list[tuple[str, int, int, None]]:
144 - df = pd.read_csv(dataset_path)
145 - gpt4_df = df[df["Model"] == "GPT-4"]
146 - # Remove the failed requests (i.e., response length is 0)
147 - gpt4_df = gpt4_df[gpt4_df["Response tokens"] > 0]
148 - # Randomly sample num_requests from the dataset
149 - if num_requests <= len(gpt4_df):
150 - gpt4_df = gpt4_df.sample(n=num_requests, random_state=random_seed)
151 - else:
152 - gpt4_df = gpt4_df.sample(n=num_requests,
153 - random_state=random_seed,
154 - replace=True)
155 - # Convert the dataframe to a list of tuples
156 - dataset = gpt4_df.values.tolist()
157 - input_requests = []
158 - for i in range(num_requests):
159 - input_len = int(dataset[i][2])
160 - output_len = int(dataset[i][3])
161 - prompt = tokenizer.decode([(i + j) % tokenizer.vocab_size
162 - for j in range(input_len)])
163 - input_requests.append((prompt, input_len, output_len, None))
164 - return input_requests
165 -
166 -
167 -def sample_sonnet_requests(
168 - dataset_path: str,
169 - num_requests: int,
170 - input_len: int,
171 - output_len: int,
172 - prefix_len: int,
173 - tokenizer: PreTrainedTokenizerBase,
174 -) -> list[tuple[str, str, int, int, None]]:
175 - assert (
176 - input_len > prefix_len
177 - ), "'args.sonnet-input-len' must be greater than 'args.sonnet-prefix-len'."
178 -
179 - # Load the dataset.
180 - with open(dataset_path, encoding='utf-8') as f:
181 - poem_lines = f.readlines()
182 -
183 - # Tokenize the poem lines.
184 - poem_token_ids = tokenizer(poem_lines).input_ids
185 - average_poem_len = sum(
186 - len(token_ids) for token_ids in poem_token_ids) / len(poem_token_ids)
187 -
188 - # Base prefix for all requests.
189 - base_prompt = "Pick as many lines as you can from these poem lines:\n"
190 - base_message = [{
191 - "role": "user",
192 - "content": base_prompt,
193 - }]
194 - base_prompt_formatted = tokenizer.apply_chat_template(
195 - base_message, add_generation_prompt=True, tokenize=False)
196 - base_prompt_offset = len(tokenizer(base_prompt_formatted).input_ids)
197 -
198 - assert (
199 - input_len > base_prompt_offset
200 - ), f"Please set 'args.sonnet-input-len' higher than {base_prompt_offset}."
201 - num_input_lines = round(
202 - (input_len - base_prompt_offset) / average_poem_len)
203 -
204 - # First approximately `prefix_len` number of tokens in the
205 - # prompt are fixed poem lines.
206 - assert (
207 - prefix_len > base_prompt_offset
208 - ), f"Please set 'args.sonnet-prefix-len' higher than {base_prompt_offset}."
209 -
210 - num_prefix_lines = round(
211 - (prefix_len - base_prompt_offset) / average_poem_len)
212 - prefix_lines = poem_lines[:num_prefix_lines]
213 -
214 - # Sample the rest of lines per request.
215 - sampled_requests: list[tuple[str, int, int]] = []
216 - for _ in range(num_requests):
217 - num_lines_needed = num_input_lines - num_prefix_lines
218 - sampled_lines = "".join(prefix_lines +
219 - random.choices(poem_lines, k=num_lines_needed))
220 -
221 - prompt = f"{base_prompt}{sampled_lines}"
222 - message = [
223 - {
224 - "role": "user",
225 - "content": prompt,
226 - },
227 - ]
228 - prompt_formatted = tokenizer.apply_chat_template(
229 - message, add_generation_prompt=True, tokenize=False)
230 - prompt_len = len(tokenizer(prompt_formatted).input_ids)
231 - sampled_requests.append(
232 - (prompt, prompt_formatted, prompt_len, output_len, None))
233 -
234 - return sampled_requests
235 -
236 -
237 -def sample_vision_arena_requests(
238 - dataset,
239 - num_requests: int,
240 - tokenizer: PreTrainedTokenizerBase,
241 - fixed_output_len: Optional[int] = None,
242 -) -> list[tuple[str, str, int, Optional[dict[str, Collection[str]]]]]:
243 - sampled_requests: list[tuple[str, int, int, dict[str,
244 - Collection[str]]]] = []
245 - for data in dataset:
246 - if len(sampled_requests) == num_requests:
247 - break
248 -
249 - prompt = data["turns"][0][0]['content']
250 -
251 - prompt_token_ids = tokenizer(prompt).input_ids
252 - if fixed_output_len is None:
253 - # Default max output len is set to 128
254 - print("--hf-output-len is not provided. Using default value 128.")
255 - fixed_output_len = 128
256 -
257 - prompt_len = len(prompt_token_ids)
258 - output_len = fixed_output_len
259 -
260 - assert isinstance(
261 - data["images"][0],
262 - Image), ("Input image format must be `PIL.Image.Image`, "
263 - f"given {type(data['image'])}.")
264 - image: Image = data["images"][0]
265 - image = image.convert("RGB")
266 - image_data = io.BytesIO()
267 - image.save(image_data, format='JPEG')
268 - image_base64 = base64.b64encode(image_data.getvalue()).decode("utf-8")
269 - mm_content = {
270 - "type": "image_url",
271 - "image_url": {
272 - "url": f"data:image/jpeg;base64,{image_base64}"
273 - },
274 - }
275 -
276 - sampled_requests.append((prompt, prompt_len, output_len, mm_content))
277 -
278 - return sampled_requests
279 -
280 -
281 -def sample_hf_requests(
282 - dataset_path: str,
283 - dataset_subset: Optional[str],
284 - dataset_split: str,
285 - num_requests: int,
286 - tokenizer: PreTrainedTokenizerBase,
287 - random_seed: int,
288 - fixed_output_len: Optional[int] = None,
289 -) -> list[tuple[str, str, int, Optional[dict[str, Collection[str]]]]]:
290 -
291 - # Special case for vision_arena dataset
292 - if dataset_path == 'lmarena-ai/vision-arena-bench-v0.1' \
293 - and dataset_subset is None:
294 - assert dataset_split == "train"
295 - dataset = load_dataset(dataset_path,
296 - name=dataset_subset,
297 - split=dataset_split,
298 - streaming=True)
299 - dataset = dataset.shuffle(seed=random_seed)
300 - return sample_vision_arena_requests(dataset, num_requests, tokenizer,
301 - fixed_output_len)
302 -
303 - dataset = load_dataset(dataset_path,
304 - name=dataset_subset,
305 - split=dataset_split,
306 - streaming=True)
307 - assert "conversations" in dataset.features, (
308 - "HF Dataset must have 'conversations' column.")
309 - filter_func = lambda x: len(x["conversations"]) >= 2
310 - filtered_dataset = dataset.shuffle(seed=random_seed).filter(filter_func)
311 - sampled_requests: list[tuple[str, int, int, dict[str,
312 - Collection[str]]]] = []
313 - for data in filtered_dataset:
314 - if len(sampled_requests) == num_requests:
315 - break
316 -
317 - # Tokenize the prompts and completions.
318 - prompt = data["conversations"][0]["value"]
319 - prompt_token_ids = tokenizer(prompt).input_ids
320 - completion = data["conversations"][1]["value"]
321 - completion_token_ids = tokenizer(completion).input_ids
322 - prompt_len = len(prompt_token_ids)
323 - output_len = len(completion_token_ids
324 - ) if fixed_output_len is None else fixed_output_len
325 - if fixed_output_len is None and (prompt_len < 4 or output_len < 4):
326 - # Prune too short sequences.
327 - continue
328 - if fixed_output_len is None and \
329 - (prompt_len > 1024 or prompt_len + output_len > 2048):
330 - # Prune too long sequences.
331 - continue
332 -
333 - if "image" in data and isinstance(data["image"], Image):
334 - image: Image = data["image"]
335 - image = image.convert("RGB")
336 - image_data = io.BytesIO()
337 - image.save(image_data, format='JPEG')
338 - image_base64 = base64.b64encode(
339 - image_data.getvalue()).decode("utf-8")
340 - mm_content = {
341 - "type": "image_url",
342 - "image_url": {
343 - "url": f"data:image/jpeg;base64,{image_base64}"
344 - },
345 - }
346 - elif "image" in data and isinstance(data["image"], str):
347 - if (data["image"].startswith("http://") or \
348 - data["image"].startswith("file://")):
349 - image_url = data["image"]
350 - else:
351 - image_url = f"file://{data['image']}"
352 -
353 - mm_content = {
354 - "type": "image_url",
355 - "image_url": {
356 - "url": image_url
357 - },
358 - }
359 - else:
360 - mm_content = None
361 -
362 - sampled_requests.append((prompt, prompt_len, output_len, mm_content))
363 -
364 - return sampled_requests
365 -
366 -
367 -def sample_random_requests(
368 - prefix_len: int,
369 - input_len: int,
370 - output_len: int,
371 - num_prompts: int,
372 - range_ratio: float,
373 - tokenizer: PreTrainedTokenizerBase,
374 -) -> list[tuple[str, int, int]]:
375 - prefix_token_ids = np.random.randint(0,
376 - tokenizer.vocab_size,
377 - size=prefix_len).tolist()
378 -
379 - input_lens = np.random.randint(
380 - int(input_len * range_ratio),
381 - input_len + 1,
382 - size=num_prompts,
383 - )
384 - output_lens = np.random.randint(
385 - int(output_len * range_ratio),
386 - output_len + 1,
387 - size=num_prompts,
388 - )
389 - offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts)
390 - input_requests = []
391 - for i in range(num_prompts):
392 - prompt = tokenizer.decode(prefix_token_ids +
393 - [(offsets[i] + i + j) % tokenizer.vocab_size
394 - for j in range(input_lens[i])])
395 -
396 - input_requests.append((prompt, int(prefix_len + input_lens[i]),
397 - int(output_lens[i]), None))
398 -
399 - return input_requests
400 -
401 -
93 async def get_request(
403 - input_requests: list[tuple[str, int, int]],
94 + input_requests: list[SampleRequest],
95 request_rate: float,
96 burstiness: float = 1.0,
406 -) -> AsyncGenerator[tuple[str, int, int], None]:
97 +) -> AsyncGenerator[SampleRequest, None]:
98 """
99 Asynchronously generates requests at a specified rate
100 with OPTIONAL burstiness.
101
102 Args:
103 input_requests:
413 - A list of input requests, each represented as a tuple.
104 + A list of input requests, each represented as a SampleRequest.
105 request_rate:
106 The rate at which requests are generated (requests/s).
107 burstiness (optional):
@@ -422,7 +113,7 @@ async def get_request(
113 in more bursty requests, while a higher burstiness value
114 (burstiness > 1) results in a more uniform arrival of requests.
115 """
425 - input_requests = iter(input_requests)
116 + input_requests: Iterable[SampleRequest] = iter(input_requests)
117
118 # Calculate scale parameter theta to maintain the desired request_rate.
119 assert burstiness > 0, (
@@ -444,7 +135,7 @@ async def get_request(
135
136
137 def calculate_metrics(
447 - input_requests: list[tuple[str, int, int]],
138 + input_requests: list[SampleRequest],
139 outputs: list[RequestFuncOutput],
140 dur_s: float,
141 tokenizer: PreTrainedTokenizerBase,
@@ -475,7 +166,7 @@ def calculate_metrics(
166 tokenizer(outputs[i].generated_text,
167 add_special_tokens=False).input_ids)
168 actual_output_lens.append(output_len)
478 - total_input += input_requests[i][1]
169 + total_input += input_requests[i].prompt_len
170 tpot = 0
171 if output_len > 1:
172 latency_minus_ttft = outputs[i].latency - outputs[i].ttft
@@ -558,18 +249,18 @@ async def benchmark(
249 model_id: str,
250 model_name: str,
251 tokenizer: PreTrainedTokenizerBase,
561 - input_requests: list[tuple[str, int, int]],
252 + input_requests: list[SampleRequest],
253 logprobs: Optional[int],
254 request_rate: float,
255 burstiness: float,
256 disable_tqdm: bool,
257 profile: bool,
258 selected_percentile_metrics: list[str],
568 - selected_percentiles: list[str],
259 + selected_percentiles: list[float],
260 ignore_eos: bool,
261 goodput_config_dict: dict[str, float],
262 max_concurrency: Optional[int],
572 - lora_modules: Optional[list[str]],
263 + lora_modules: Optional[Iterable[str]],
264 ):
265 if backend in ASYNC_REQUEST_FUNCS:
266 request_func = ASYNC_REQUEST_FUNCS[backend]
@@ -577,12 +268,16 @@ async def benchmark(
268 raise ValueError(f"Unknown backend: {backend}")
269
270 print("Starting initial single prompt test run...")
580 - test_prompt, test_prompt_len, test_output_len, test_mm_content = (
581 - input_requests[0])
271 + test_prompt, test_prompt_len, test_output_len, test_mm_content = \
272 + input_requests[0].prompt, input_requests[0].prompt_len, \
273 + input_requests[0].expected_output_len, \
274 + input_requests[0].multi_modal_data
275 +
276 if backend != "openai-chat" and test_mm_content is not None:
277 # multi-modal benchmark is only available on OpenAI Chat backend.
278 raise ValueError(
279 "Multi-modal content is only supported on 'openai-chat' backend.")
280 + assert test_mm_content is None or isinstance(test_mm_content, dict)
281 test_input = RequestFuncInput(
282 model=model_id,
283 model_name=model_name,
@@ -606,7 +301,8 @@ async def benchmark(
301 if lora_modules:
302 # For each input request, choose a LoRA module at random.
303 lora_modules = iter(
609 - [random.choice(lora_modules) for _ in range(len(input_requests))])
304 + [random.choice(lora_modules) \
305 + for _ in range(len(input_requests))])
306
307 if profile:
308 print("Starting profiler...")
@@ -652,7 +348,9 @@ async def benchmark(
348 benchmark_start_time = time.perf_counter()
349 tasks: list[asyncio.Task] = []
350 async for request in get_request(input_requests, request_rate, burstiness):
655 - prompt, prompt_len, output_len, mm_content = request
351 + prompt, prompt_len, output_len, mm_content = request.prompt, \
352 + request.prompt_len, request.expected_output_len, \
353 + request.multi_modal_data
354 req_model_id, req_model_name = model_id, model_name
355 if lora_modules:
356 req_lora_module = next(lora_modules)
@@ -867,76 +565,72 @@ def main(args: argparse.Namespace):
565 "Please specify '--dataset-name' and the corresponding "
566 "'--dataset-path' if required.")
567
870 - elif args.dataset_name == "sharegpt":
871 - input_requests = sample_sharegpt_requests(
872 - dataset_path=args.dataset_path,
873 - num_requests=args.num_prompts,
874 - tokenizer=tokenizer,
875 - fixed_output_len=args.sharegpt_output_len,
876 - )
877 -
878 - elif args.dataset_name == "burstgpt":
879 - input_requests = sample_burstgpt_requests(
880 - dataset_path=args.dataset_path,
881 - num_requests=args.num_prompts,
882 - random_seed=args.seed,
883 - tokenizer=tokenizer,
884 - )
885 -
886 - elif args.dataset_name == "sonnet":
887 - # Do not format the prompt, pass to message directly
568 + if args.dataset_name == "sonnet":
569 + dataset = SonnetDataset(dataset_path=args.dataset_path)
570 + # For the "sonnet" dataset, formatting depends on the backend.
571 if args.backend == "openai-chat":
889 - input_requests = sample_sonnet_requests(
890 - dataset_path=args.dataset_path,
891 - num_requests=args.num_prompts,
892 - input_len=args.sonnet_input_len,
893 - output_len=args.sonnet_output_len,
894 - prefix_len=args.sonnet_prefix_len,
895 - tokenizer=tokenizer,
896 - )
897 - input_requests = [(prompt, prompt_len, output_len, None)
898 - for prompt, prompt_formatted, prompt_len,
899 - output_len, _ in input_requests]
572 + input_requests = dataset.sample(num_requests=args.num_prompts,
573 + input_len=args.sonnet_input_len,
574 + output_len=args.sonnet_output_len,
575 + prefix_len=args.sonnet_prefix_len,
576 + tokenizer=tokenizer,
577 + return_prompt_formatted=False)
578 else:
901 - assert (
902 - tokenizer.chat_template or tokenizer.default_chat_template
903 - ), "Tokenizer/model must have chat template for sonnet dataset."
904 - input_requests = sample_sonnet_requests(
905 - dataset_path=args.dataset_path,
906 - num_requests=args.num_prompts,
907 - input_len=args.sonnet_input_len,
908 - output_len=args.sonnet_output_len,
909 - prefix_len=args.sonnet_prefix_len,
910 - tokenizer=tokenizer,
911 - )
912 - input_requests = [(prompt_formatted, prompt_len, output_len, None)
913 - for prompt, prompt_formatted, prompt_len,
914 - output_len, _ in input_requests]
579 + assert tokenizer.chat_template or tokenizer.default_chat_template, (
580 + "Tokenizer/model must have chat template for sonnet dataset.")
581 + input_requests = dataset.sample(num_requests=args.num_prompts,
582 + input_len=args.sonnet_input_len,
583 + output_len=args.sonnet_output_len,
584 + prefix_len=args.sonnet_prefix_len,
585 + tokenizer=tokenizer,
586 + return_prompt_formatted=True)
587
588 elif args.dataset_name == "hf":
917 - input_requests = sample_hf_requests(
589 + # Choose between VisionArenaDataset
590 + # and HuggingFaceDataset based on provided parameters.
591 + dataset_class = (VisionArenaDataset if args.dataset_path
592 + == VisionArenaDataset.VISION_ARENA_DATASET_PATH
593 + and args.hf_subset is None else HuggingFaceDataset)
594 + input_requests = dataset_class(
595 dataset_path=args.dataset_path,
596 dataset_subset=args.hf_subset,
597 dataset_split=args.hf_split,
598 + ).sample(
599 num_requests=args.num_prompts,
600 tokenizer=tokenizer,
601 random_seed=args.seed,
924 - fixed_output_len=args.hf_output_len,
925 - )
926 -
927 - elif args.dataset_name == "random":
928 - input_requests = sample_random_requests(
929 - prefix_len=args.random_prefix_len,
930 - input_len=args.random_input_len,
931 - output_len=args.random_output_len,
932 - num_prompts=args.num_prompts,
933 - range_ratio=args.random_range_ratio,
934 - tokenizer=tokenizer,
602 + output_len=args.hf_output_len,
603 )
604
605 else:
938 - raise ValueError(f"Unknown dataset: {args.dataset_name}")
606 + # For datasets that follow a similar structure, use a mapping.
607 + dataset_mapping = {
608 + "sharegpt":
609 + lambda: ShareGPTDataset(random_seed=args.seed,
610 + dataset_path=args.dataset_path).sample(
611 + tokenizer=tokenizer,
612 + num_requests=args.num_prompts,
613 + output_len=args.sharegpt_output_len,
614 + ),
615 + "burstgpt":
616 + lambda: BurstGPTDataset(random_seed=args.seed,
617 + dataset_path=args.dataset_path).
618 + sample(tokenizer=tokenizer, num_requests=args.num_prompts),
619 + "random":
620 + lambda: RandomDataset(dataset_path=args.dataset_path).sample(
621 + tokenizer=tokenizer,
622 + num_requests=args.num_prompts,
623 + prefix_len=args.random_prefix_len,
624 + input_len=args.random_input_len,
625 + output_len=args.random_output_len,
626 + range_ratio=args.random_range_ratio,
627 + )
628 + }
629
630 + try:
631 + input_requests = dataset_mapping[args.dataset_name]()
632 + except KeyError as err:
633 + raise ValueError(f"Unknown dataset: {args.dataset_name}") from err
634 goodput_config_dict = check_goodput_args(args)
635
636 # Avoid GC processing "static" data - reduce pause times.
@@ -1298,4 +992,5 @@ if __name__ == "__main__":
992 "script chooses a LoRA module at random.")
993
994 args = parser.parse_args()
995 +
996 main(args)
benchmarks/benchmark_throughput.py
+81 -197
@@ -6,13 +6,14 @@ import json
6 import os
7 import random
8 import time
9 -from functools import cache
9 +import warnings
10 from typing import Any, Optional, Union
11
12 import torch
13 import uvloop
14 +from benchmark_dataset import (BurstGPTDataset, RandomDataset, SampleRequest,
15 + ShareGPTDataset, SonnetDataset)
16 from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json
15 -from PIL import Image
17 from tqdm import tqdm
18 from transformers import (AutoModelForCausalLM, AutoTokenizer,
19 PreTrainedTokenizerBase)
@@ -22,148 +23,10 @@ from vllm.entrypoints.openai.api_server import (
23 build_async_engine_client_from_engine_args)
24 from vllm.inputs import TextPrompt, TokensPrompt
25 from vllm.lora.request import LoRARequest
25 -from vllm.lora.utils import get_adapter_absolute_path
26 -from vllm.multimodal import MultiModalDataDict
26 from vllm.sampling_params import BeamSearchParams
28 -from vllm.transformers_utils.tokenizer import AnyTokenizer, get_lora_tokenizer
27 from vllm.utils import FlexibleArgumentParser, merge_async_iterators
28
29
32 -@dataclasses.dataclass
33 -class SampleRequest:
34 - """A class representing a single inference request for benchmarking.
35 -
36 - Attributes:
37 - prompt: The input text prompt for the model.
38 - prompt_len: The length of the prompt in tokens.
39 - expected_output_len: The expected length of the output in tokens.
40 - multi_modal_data: Optional dictionary containing multi-modal data (e.g.
41 - images).
42 - lora_request: Optional LoRARequest specifying the LoRA to use.
43 - """
44 - prompt: str
45 - prompt_len: int
46 - expected_output_len: int
47 - multi_modal_data: Optional[MultiModalDataDict] = None
48 - lora_request: Optional[LoRARequest] = None
49 -
50 -
51 -def _get_prompt_for_image_model(question: str, *, model: str) -> str:
52 - """Prepend and append special tokens around the question to form a prompt.
53 -
54 - Args:
55 - question: The input question text to wrap with special tokens
56 - model: The name of the model being used, to determine which special
57 - tokens to add
58 -
59 - Returns:
60 - The formatted prompt string with appropriate special tokens for the
61 - model
62 -
63 - Raises:
64 - ValueError: If an unsupported model name is provided
65 - """
66 - model = model.lower()
67 - if "pixtral" in model:
68 - return f"<s>[INST]{question}\n[IMG][/INST]"
69 - raise ValueError(f"Unsupported model {model}")
70 -
71 -
72 -@cache
73 -def lora_path_on_disk(lora_path: str) -> str:
74 - return get_adapter_absolute_path(lora_path)
75 -
76 -
77 -lora_tokenizer_cache: dict[int, AnyTokenizer] = {}
78 -
79 -
80 -def get_random_lora_request(
81 - args: argparse.Namespace
82 -) -> tuple[LoRARequest, Optional[AnyTokenizer]]:
83 - global lora_tokenizer_cache
84 - lora_id = random.randint(1, args.max_loras)
85 - lora_request = LoRARequest(lora_name=str(lora_id),
86 - lora_int_id=lora_id,
87 - lora_path=lora_path_on_disk(args.lora_path))
88 - if lora_id not in lora_tokenizer_cache:
89 - lora_tokenizer_cache[lora_id] = get_lora_tokenizer(lora_request)
90 - return lora_request, lora_tokenizer_cache[lora_id]
91 -
92 -
93 -def sample_requests(tokenizer: PreTrainedTokenizerBase,
94 - args: argparse.Namespace) -> list[SampleRequest]:
95 -
96 - dataset_path: str = args.dataset
97 - num_requests: int = args.num_prompts
98 - fixed_output_len: Optional[int] = args.output_len
99 - model: str = args.model
100 - if fixed_output_len is not None and fixed_output_len < 4:
101 - raise ValueError("output_len too small")
102 -
103 - # Load the dataset.
104 - with open(dataset_path) as f:
105 - dataset = json.load(f)
106 - # Filter out the conversations with less than 2 turns.
107 - dataset = [data for data in dataset if len(data["conversations"]) >= 2]
108 - # Shuffle the dataset.
109 - random.shuffle(dataset)
110 -
111 - # Filter out sequences that are too long or too short
112 - filtered_dataset: list[SampleRequest] = []
113 - for data in tqdm(dataset,
114 - total=len(filtered_dataset),
115 - desc="sampling requests"):
116 - if len(filtered_dataset) == num_requests:
117 - break
118 -
119 - # Only keep the first two turns of each conversation.
120 - prompt = data["conversations"][0]["value"]
121 - completion = data["conversations"][1]["value"]
122 -
123 - multi_modal_data: Optional[MultiModalDataDict] = None
124 - if "image" in data:
125 - multi_modal_data = multi_modal_data or {}
126 - image_path = data["image"]
127 - # TODO(vllm-project/vllm/issues/9778): Support multiple images.
128 - assert isinstance(image_path,
129 - str), "Only support single image input"
130 - try:
131 - multi_modal_data["image"] = Image.open(image_path).convert(
132 - "RGB")
133 - except FileNotFoundError:
134 - # Ignore datapoint where asset is missing
135 - continue
136 - prompt = _get_prompt_for_image_model(question=prompt, model=model)
137 -
138 - request_tokenizer = tokenizer
139 - lora_request: Optional[LoRARequest] = None
140 - if args.enable_lora:
141 - lora_request, lora_tokenizer = get_random_lora_request(args)
142 - if lora_tokenizer:
143 - request_tokenizer = lora_tokenizer
144 -
145 - # Tokenize the prompts and completions.
146 - prompt_token_ids = request_tokenizer(prompt).input_ids
147 - completion_token_ids = request_tokenizer(completion).input_ids
148 - prompt_len = len(prompt_token_ids)
149 - output_len = len(completion_token_ids
150 - ) if fixed_output_len is None else fixed_output_len
151 - if prompt_len < 4 or output_len < 4:
152 - # Prune too short sequences.
153 - continue
154 - if prompt_len > 1024 or prompt_len + output_len > 2048:
155 - # Prune too long sequences.
156 - continue
157 - filtered_dataset.append(
158 - SampleRequest(prompt=prompt,
159 - prompt_len=prompt_len,
160 - expected_output_len=output_len,
161 - multi_modal_data=multi_modal_data,
162 - lora_request=lora_request))
163 -
164 - return filtered_dataset
165 -
166 -
30 def run_vllm(
31 requests: list[SampleRequest],
32 n: int,
@@ -381,61 +244,50 @@ def save_to_pytorch_benchmark_format(args: argparse.Namespace,
244 write_to_json(pt_file, pt_records)
245
246
247 +def get_requests(args, tokenizer):
248 + # Common parameters for all dataset types.
249 + common_kwargs = {
250 + "dataset_path": args.dataset_path,
251 + "random_seed": args.seed,
252 + }
253 + sample_kwargs = {
254 + "tokenizer": tokenizer,
255 + "lora_path": args.lora_path,
256 + "max_loras": args.max_loras,
257 + "num_requests": args.num_prompts,
258 + "input_len": args.input_len,
259 + "output_len": args.output_len,
260 + }
261 + if args.dataset_path is None or args.dataset_name == "random":
262 + sample_kwargs["range_ratio"] = args.random_range_ratio
263 + sample_kwargs["prefix_len"] = args.prefix_len
264 + dataset_cls = RandomDataset
265 + elif args.dataset_name == "sharegpt":
266 + dataset_cls = ShareGPTDataset
267 + elif args.dataset_name == "sonnet":
268 + assert tokenizer.chat_template or tokenizer.default_chat_template, (
269 + "Tokenizer/model must have chat template for sonnet dataset.")
270 + dataset_cls = SonnetDataset
271 + sample_kwargs["prefix_len"] = args.prefix_len
272 + sample_kwargs["return_prompt_formatted"] = True
273 + elif args.dataset_name == "burstgpt":
274 + dataset_cls = BurstGPTDataset
275 + else:
276 + raise ValueError(f"Unknown dataset name: {args.dataset_name}")
277 + # Remove None values
278 + sample_kwargs = {k: v for k, v in sample_kwargs.items() if v is not None}
279 + return dataset_cls(**common_kwargs).sample(**sample_kwargs)
280 +
281 +
282 def main(args: argparse.Namespace):
283 + if args.seed is None:
284 + args.seed = 0
285 print(args)
286 random.seed(args.seed)
387 -
287 # Sample the requests.
288 tokenizer = AutoTokenizer.from_pretrained(
289 args.tokenizer, trust_remote_code=args.trust_remote_code)
391 - if args.dataset is None:
392 - vocab_size = tokenizer.vocab_size
393 - requests = []
394 - for _ in range(args.num_prompts):
395 -
396 - request_tokenizer = tokenizer
397 - lora_request: Optional[LoRARequest] = None
398 - if args.enable_lora:
399 - lora_request, lora_tokenizer = get_random_lora_request(args)
400 - if lora_tokenizer:
401 - request_tokenizer = lora_tokenizer
402 -
403 - # Synthesize a prompt with the given input length.
404 - candidate_ids = [
405 - random.randint(0, vocab_size - 1)
406 - for _ in range(args.input_len)
407 - ]
408 -
409 - candidate_prompt = {"prompt_token_ids": candidate_ids}
410 -
411 - if not args.skip_tokenizer_init:
412 - # As tokenizer may add additional tokens like BOS, we need
413 - # to try different lengths to get the desired input length.
414 - for _ in range(5): # Max attempts to correct
415 - candidate_prompt = request_tokenizer.decode(candidate_ids)
416 - tokenized_len = len(
417 - request_tokenizer.encode(candidate_prompt))
418 -
419 - if tokenized_len == args.input_len:
420 - break
421 -
422 - # Adjust length based on difference
423 - diff = args.input_len - tokenized_len
424 - if diff > 0:
425 - candidate_ids.extend([
426 - random.randint(100, vocab_size - 100)
427 - for _ in range(diff)
428 - ])
429 - else:
430 - candidate_ids = candidate_ids[:diff]
431 - requests.append(
432 - SampleRequest(prompt=candidate_prompt,
433 - prompt_len=args.input_len,
434 - expected_output_len=args.output_len,
435 - lora_request=lora_request))
436 - else:
437 - requests = sample_requests(tokenizer, args)
438 -
290 + requests = get_requests(args, tokenizer)
291 is_multi_modal = any(request.multi_modal_data is not None
292 for request in requests)
293 if args.backend == "vllm":
@@ -470,7 +322,7 @@ def main(args: argparse.Namespace):
322 print("\033[91mWARNING\033[0m: Multi-modal request detected. The "
323 "following metrics are not accurate because image tokens are not"
324 " counted. See vllm-project/vllm/issues/9778 for details.")
473 - # TODO(vllm-project/vllm/issues/9778): Count molti-modal token length.
325 + # TODO(vllm-project/vllm/issues/9778): Count multi-modal token length.
326 print(f"Throughput: {len(requests) / elapsed_time:.2f} requests/s, "
327 f"{total_num_tokens / elapsed_time:.2f} total tokens/s, "
328 f"{total_output_tokens / elapsed_time:.2f} output tokens/s")
@@ -495,12 +347,23 @@ if __name__ == "__main__":
347 type=str,
348 choices=["vllm", "hf", "mii"],
349 default="vllm")
498 - parser.add_argument("--dataset",
350 + parser.add_argument("--dataset-name",
351 + type=str,
352 + choices=["sharegpt", "random", "sonnet", "burstgpt"],
353 + help="Name of the dataset to benchmark on.",
354 + default="sharegpt")
355 + parser.add_argument(
356 + "--dataset",
357 + type=str,
358 + default=None,
359 + help="Path to the ShareGPT dataset, will be deprecated in\
360 + the next release. The dataset is expected to "
361 + "be a json in form of list[dict[..., conversations: "
362 + "list[dict[..., value: <prompt_or_response>]]]]")
363 + parser.add_argument("--dataset-path",
364 type=str,
365 default=None,
501 - help="Path to the dataset. The dataset is expected to "
502 - "be a json in form of list[dict[..., conversations: "
503 - "list[dict[..., value: <prompt_or_response>]]]]")
366 + help="Path to the dataset")
367 parser.add_argument("--input-len",
368 type=int,
369 default=None,
@@ -547,14 +410,35 @@ if __name__ == "__main__":
410 default=None,
411 help="Path to the lora adapters to use. This can be an absolute path, "
412 "a relative path, or a Hugging Face model identifier.")
413 + parser.add_argument("--prefix-len",
414 + type=int,
415 + default=None,
416 + help="Number of prefix tokens per request."
417 + "This is for the RandomDataset and SonnetDataset")
418 + # random dataset
419 + parser.add_argument(
420 + "--random-range-ratio",
421 + type=float,
422 + default=1.0,
423 + help="Range of sampled ratio of input/output length, "
424 + "used only for RandomDataSet.",
425 + )
426
427 parser = AsyncEngineArgs.add_cli_args(parser)
428 args = parser.parse_args()
429 if args.tokenizer is None:
430 args.tokenizer = args.model
555 - if args.dataset is None:
556 - assert args.input_len is not None
557 - assert args.output_len is not None
431 + if args.dataset is not None:
432 + warnings.warn(
433 + "The '--dataset' argument will be deprecated in the next "
434 + "release. Please use '--dataset-name' and "
435 + "'--dataset-path' in the future runs.",
436 + stacklevel=2)
437 + args.dataset_path = args.dataset
438 + if args.dataset is None and args.dataset_path is None:
439 + # for random dataset, the default sampling setting is in
440 + # benchmark_dataset.RandomDataset
441 + print("When dataset is not set, it will default to random dataset")
442 else:
443 assert args.input_len is None
444 if args.enable_lora: