| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from transformers import AutoTokenizer |
| 5 | |
| 6 | from vllm import LLM, SamplingParams |
| 7 | from vllm.benchmarks.datasets import add_dataset_parser, get_samples |
| 8 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 9 | from vllm.v1.metrics.reader import Counter, Vector |
| 10 | |
| 11 | QUESTION = "What is the content of each image?" |
| 12 | IMAGE_URLS = [ |
| 13 | "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/duck.jpg", |
| 14 | "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/lion.jpg", |
| 15 | "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/flycatcher.jpeg", |
| 16 | "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/somefish.jpg", |
| 17 | "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/starfish.jpg", |
| 18 | "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/snail.jpg", |
| 19 | "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/thistle.jpg", |
| 20 | "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/husky.jpg", |
| 21 | "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/orangetabbycat.jpg", |
| 22 | "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/guineapig.jpg", |
| 23 | "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/rabbit.jpg", |
| 24 | "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/horsepony.jpg", |
| 25 | ] |
| 26 | |
| 27 | |
| 28 | def get_custom_mm_prompts(num_prompts): |
| 29 | prompts = [] |
| 30 | for url in IMAGE_URLS: |
| 31 | prompts.append( |
| 32 | [ |
| 33 | {"type": "image_url", "image_url": {"url": url}}, |
| 34 | {"type": "text", "text": QUESTION}, |
| 35 | ] |
| 36 | ) |
| 37 | if num_prompts > len(IMAGE_URLS): |
| 38 | prompts = prompts * (num_prompts // len(IMAGE_URLS) + 1) |
| 39 | |
| 40 | return [[{"role": "user", "content": prompt}] for prompt in prompts[:num_prompts]] |
| 41 | |
| 42 | |
| 43 | def parse_args(): |
| 44 | parser = FlexibleArgumentParser() |
| 45 | add_dataset_parser(parser) |
| 46 | parser.add_argument("--test", action="store_true") |
| 47 | parser.add_argument( |
| 48 | "--method", |
| 49 | type=str, |
| 50 | default="eagle", |
| 51 | choices=["ngram", "eagle", "eagle3", "mtp", "draft_model"], |
| 52 | ) |
| 53 | parser.add_argument("--backend", type=str, default="openai") |
| 54 | parser.add_argument("--num-spec-tokens", type=int, default=2) |
| 55 | parser.add_argument("--prompt-lookup-max", type=int, default=5) |
| 56 | parser.add_argument("--prompt-lookup-min", type=int, default=2) |
| 57 | parser.add_argument("--tp", type=int, default=1) |
| 58 | parser.add_argument("--enforce-eager", action="store_true") |
| 59 | parser.add_argument("--enable-chunked-prefill", action="store_true") |
| 60 | parser.add_argument("--max-model-len", type=int, default=16384) |
| 61 | parser.add_argument("--temp", type=float, default=0) |
| 62 | parser.add_argument("--top-p", type=float, default=1.0) |
| 63 | parser.add_argument("--top-k", type=int, default=-1) |
| 64 | parser.add_argument("--print-output", action="store_true") |
| 65 | parser.add_argument("--output-len", type=int, default=256) |
| 66 | parser.add_argument("--model-dir", type=str, default=None) |
| 67 | parser.add_argument("--eagle-dir", type=str, default=None) |
| 68 | parser.add_argument("--draft-model", type=str, default=None) |
| 69 | parser.add_argument("--custom-mm-prompts", action="store_true") |
| 70 | parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) |
| 71 | parser.add_argument("--disable-padded-drafter-batch", action="store_true") |
| 72 | parser.add_argument("--max-num-seqs", type=int, default=None) |
| 73 | parser.add_argument("--parallel-drafting", action="store_true") |
| 74 | parser.add_argument("--allowed-local-media-path", type=str, default="") |
| 75 | parser.add_argument("--use-heterogeneous-vocab", action="store_true") |
| 76 | return parser.parse_args() |
| 77 | |
| 78 | |
| 79 | def main(args): |
| 80 | model_dir = args.model_dir |
| 81 | if args.model_dir is None: |
| 82 | if args.custom_mm_prompts: |
| 83 | raise ValueError( |
| 84 | "custom_mm_prompts requires mm based models" |
| 85 | "default llama3.1-8b-instruct is not mm based" |
| 86 | "please specify model_dir to give a mm based model" |
| 87 | ) |
| 88 | model_dir = "meta-llama/Llama-3.1-8B-Instruct" |
| 89 | tokenizer = AutoTokenizer.from_pretrained(model_dir) |
| 90 | |
| 91 | if args.custom_mm_prompts: |
| 92 | prompts = llm_prompts = get_custom_mm_prompts(args.num_prompts) |
| 93 | else: |
| 94 | prompts = get_samples(args, tokenizer) |
| 95 | if args.enable_multimodal_chat: |
| 96 | llm_prompts = [p.prompt for p in prompts] |
| 97 | else: |
| 98 | # add_special_tokens is False to avoid adding bos twice |
| 99 | # when using chat templates |
| 100 | llm_prompts = [ |
| 101 | { |
| 102 | "prompt_token_ids": tokenizer.encode( |
| 103 | prompt.prompt, add_special_tokens=False |
| 104 | ), |
| 105 | "multi_modal_data": prompt.multi_modal_data, |
| 106 | } |
| 107 | for prompt in prompts |
| 108 | ] |
| 109 | if args.method == "eagle" or args.method == "eagle3": |
| 110 | eagle_dir = args.eagle_dir |
| 111 | if args.method == "eagle" and eagle_dir is None: |
| 112 | eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B" |
| 113 | |
| 114 | elif args.method == "eagle3" and eagle_dir is None: |
| 115 | eagle_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" |
| 116 | speculative_config = { |
| 117 | "method": args.method, |
| 118 | "model": eagle_dir, |
| 119 | "num_speculative_tokens": args.num_spec_tokens, |
| 120 | "disable_padded_drafter_batch": args.disable_padded_drafter_batch, |
| 121 | "parallel_drafting": args.parallel_drafting, |
| 122 | } |
| 123 | elif args.method == "ngram": |
| 124 | speculative_config = { |
| 125 | "method": "ngram", |
| 126 | "num_speculative_tokens": args.num_spec_tokens, |
| 127 | "prompt_lookup_max": args.prompt_lookup_max, |
| 128 | "prompt_lookup_min": args.prompt_lookup_min, |
| 129 | } |
| 130 | elif args.method == "draft_model": |
| 131 | assert args.draft_model is not None and args.draft_model != "" |
| 132 | speculative_config = { |
| 133 | "method": args.method, |
| 134 | "model": args.draft_model, |
| 135 | "num_speculative_tokens": args.num_spec_tokens, |
| 136 | "enforce_eager": args.enforce_eager, |
| 137 | "max_model_len": args.max_model_len, |
| 138 | "parallel_drafting": args.parallel_drafting, |
| 139 | "use_heterogeneous_vocab": args.use_heterogeneous_vocab, |
| 140 | } |
| 141 | elif args.method == "mtp": |
| 142 | speculative_config = { |
| 143 | "method": "mtp", |
| 144 | "num_speculative_tokens": args.num_spec_tokens, |
| 145 | } |
| 146 | else: |
| 147 | raise ValueError(f"unknown method: {args.method}") |
| 148 | |
| 149 | llm = LLM( |
| 150 | model=model_dir, |
| 151 | trust_remote_code=True, |
| 152 | tensor_parallel_size=args.tp, |
| 153 | enable_chunked_prefill=args.enable_chunked_prefill, |
| 154 | enforce_eager=args.enforce_eager, |
| 155 | gpu_memory_utilization=args.gpu_memory_utilization, |
| 156 | speculative_config=speculative_config, |
| 157 | disable_log_stats=False, |
| 158 | max_model_len=args.max_model_len, |
| 159 | limit_mm_per_prompt={"image": 5}, |
| 160 | disable_chunked_mm_input=True, |
| 161 | max_num_seqs=args.max_num_seqs, |
| 162 | allowed_local_media_path=args.allowed_local_media_path, |
| 163 | ) |
| 164 | |
| 165 | sampling_params = SamplingParams(temperature=args.temp, max_tokens=args.output_len) |
| 166 | if args.backend == "openai-chat": |
| 167 | outputs = llm.chat(llm_prompts, sampling_params=sampling_params) |
| 168 | else: |
| 169 | outputs = llm.generate( |
| 170 | llm_prompts, |
| 171 | sampling_params=sampling_params, |
| 172 | ) |
| 173 | |
| 174 | # print the generated text |
| 175 | if args.print_output: |
| 176 | for i, output in enumerate(outputs): |
| 177 | print("-" * 50) |
| 178 | if not args.custom_mm_prompts: |
| 179 | print(f"prompt: {prompts[i].prompt}") |
| 180 | else: |
| 181 | print(f"prompt: {prompts[i]}") |
| 182 | print(f"generated text: {output.outputs[0].text}") |
| 183 | print("-" * 50) |
| 184 | |
| 185 | metrics = llm.get_metrics() |
| 186 | |
| 187 | total_num_output_tokens = sum( |
| 188 | len(output.outputs[0].token_ids) for output in outputs |
| 189 | ) |
| 190 | num_drafts = 0 |
| 191 | num_draft_tokens = 0 |
| 192 | num_accepted_tokens = 0 |
| 193 | acceptance_counts = [0] * args.num_spec_tokens |
| 194 | for metric in metrics: |
| 195 | if metric.name == "vllm:spec_decode_num_drafts": |
| 196 | assert isinstance(metric, Counter) |
| 197 | num_drafts += metric.value |
| 198 | elif metric.name == "vllm:spec_decode_num_draft_tokens": |
| 199 | assert isinstance(metric, Counter) |
| 200 | num_draft_tokens += metric.value |
| 201 | elif metric.name == "vllm:spec_decode_num_accepted_tokens": |
| 202 | assert isinstance(metric, Counter) |
| 203 | num_accepted_tokens += metric.value |
| 204 | elif metric.name == "vllm:spec_decode_num_accepted_tokens_per_pos": |
| 205 | assert isinstance(metric, Vector) |
| 206 | for pos in range(len(metric.values)): |
| 207 | acceptance_counts[pos] += metric.values[pos] |
| 208 | |
| 209 | print("-" * 50) |
| 210 | print(f"total_num_output_tokens: {total_num_output_tokens}") |
| 211 | print(f"num_drafts: {num_drafts}") |
| 212 | print(f"num_draft_tokens: {num_draft_tokens}") |
| 213 | print(f"num_accepted_tokens: {num_accepted_tokens}") |
| 214 | acceptance_length = 1 + (num_accepted_tokens / num_drafts) if num_drafts > 0 else 1 |
| 215 | print(f"mean acceptance length: {acceptance_length:.2f}") |
| 216 | print("-" * 50) |
| 217 | |
| 218 | # print acceptance at each token position |
| 219 | for i in range(len(acceptance_counts)): |
| 220 | acceptance_rate = acceptance_counts[i] / num_drafts if num_drafts > 0 else 0 |
| 221 | print(f"acceptance at token {i}: {acceptance_rate:.2f}") |
| 222 | |
| 223 | return acceptance_length |
| 224 | |
| 225 | |
| 226 | if __name__ == "__main__": |
| 227 | args = parse_args() |
| 228 | args.enable_multimodal_chat = args.backend == "openai-chat" |
| 229 | |
| 230 | acceptance_length = main(args) |
| 231 | |
| 232 | if args.test: |
| 233 | # takes ~30s to run on 1xH100 |
| 234 | assert args.method in ["eagle", "eagle3"] |
| 235 | assert args.tp == 1 |
| 236 | assert args.num_spec_tokens == 3 |
| 237 | assert args.dataset_name == "hf" |
| 238 | assert args.dataset_path == "philschmid/mt-bench" |
| 239 | assert args.num_prompts == 80 |
| 240 | assert args.temp == 0 |
| 241 | assert args.top_p == 1.0 |
| 242 | assert args.top_k == -1 |
| 243 | assert args.enable_chunked_prefill |
| 244 | |
| 245 | # check acceptance length is within 2% of expected value |
| 246 | rtol = 0.02 |
| 247 | expected_acceptance_length = 2.296 if args.method == "eagle" else 2.811 |
| 248 | |
| 249 | assert ( |
| 250 | acceptance_length <= (1 + rtol) * expected_acceptance_length |
| 251 | and acceptance_length >= (1 - rtol) * expected_acceptance_length |
| 252 | ), ( |
| 253 | f"acceptance_length {acceptance_length} is not " |
| 254 | f"within {rtol * 100}% of {expected_acceptance_length}" |
| 255 | ) |
| 256 | |
| 257 | print( |
| 258 | f"Test passed! Expected AL: " |
| 259 | f"{expected_acceptance_length}, got {acceptance_length}" |
| 260 | ) |