benchmark_throughput.py 22.3 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
"""Benchmark offline inference throughput."""
3
import argparse
4
import dataclasses
5
6
7
import json
import random
import time
zhuwenwen's avatar
zhuwenwen committed
8
from pathlib import Path
9
10
from functools import cache
from typing import Dict, List, Optional, Tuple
11

zhuwenwen's avatar
zhuwenwen committed
12
import numpy as np
13
import torch
14
import uvloop
15
from PIL import Image
16
from tqdm import tqdm
17
18
from transformers import (AutoModelForCausalLM, AutoTokenizer,
                          PreTrainedTokenizerBase)
19

zhuwenwen's avatar
zhuwenwen committed
20
21

from vllm.inputs import PromptType
22
from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
23
24
from vllm.entrypoints.openai.api_server import (
    build_async_engine_client_from_engine_args)
25
from vllm.inputs import TextPrompt
26
27
from vllm.lora.request import LoRARequest
from vllm.lora.utils import get_adapter_absolute_path
28
from vllm.multimodal import MultiModalDataDict
29
from vllm.sampling_params import BeamSearchParams
30
from vllm.transformers_utils.tokenizer import AnyTokenizer, get_lora_tokenizer
31
from vllm.utils import FlexibleArgumentParser, merge_async_iterators
32

33

34
35
36
37
38
39
40
41
@dataclasses.dataclass
class SampleRequest:
    """A class representing a single inference request for benchmarking.

    Attributes:
        prompt: The input text prompt for the model.
        prompt_len: The length of the prompt in tokens.
        expected_output_len: The expected length of the output in tokens.
42
43
44
        multi_modal_data: Optional dictionary containing multi-modal data (e.g.
            images).
        lora_request: Optional LoRARequest specifying the LoRA to use. 
45
46
47
48
49
    """
    prompt: str
    prompt_len: int
    expected_output_len: int
    multi_modal_data: Optional[MultiModalDataDict] = None
50
    lora_request: Optional[LoRARequest] = None
51
52


53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def _get_prompt_for_image_model(question: str, *, model: str) -> str:
    """Prepend and append special tokens around the question to form a prompt.

    Args:
        question: The input question text to wrap with special tokens
        model: The name of the model being used, to determine which special
            tokens to add

    Returns:
        The formatted prompt string with appropriate special tokens for the
            model

    Raises:
        ValueError: If an unsupported model name is provided
    """
    model = model.lower()
    if "pixtral" in model:
        return f"<s>[INST]{question}\n[IMG][/INST]"
    raise ValueError(f"Unsupported model {model}")


74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
@cache
def lora_path_on_disk(lora_path: str) -> str:
    return get_adapter_absolute_path(lora_path)


lora_tokenizer_cache: Dict[int, AnyTokenizer] = {}


def get_random_lora_request(
        args: argparse.Namespace
) -> Tuple[LoRARequest, Optional[AnyTokenizer]]:
    global lora_tokenizer_cache
    lora_id = random.randint(1, args.max_loras)
    lora_request = LoRARequest(lora_name=str(lora_id),
                               lora_int_id=lora_id,
                               lora_path=lora_path_on_disk(args.lora_path))
    if lora_id not in lora_tokenizer_cache:
        lora_tokenizer_cache[lora_id] = get_lora_tokenizer(lora_request)
    return lora_request, lora_tokenizer_cache[lora_id]


95
96
def sample_requests(tokenizer: PreTrainedTokenizerBase,
                    args: argparse.Namespace) -> List[SampleRequest]:
97

98
99
100
101
    dataset_path: str = args.dataset
    num_requests: int = args.num_prompts
    fixed_output_len: Optional[int] = args.output_len
    model: str = args.model
102
103
    if fixed_output_len is not None and fixed_output_len < 4:
        raise ValueError("output_len too small")
104

105
106
107
108
    # Load the dataset.
    with open(dataset_path) as f:
        dataset = json.load(f)
    # Filter out the conversations with less than 2 turns.
109
    dataset = [data for data in dataset if len(data["conversations"]) >= 2]
110
111
    # Shuffle the dataset.
    random.shuffle(dataset)
112

113
    # Filter out sequences that are too long or too short
114
    filtered_dataset: List[SampleRequest] = []
115
116
117
    for data in tqdm(dataset,
                     total=len(filtered_dataset),
                     desc="sampling requests"):
118
119
120
        if len(filtered_dataset) == num_requests:
            break

121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
        # Only keep the first two turns of each conversation.
        prompt = data["conversations"][0]["value"]
        completion = data["conversations"][1]["value"]

        multi_modal_data: Optional[MultiModalDataDict] = None
        if "image" in data:
            multi_modal_data = multi_modal_data or {}
            image_path = data["image"]
            # TODO(vllm-project/vllm/issues/9778): Support multiple images.
            assert isinstance(image_path,
                              str), "Only support single image input"
            try:
                multi_modal_data["image"] = Image.open(image_path).convert(
                    "RGB")
            except FileNotFoundError:
                # Ignore datapoint where asset is missing
                continue
            prompt = _get_prompt_for_image_model(question=prompt, model=model)

140
141
142
143
144
145
146
        request_tokenizer = tokenizer
        lora_request: Optional[LoRARequest] = None
        if args.enable_lora:
            lora_request, lora_tokenizer = get_random_lora_request(args)
            if lora_tokenizer:
                request_tokenizer = lora_tokenizer

147
        # Tokenize the prompts and completions.
148
149
        prompt_token_ids = request_tokenizer(prompt).input_ids
        completion_token_ids = request_tokenizer(completion).input_ids
150
        prompt_len = len(prompt_token_ids)
151
152
        output_len = len(completion_token_ids
                         ) if fixed_output_len is None else fixed_output_len
153
154
155
156
157
158
        if prompt_len < 4 or output_len < 4:
            # Prune too short sequences.
            continue
        if prompt_len > 1024 or prompt_len + output_len > 2048:
            # Prune too long sequences.
            continue
159
160
161
        filtered_dataset.append(
            SampleRequest(prompt=prompt,
                          prompt_len=prompt_len,
162
                          expected_output_len=output_len,
163
164
                          multi_modal_data=multi_modal_data,
                          lora_request=lora_request))
165

166
    return filtered_dataset
167
168


Woosuk Kwon's avatar
Woosuk Kwon committed
169
def run_vllm(
170
    requests: List[SampleRequest],
171
    n: int,
172
    num_iters_warmup: int,
173
    engine_args: EngineArgs,
174
) -> float:
175
    from vllm import LLM, SamplingParams
176
    llm = LLM(**dataclasses.asdict(engine_args))
177

Zhuohan Li's avatar
Zhuohan Li committed
178
    # Add the requests to the engine.
179
    prompts: List[TextPrompt] = []
180
    sampling_params: List[SamplingParams] = []
181
    for request in requests:
182
183
184
        prompts.append(
            TextPrompt(prompt=request.prompt,
                       multi_modal_data=request.multi_modal_data))
185
186
187
        sampling_params.append(
            SamplingParams(
                n=n,
188
                temperature=1.0,
189
190
                top_p=1.0,
                ignore_eos=True,
191
                max_tokens=request.expected_output_len,
192
            ))
193
194
195
    lora_requests: Optional[List[LoRARequest]] = None
    if engine_args.enable_lora:
        lora_requests = [request.lora_request for request in requests]
196

zhuwenwen's avatar
zhuwenwen committed
197
    # warmup
198
199
200
201
202
203
204
205
206
207
208
    warmup_sampling_params = SamplingParams(
        n=args.n,
        temperature=1.0,
        top_p=1.0,
        ignore_eos=True,
        max_tokens=10,
    )
    dummy_prompt_token_ids = np.random.randint(10000, size=(1,10))
    dummy_prompts: List[PromptType] = [{
        "prompt_token_ids": batch
    } for batch in dummy_prompt_token_ids.tolist()]
zhuwenwen's avatar
zhuwenwen committed
209
    
210
211
212
213
214
    print("Warming up...")
    for _ in tqdm(range(num_iters_warmup), desc="Warmup iterations"):
        llm.generate(dummy_prompts,
                        sampling_params=warmup_sampling_params,
                        use_tqdm=False)
zhuwenwen's avatar
zhuwenwen committed
215

216
217
218
    use_beam_search = False

    if not use_beam_search:
zhuwenwen's avatar
zhuwenwen committed
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
        if args.profile:
            profile_dir = args.profile_result_dir
            if not profile_dir:
                profile_dir = Path(
                    "."
                ) / "vllm_benchmark_result" / f"latency_result_{time.time()}"
            print(f"Profiling (results will be saved to '{profile_dir}')...")
            with torch.profiler.profile(
                        activities=[torch.profiler.ProfilerActivity.CPU,
                                    torch.profiler.ProfilerActivity.CUDA,
                        ],record_shapes=True,
                        on_trace_ready=torch.profiler.tensorboard_trace_handler(str(profile_dir))
                        ) as prof:
                start = time.perf_counter()
                llm.generate(prompts,
                        sampling_params,
                        lora_request=lora_requests,
                        use_tqdm=True)
                end = time.perf_counter()
            print('Prepare time report')
            print(prof.key_averages(group_by_input_shape=True).table(sort_by="self_cuda_time_total", row_limit=-1))
        else:
            start = time.perf_counter()
            llm.generate(prompts,
                        sampling_params,
                        lora_request=lora_requests,
                        use_tqdm=True)
            end = time.perf_counter()  
247
    else:
248
        assert lora_requests is None, "BeamSearch API does not support LoRA"
249
        prompts = [request.prompt for request in requests]
250
251
        # output_len should be the same for all requests.
        output_len = requests[0][2]
252
253
        for request in requests:
            assert request.expected_output_len == output_len
254
        start = time.perf_counter()
255
256
257
258
259
260
261
        llm.beam_search(
            prompts,
            BeamSearchParams(
                beam_width=n,
                max_tokens=output_len,
                ignore_eos=True,
            ))
262
        end = time.perf_counter()
263
264
265
    return end - start


266
async def run_vllm_async(
267
    requests: List[SampleRequest],
268
    n: int,
269
    engine_args: AsyncEngineArgs,
270
271
272
273
274
275
276
277
    disable_frontend_multiprocessing: bool = False,
) -> float:
    from vllm import SamplingParams

    async with build_async_engine_client_from_engine_args(
            engine_args, disable_frontend_multiprocessing) as llm:

        # Add the requests to the engine.
278
        prompts: List[TextPrompt] = []
279
        sampling_params: List[SamplingParams] = []
280
        lora_requests: List[Optional[LoRARequest]] = []
281
        for request in requests:
282
283
284
            prompts.append(
                TextPrompt(prompt=request.prompt,
                           multi_modal_data=request.multi_modal_data))
285
286
287
            sampling_params.append(
                SamplingParams(
                    n=n,
288
                    temperature=1.0,
289
290
                    top_p=1.0,
                    ignore_eos=True,
291
                    max_tokens=request.expected_output_len,
292
                ))
293
            lora_requests.append(request.lora_request)
294
295
296

        generators = []
        start = time.perf_counter()
297
298
299
300
301
302
        for i, (prompt, sp,
                lr) in enumerate(zip(prompts, sampling_params, lora_requests)):
            generator = llm.generate(prompt,
                                     sp,
                                     lora_request=lr,
                                     request_id=f"test{i}")
303
304
305
306
307
308
309
310
            generators.append(generator)
        all_gens = merge_async_iterators(*generators)
        async for i, res in all_gens:
            pass
        end = time.perf_counter()
        return end - start


311
def run_hf(
312
    requests: List[SampleRequest],
313
314
315
316
    model: str,
    tokenizer: PreTrainedTokenizerBase,
    n: int,
    max_batch_size: int,
317
    trust_remote_code: bool,
318
) -> float:
319
320
    llm = AutoModelForCausalLM.from_pretrained(
        model, torch_dtype=torch.float16, trust_remote_code=trust_remote_code)
321
322
323
    if llm.config.model_type == "llama":
        # To enable padding in the HF backend.
        tokenizer.pad_token = tokenizer.eos_token
324
325
326
    llm = llm.cuda()

    pbar = tqdm(total=len(requests))
327
    start = time.perf_counter()
328
329
330
331
332
333
334
335
336
337
338
339
    batch: List[str] = []
    max_prompt_len = 0
    max_output_len = 0
    for i in range(len(requests)):
        prompt, prompt_len, output_len = requests[i]
        # Add the prompt to the batch.
        batch.append(prompt)
        max_prompt_len = max(max_prompt_len, prompt_len)
        max_output_len = max(max_output_len, output_len)
        if len(batch) < max_batch_size and i != len(requests) - 1:
            # Check if we can add more requests to the batch.
            _, next_prompt_len, next_output_len = requests[i + 1]
340
341
            if (max(max_prompt_len, next_prompt_len) +
                    max(max_output_len, next_output_len)) <= 2048:
342
343
344
345
                # We can add more requests to the batch.
                continue

        # Generate the sequences.
346
347
        input_ids = tokenizer(batch, return_tensors="pt",
                              padding=True).input_ids
348
349
        llm_outputs = llm.generate(
            input_ids=input_ids.cuda(),
350
            do_sample=True,
351
352
353
354
355
356
357
358
359
360
361
362
363
364
            num_return_sequences=n,
            temperature=1.0,
            top_p=1.0,
            use_cache=True,
            max_new_tokens=max_output_len,
        )
        # Include the decoding time.
        tokenizer.batch_decode(llm_outputs, skip_special_tokens=True)
        pbar.update(len(batch))

        # Clear the batch.
        batch = []
        max_prompt_len = 0
        max_output_len = 0
365
    end = time.perf_counter()
366
367
368
    return end - start


369
def run_mii(
370
    requests: List[SampleRequest],
371
372
373
374
    model: str,
    tensor_parallel_size: int,
    output_len: int,
) -> float:
375
376
    from mii import client, serve
    llm = serve(model, tensor_parallel=tensor_parallel_size)
377
    prompts = [request.prompt for request in requests]
378
379

    start = time.perf_counter()
380
    llm.generate(prompts, max_new_tokens=output_len)
381
    end = time.perf_counter()
382
383
    client = client(model)
    client.terminate_server()
384
385
386
    return end - start


387
388
389
390
391
def main(args: argparse.Namespace):
    print(args)
    random.seed(args.seed)

    # Sample the requests.
392
393
394
    tokenizer = AutoTokenizer.from_pretrained(
        args.tokenizer, trust_remote_code=args.trust_remote_code)
    if args.dataset is None:
395
396
397
        vocab_size = tokenizer.vocab_size
        requests = []
        for _ in range(args.num_prompts):
398
399
400
401
402
403
404
405

            request_tokenizer = tokenizer
            lora_request: Optional[LoRARequest] = None
            if args.enable_lora:
                lora_request, lora_tokenizer = get_random_lora_request(args)
                if lora_tokenizer:
                    request_tokenizer = lora_tokenizer

406
407
408
409
410
411
412
413
            # Synthesize a prompt with the given input length.
            candidate_ids = [
                random.randint(0, vocab_size - 1)
                for _ in range(args.input_len)
            ]
            # As tokenizer may add additional tokens like BOS, we need to try
            # different lengths to get the desired input length.
            for _ in range(5):  # Max attempts to correct
414
415
                candidate_prompt = request_tokenizer.decode(candidate_ids)
                tokenized_len = len(request_tokenizer.encode(candidate_prompt))
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431

                if tokenized_len == args.input_len:
                    break

                # Adjust length based on difference
                diff = args.input_len - tokenized_len
                if diff > 0:
                    candidate_ids.extend([
                        random.randint(100, vocab_size - 100)
                        for _ in range(diff)
                    ])
                else:
                    candidate_ids = candidate_ids[:diff]
            requests.append(
                SampleRequest(prompt=candidate_prompt,
                              prompt_len=args.input_len,
432
433
                              expected_output_len=args.output_len,
                              lora_request=lora_request))
434
    else:
435
        requests = sample_requests(tokenizer, args)
436

437
438
    is_multi_modal = any(request.multi_modal_data is not None
                         for request in requests)
Woosuk Kwon's avatar
Woosuk Kwon committed
439
    if args.backend == "vllm":
440
        if args.async_engine:
441
442
443
444
445
446
447
            elapsed_time = uvloop.run(
                run_vllm_async(
                    requests,
                    args.n,
                    AsyncEngineArgs.from_cli_args(args),
                    args.disable_frontend_multiprocessing,
                ))
448
        else:
449
            elapsed_time = run_vllm(requests, args.n, args.num_iters_warmup,
450
                                    EngineArgs.from_cli_args(args))
451
452
    elif args.backend == "hf":
        assert args.tensor_parallel_size == 1
453
        elapsed_time = run_hf(requests, args.model, tokenizer, args.n,
454
                              args.hf_max_batch_size, args.trust_remote_code)
455
456
457
    elif args.backend == "mii":
        elapsed_time = run_mii(requests, args.model, args.tensor_parallel_size,
                               args.output_len)
458
459
    else:
        raise ValueError(f"Unknown backend: {args.backend}")
460
461
462
    total_num_tokens = sum(request.prompt_len + request.expected_output_len
                           for request in requests)
    total_output_tokens = sum(request.expected_output_len
zhuwenwen's avatar
zhuwenwen committed
463
                            for request in requests)
464
465
466
467
468
    if is_multi_modal:
        print("\033[91mWARNING\033[0m: Multi-modal request detected. The "
              "following metrics are not accurate because image tokens are not"
              " counted. See vllm-project/vllm/issues/9778 for details.")
        # TODO(vllm-project/vllm/issues/9778): Count molti-modal token length.
zhuwenwen's avatar
zhuwenwen committed
469
    print(f"Latency: {elapsed_time:.2f} s")
Woosuk Kwon's avatar
Woosuk Kwon committed
470
    print(f"Throughput: {len(requests) / elapsed_time:.2f} requests/s, "
471
472
          f"{total_num_tokens / elapsed_time:.2f} total tokens/s, "
          f"{total_output_tokens / elapsed_time:.2f} output tokens/s")
473

474
475
476
477
478
479
480
481
482
483
484
485
    # Output JSON results if specified
    if args.output_json:
        results = {
            "elapsed_time": elapsed_time,
            "num_requests": len(requests),
            "total_num_tokens": total_num_tokens,
            "requests_per_second": len(requests) / elapsed_time,
            "tokens_per_second": total_num_tokens / elapsed_time,
        }
        with open(args.output_json, "w") as f:
            json.dump(results, f, indent=4)

486
487

if __name__ == "__main__":
488
    parser = FlexibleArgumentParser(description="Benchmark the throughput.")
489
490
    parser.add_argument("--backend",
                        type=str,
491
                        choices=["vllm", "hf", "mii"],
Woosuk Kwon's avatar
Woosuk Kwon committed
492
                        default="vllm")
493
494
    parser.add_argument("--dataset",
                        type=str,
495
                        default=None,
496
497
498
                        help="Path to the dataset. The dataset is expected to "
                        "be a json in form of List[Dict[..., conversations: "
                        "List[Dict[..., value: <prompt_or_response>]]]]")
499
500
501
502
503
504
505
506
507
    parser.add_argument("--input-len",
                        type=int,
                        default=None,
                        help="Input prompt length for each request")
    parser.add_argument("--output-len",
                        type=int,
                        default=None,
                        help="Output length for each request. Overrides the "
                        "output length from the dataset.")
508
509
510
    parser.add_argument("--n",
                        type=int,
                        default=1,
511
                        help="Number of generated sequences per prompt.")
zhuwenwen's avatar
zhuwenwen committed
512
513
514
515
    parser.add_argument('--num-iters-warmup',
                        type=int,
                        default=1,
                        help='Number of iterations to run for warmup.')
516
517
518
    parser.add_argument("--num-prompts",
                        type=int,
                        default=1000,
519
                        help="Number of prompts to process.")
520
521
522
    parser.add_argument("--hf-max-batch-size",
                        type=int,
                        default=None,
523
                        help="Maximum batch size for HF backend.")
zhuwenwen's avatar
zhuwenwen committed
524
525
526
527
528
529
530
531
532
533
    parser.add_argument(
        '--profile',
        action='store_true',
        help='profile the generation process of a single batch')
    parser.add_argument(
        '--profile-result-dir',
        type=str,
        default=None,
        help=('path to save the pytorch profiler output. Can be visualized '
              'with ui.perfetto.dev or Tensorboard.'))
534
535
536
537
538
    parser.add_argument(
        '--output-json',
        type=str,
        default=None,
        help='Path to save the throughput results in JSON format.')
539
540
541
542
543
544
545
546
    parser.add_argument("--async-engine",
                        action='store_true',
                        default=False,
                        help="Use vLLM async engine rather than LLM class.")
    parser.add_argument("--disable-frontend-multiprocessing",
                        action='store_true',
                        default=False,
                        help="Disable decoupled async engine frontend.")
547
548
549
550
551
552
553
554
    # LoRA
    parser.add_argument(
        "--lora-path",
        type=str,
        default=None,
        help="Path to the lora adapters to use. This can be an absolute path, "
        "a relative path, or a Hugging Face model identifier.")

555
    parser = AsyncEngineArgs.add_cli_args(parser)
556
    args = parser.parse_args()
557
558
559
560
561
562
563
    if args.tokenizer is None:
        args.tokenizer = args.model
    if args.dataset is None:
        assert args.input_len is not None
        assert args.output_len is not None
    else:
        assert args.input_len is None
564
565
    if args.enable_lora:
        assert args.lora_path is not None
566

Woosuk Kwon's avatar
Woosuk Kwon committed
567
    if args.backend == "vllm":
568
569
570
571
572
        if args.hf_max_batch_size is not None:
            raise ValueError("HF max batch size is only for HF backend.")
    elif args.backend == "hf":
        if args.hf_max_batch_size is None:
            raise ValueError("HF max batch size is required for HF backend.")
573
574
        if args.quantization is not None:
            raise ValueError("Quantization is only for vLLM backend.")
575
576
577
        if args.enable_lora is not None:
            raise ValueError("LoRA benchmarking is only supported for vLLM"
                             " backend")
578
579
580
581
582
583
584
585
586
587
588
589
    elif args.backend == "mii":
        if args.dtype != "auto":
            raise ValueError("dtype must be auto for MII backend.")
        if args.n != 1:
            raise ValueError("n must be 1 for MII backend.")
        if args.quantization is not None:
            raise ValueError("Quantization is only for vLLM backend.")
        if args.hf_max_batch_size is not None:
            raise ValueError("HF max batch size is only for HF backend.")
        if args.tokenizer != args.model:
            raise ValueError("Tokenizer must be the same as the model for MII "
                             "backend.")
590
591
592
        if args.enable_lora is not None:
            raise ValueError("LoRA benchmarking is only supported for vLLM"
                             " backend")
593
    main(args)