benchmark_latency.py 8.33 KB
Newer Older
1
"""Benchmark the latency of processing a single batch of requests."""
2
3
import argparse
import time
4
5
from pathlib import Path
from typing import Optional
6
7
8

import numpy as np
import torch
9
from tqdm import tqdm
10

Woosuk Kwon's avatar
Woosuk Kwon committed
11
from vllm import LLM, SamplingParams
12
from vllm.model_executor.layers.quantization import QUANTIZATION_METHODS
13
14
15


def main(args: argparse.Namespace):
16
17
18
    print(args)

    # NOTE(woosuk): If the request cannot be processed in a single batch,
Zhuohan Li's avatar
Zhuohan Li committed
19
    # the engine will automatically process the request in multiple batches.
20
    llm = LLM(model=args.model,
21
22
              speculative_model=args.speculative_model,
              num_speculative_tokens=args.num_speculative_tokens,
23
24
25
26
27
28
29
              tokenizer=args.tokenizer,
              quantization=args.quantization,
              tensor_parallel_size=args.tensor_parallel_size,
              trust_remote_code=args.trust_remote_code,
              dtype=args.dtype,
              enforce_eager=args.enforce_eager,
              kv_cache_dtype=args.kv_cache_dtype,
30
              quantization_param_path=args.quantization_param_path,
31
32
              device=args.device,
              ray_workers_use_nsight=args.ray_workers_use_nsight,
33
              use_v2_block_manager=args.use_v2_block_manager,
34
35
36
              enable_chunked_prefill=args.enable_chunked_prefill,
              download_dir=args.download_dir,
              block_size=args.block_size)
37

Woosuk Kwon's avatar
Woosuk Kwon committed
38
39
40
41
42
    sampling_params = SamplingParams(
        n=args.n,
        temperature=0.0 if args.use_beam_search else 1.0,
        top_p=1.0,
        use_beam_search=args.use_beam_search,
43
        ignore_eos=True,
Woosuk Kwon's avatar
Woosuk Kwon committed
44
45
        max_tokens=args.output_len,
    )
46
    print(sampling_params)
47
48
49
50
    dummy_prompt_token_ids = np.random.randint(10000,
                                               size=(args.batch_size,
                                                     args.input_len))
    dummy_prompt_token_ids = dummy_prompt_token_ids.tolist()
51

52
53
54
55
56
57
58
59
60
    def run_to_completion(profile_dir: Optional[str] = None):
        if profile_dir:
            with torch.profiler.profile(
                    activities=[
                        torch.profiler.ProfilerActivity.CPU,
                        torch.profiler.ProfilerActivity.CUDA,
                    ],
                    on_trace_ready=torch.profiler.tensorboard_trace_handler(
                        str(profile_dir))) as p:
61
62
63
64
65
66
67
68
69
70
71
72
                llm.generate(prompt_token_ids=dummy_prompt_token_ids,
                             sampling_params=sampling_params,
                             use_tqdm=False)
            print(p.key_averages())
        else:
            start_time = time.perf_counter()
            llm.generate(prompt_token_ids=dummy_prompt_token_ids,
                         sampling_params=sampling_params,
                         use_tqdm=False)
            end_time = time.perf_counter()
            latency = end_time - start_time
            return latency
73

74
    print("Warming up...")
75
76
    for _ in tqdm(range(args.num_iters_warmup), desc="Warmup iterations"):
        run_to_completion(profile_dir=None)
77

78
    if args.profile:
79
80
        profile_dir = args.profile_result_dir
        if not profile_dir:
81
82
83
            profile_dir = Path(
                "."
            ) / "vllm_benchmark_result" / f"latency_result_{time.time()}"
84
        print(f"Profiling (results will be saved to '{profile_dir}')...")
85
        run_to_completion(profile_dir=profile_dir)
86
87
        return

88
89
    # Benchmark.
    latencies = []
90
    for _ in tqdm(range(args.num_iters), desc="Profiling iterations"):
91
        latencies.append(run_to_completion(profile_dir=None))
92
93
94
    latencies = np.array(latencies)
    percentages = [10, 25, 50, 75, 90]
    percentiles = np.percentile(latencies, percentages)
95
    print(f'Avg latency: {np.mean(latencies)} seconds')
96
97
    for percentage, percentile in zip(percentages, percentiles):
        print(f'{percentage}% percentile latency: {percentile} seconds')
98
99
100


if __name__ == '__main__':
101
    parser = argparse.ArgumentParser(
102
        description='Benchmark the latency of processing a single batch of '
103
        'requests till completion.')
104
    parser.add_argument('--model', type=str, default='facebook/opt-125m')
105
106
    parser.add_argument('--speculative-model', type=str, default=None)
    parser.add_argument('--num-speculative-tokens', type=int, default=None)
107
    parser.add_argument('--tokenizer', type=str, default=None)
108
109
    parser.add_argument('--quantization',
                        '-q',
110
                        choices=[*QUANTIZATION_METHODS, None],
111
                        default=None)
112
    parser.add_argument('--tensor-parallel-size', '-tp', type=int, default=1)
113
114
115
    parser.add_argument('--input-len', type=int, default=32)
    parser.add_argument('--output-len', type=int, default=128)
    parser.add_argument('--batch-size', type=int, default=8)
116
117
118
    parser.add_argument('--n',
                        type=int,
                        default=1,
119
                        help='Number of generated sequences per prompt.')
120
    parser.add_argument('--use-beam-search', action='store_true')
121
122
123
124
    parser.add_argument('--num-iters-warmup',
                        type=int,
                        default=10,
                        help='Number of iterations to run for warmup.')
125
126
    parser.add_argument('--num-iters',
                        type=int,
127
                        default=30,
128
                        help='Number of iterations to run.')
129
130
    parser.add_argument('--trust-remote-code',
                        action='store_true',
131
                        help='trust remote code from huggingface')
132
133
134
135
136
137
138
139
140
    parser.add_argument(
        '--dtype',
        type=str,
        default='auto',
        choices=['auto', 'half', 'float16', 'bfloat16', 'float', 'float32'],
        help='data type for model weights and activations. '
        'The "auto" option will use FP16 precision '
        'for FP32 and FP16 models, and BF16 precision '
        'for BF16 models.')
141
142
143
    parser.add_argument('--enforce-eager',
                        action='store_true',
                        help='enforce eager mode and disable CUDA graph')
144
145
146
    parser.add_argument(
        "--kv-cache-dtype",
        type=str,
147
        choices=['auto', 'fp8'],
148
149
        default='auto',
        help=
150
151
152
153
154
155
156
157
158
159
160
161
162
163
        'Data type for kv cache storage. If "auto", will use model data type. '
        'FP8_E5M2 (without scaling) is only supported on cuda version greater '
        'than 11.8. On ROCm (AMD GPU), FP8_E4M3 is instead supported for '
        'common inference criteria.')
    parser.add_argument(
        '--quantization-param-path',
        type=str,
        default=None,
        help='Path to the JSON file containing the KV cache scaling factors. '
        'This should generally be supplied, when KV cache dtype is FP8. '
        'Otherwise, KV cache scaling factors default to 1.0, which may cause '
        'accuracy issues. FP8_E5M2 (without scaling) is only supported on '
        'cuda version greater than 11.8. On ROCm (AMD GPU), FP8_E4M3 is '
        'instead supported for common inference criteria.')
164
165
166
167
    parser.add_argument(
        '--profile',
        action='store_true',
        help='profile the generation process of a single batch')
168
169
170
171
    parser.add_argument(
        '--profile-result-dir',
        type=str,
        default=None,
172
173
        help=('path to save the pytorch profiler output. Can be visualized '
              'with ui.perfetto.dev or Tensorboard.'))
174
175
176
177
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
178
179
        choices=["cuda", "cpu"],
        help='device type for vLLM execution, supporting CUDA and CPU.')
180
181
182
183
184
185
    parser.add_argument('--block-size',
                        type=int,
                        default=16,
                        help='block size of key/value cache')
    parser.add_argument(
        '--enable-chunked-prefill',
186
        action='store_true',
187
188
        help='If True, the prefill requests can be chunked based on the '
        'max_num_batched_tokens')
189
    parser.add_argument('--use-v2-block-manager', action='store_true')
190
191
192
193
194
    parser.add_argument(
        "--ray-workers-use-nsight",
        action='store_true',
        help="If specified, use nsight to profile ray workers",
    )
195
196
197
198
199
    parser.add_argument('--download-dir',
                        type=str,
                        default=None,
                        help='directory to download and load the weights, '
                        'default to the default cache dir of huggingface')
200
201
    args = parser.parse_args()
    main(args)