@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 212 lines 6.41 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import gc
4 import time
5 from unittest import mock
6
7 import numpy as np
8 from benchmark_utils import TimeCollector
9 from tabulate import tabulate
10
11 from vllm.config import (
12 CacheConfig,
13 DeviceConfig,
14 LoadConfig,
15 ModelConfig,
16 ParallelConfig,
17 SchedulerConfig,
18 SpeculativeConfig,
19 VllmConfig,
20 )
21 from vllm.platforms import current_platform
22 from vllm.utils.argparse_utils import FlexibleArgumentParser
23 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
24 from vllm.v1.worker.gpu_input_batch import InputBatch
25 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
26
27
28 def benchmark_propose(args):
29 rows = []
30 for max_ngram in args.max_ngram:
31 collector = TimeCollector(TimeCollector.US)
32
33 model_config = ModelConfig(
34 model="facebook/opt-125m",
35 max_model_len=args.num_token + args.num_spec_token,
36 tokenizer="facebook/opt-125m",
37 tokenizer_mode="auto",
38 dtype="auto",
39 seed=0,
40 trust_remote_code=False,
41 )
42 proposer = NgramProposer(
43 vllm_config=VllmConfig(
44 model_config=model_config,
45 speculative_config=SpeculativeConfig(
46 prompt_lookup_min=args.min_ngram,
47 prompt_lookup_max=max_ngram,
48 num_speculative_tokens=args.num_spec_token,
49 method="ngram",
50 ),
51 )
52 )
53
54 # Warm up
55 proposer.propose(np.random.randint(0, 20, (args.num_token,)))
56
57 gc.collect()
58 for _ in range(args.num_iteration):
59 tokens = np.random.randint(0, 20, (args.num_req, args.num_token))
60 with collector:
61 for i in range(args.num_req):
62 proposer.propose(tokens[i, :])
63 rows.append(
64 [args.num_req, args.num_token, args.min_ngram, max_ngram]
65 + collector.dump_avg_max()
66 )
67
68 print(
69 tabulate(
70 rows,
71 headers=[
72 "# Request",
73 "# Token",
74 "Min Ngram",
75 "Max Ngram",
76 "Avg (us)",
77 "Max (us)",
78 ],
79 tablefmt="grid",
80 floatfmt=".3f",
81 )
82 )
83
84
85 def benchmark_batched_propose(args):
86 NUM_SPECULATIVE_TOKENS_NGRAM = 10
87 PROMPT_LOOKUP_MIN = 5
88 PROMPT_LOOKUP_MAX = 15
89 MAX_MODEL_LEN = int(1e7)
90 DEVICE = current_platform.device_type
91
92 model_config = ModelConfig(model="facebook/opt-125m", runner="generate")
93
94 speculative_config = SpeculativeConfig(
95 target_model_config=model_config,
96 target_parallel_config=ParallelConfig(),
97 method="ngram",
98 num_speculative_tokens=NUM_SPECULATIVE_TOKENS_NGRAM,
99 prompt_lookup_max=PROMPT_LOOKUP_MAX,
100 prompt_lookup_min=PROMPT_LOOKUP_MIN,
101 )
102
103 vllm_config = VllmConfig(
104 model_config=model_config,
105 cache_config=CacheConfig(),
106 speculative_config=speculative_config,
107 device_config=DeviceConfig(device=current_platform.device_type),
108 parallel_config=ParallelConfig(),
109 load_config=LoadConfig(),
110 scheduler_config=SchedulerConfig(
111 max_model_len=model_config.max_model_len,
112 is_encoder_decoder=model_config.is_encoder_decoder,
113 ),
114 )
115
116 # monkey patch vllm.v1.worker.gpu_model_runner.get_pp_group
117 mock_pp_group = mock.MagicMock()
118 mock_pp_group.world_size = 1
119 with mock.patch(
120 "vllm.v1.worker.gpu_model_runner.get_pp_group", return_value=mock_pp_group
121 ):
122 runner = GPUModelRunner(vllm_config, DEVICE)
123
124 # hack max model len
125 runner.max_model_len = MAX_MODEL_LEN
126 runner.drafter.max_model_len = MAX_MODEL_LEN
127
128 dummy_input_batch = InputBatch(
129 max_num_reqs=args.num_req,
130 max_model_len=MAX_MODEL_LEN,
131 max_num_batched_tokens=args.num_req * args.num_token,
132 device=DEVICE,
133 pin_memory=False,
134 vocab_size=256000,
135 block_sizes=[16],
136 )
137 dummy_input_batch._req_ids = list(str(id) for id in range(args.num_req))
138 dummy_input_batch.num_tokens_no_spec = [args.num_token] * args.num_req
139 dummy_input_batch.token_ids_cpu = np.random.randint(
140 0, 20, (args.num_req, args.num_token)
141 )
142
143 runner.input_batch = dummy_input_batch
144
145 sampled_token_ids = [[0]] * args.num_req
146
147 print("Starting benchmark")
148 # first run is warmup so ignore it
149 for _ in range(args.num_iteration):
150 start = time.time()
151 runner.drafter.propose(
152 sampled_token_ids,
153 dummy_input_batch.num_tokens_no_spec,
154 dummy_input_batch.token_ids_cpu,
155 )
156 end = time.time()
157 print(f"Iteration time (s): {end - start}")
158
159
160 def invoke_main() -> None:
161 parser = FlexibleArgumentParser(
162 description="Benchmark the performance of N-gram speculative decode drafting"
163 )
164 parser.add_argument(
165 "--batched", action="store_true", help="consider time to prepare batch"
166 )
167 parser.add_argument(
168 "--num-iteration",
169 type=int,
170 default=100,
171 help="Number of iterations to run to stabilize final data readings",
172 )
173 parser.add_argument(
174 "--num-req", type=int, default=128, help="Number of requests in the batch"
175 )
176 parser.add_argument(
177 "--num-token", type=int, default=1500, help="Number of tokens for each request"
178 )
179 parser.add_argument(
180 "--min-ngram",
181 type=int,
182 default=3,
183 help="Minimum n-gram to match",
184 )
185 parser.add_argument(
186 "--max-ngram",
187 type=int,
188 nargs="*",
189 default=[5, 7, 10, 15, 20],
190 help="Maximum n-gram to match",
191 )
192 parser.add_argument(
193 "--num-spec-token",
194 type=int,
195 default=3,
196 help="Number of speculative tokens to generate",
197 )
198 args = parser.parse_args()
199
200 if not args.batched:
201 benchmark_propose(args)
202 else:
203 benchmark_batched_propose(args)
204
205
206 """
207 # Example command lines:
208 # time python3 benchmarks/benchmark_ngram_proposer.py
209 # time python3 benchmarks/benchmark_ngram_proposer.py --batched --num-iteration 4 --num-token 1000000 --num-req 128
210 """ # noqa: E501
211 if __name__ == "__main__":
212 invoke_main() # pragma: no cover