| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import random |
| 5 | from collections.abc import Iterable, Sequence |
| 6 | from typing import Any |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_offline |
| 12 | from vllm import LLM, SamplingParams |
| 13 | from vllm.assets.base import VLLM_S3_BUCKET_URL |
| 14 | from vllm.assets.image import VLM_IMAGES_DIR |
| 15 | from vllm.outputs import RequestOutput |
| 16 | from vllm.utils.torch_utils import set_random_seed |
| 17 | from vllm.v1.metrics.reader import Metric |
| 18 | |
| 19 | |
| 20 | def _skip_if_insufficient_gpus_for_tp(tp_size: int): |
| 21 | """Skip if fewer than ``tp_size`` accelerators are available.""" |
| 22 | available_gpus = torch.accelerator.device_count() |
| 23 | if available_gpus < tp_size: |
| 24 | pytest.skip( |
| 25 | f"Test requires {tp_size} GPUs, but only {available_gpus} available" |
| 26 | ) |
| 27 | |
| 28 | |
| 29 | Messages = list[dict[str, Any]] |
| 30 | |
| 31 | |
| 32 | def get_test_prompts(mm_enabled: bool, num_prompts: int = 100) -> list[Messages]: |
| 33 | prompt_types = ["repeat", "gsm8k"] |
| 34 | if mm_enabled: |
| 35 | prompt_types.append("mm") |
| 36 | prompts: list[Messages] = [] |
| 37 | |
| 38 | num_repeat_prompts = num_prompts // len(prompt_types) |
| 39 | if mm_enabled: |
| 40 | num_gsm8k_prompts = num_prompts // len(prompt_types) |
| 41 | num_mm_prompts = num_prompts - num_repeat_prompts - num_gsm8k_prompts |
| 42 | else: |
| 43 | num_mm_prompts = 0 |
| 44 | num_gsm8k_prompts = num_prompts - num_repeat_prompts |
| 45 | |
| 46 | # Generate a mixed batch of prompts, some of which can be easily |
| 47 | # predicted by n-gram matching and some which likely cannot. |
| 48 | set_random_seed(0) |
| 49 | for _ in range(num_repeat_prompts): |
| 50 | word_choices = ["test", "temp", "hello", "where"] |
| 51 | word = random.choice(word_choices) |
| 52 | prompts.append( |
| 53 | [ |
| 54 | { |
| 55 | "role": "user", |
| 56 | "content": f""" |
| 57 | please repeat the word '{word}' 10 times. |
| 58 | give no other output than the word at least ten times in a row, |
| 59 | in lowercase with spaces between each word and without quotes. |
| 60 | """, |
| 61 | } |
| 62 | ] |
| 63 | ) |
| 64 | prompts.extend( |
| 65 | [{"role": "user", "content": prompt}] |
| 66 | for prompt in _build_gsm8k_prompts( |
| 67 | num_questions=num_gsm8k_prompts, num_shots=5 |
| 68 | )[0] |
| 69 | ) |
| 70 | for _ in range(num_mm_prompts): |
| 71 | placeholders = [ |
| 72 | { |
| 73 | "type": "image_url", |
| 74 | "image_url": { |
| 75 | "url": f"{VLLM_S3_BUCKET_URL}/{VLM_IMAGES_DIR}/stop_sign.jpg" |
| 76 | }, |
| 77 | } |
| 78 | ] |
| 79 | prompt = [ |
| 80 | *placeholders, |
| 81 | {"type": "text", "text": "The meaning of the image is"}, |
| 82 | ] |
| 83 | prompts.append([{"role": "user", "content": prompt}]) |
| 84 | |
| 85 | return prompts |
| 86 | |
| 87 | |
| 88 | def get_instruct_coder_messages(n: int) -> list[Messages]: |
| 89 | from vllm.benchmarks.datasets import InstructCoderDataset |
| 90 | |
| 91 | dataset = InstructCoderDataset( |
| 92 | dataset_path="likaixin/InstructCoder", dataset_split="train" |
| 93 | ) |
| 94 | prompts: Iterable[str] = dataset.sample_prompts(n=n) |
| 95 | return [[{"role": "user", "content": prompt}] for prompt in prompts] |
| 96 | |
| 97 | |
| 98 | def greedy_sampling() -> SamplingParams: |
| 99 | return SamplingParams(temperature=0, max_tokens=10, ignore_eos=False) |
| 100 | |
| 101 | |
| 102 | def stochastic_sampling() -> SamplingParams: |
| 103 | return SamplingParams(temperature=1.0, max_tokens=10, ignore_eos=False) |
| 104 | |
| 105 | |
| 106 | def evaluate_llm_for_gsm8k(llm: LLM, expected_accuracy_threshold: float = 0.70) -> None: |
| 107 | """Evaluate the LLM on GSM8K and check that accuracy is above a sanity threshold. |
| 108 | |
| 109 | The default threshold assumes the LLM uses the same target model as the "model_name" |
| 110 | fixture, with max model len == 4096. Precomputed reference value is 75% to 80% |
| 111 | on GSM8K with greedy decoding, so we check that it's above a sanity threshold of 70% |
| 112 | to verify that the model is correct. |
| 113 | """ |
| 114 | if expected_accuracy_threshold <= 0.0: |
| 115 | print("Skipping GSM8K evaluation") |
| 116 | return |
| 117 | results = evaluate_gsm8k_offline(llm) |
| 118 | accuracy = results["accuracy"] |
| 119 | print(f"GSM8K accuracy: {accuracy:.3f}") |
| 120 | assert accuracy >= expected_accuracy_threshold, ( |
| 121 | f"Expected GSM8K accuracy >= {expected_accuracy_threshold}, got {accuracy:.3f}" |
| 122 | ) |
| 123 | |
| 124 | |
| 125 | def assert_request_outputs_match( |
| 126 | ref_outputs: Sequence[RequestOutput], |
| 127 | spec_outputs: Sequence[RequestOutput], |
| 128 | *, |
| 129 | required_matches: int, |
| 130 | context: str, |
| 131 | max_mismatches: int = 3, |
| 132 | ) -> None: |
| 133 | """Assert a thresholded exact-text match with bounded failure diagnostics.""" |
| 134 | assert ref_outputs, f"{context}: no reference outputs" |
| 135 | assert len(ref_outputs) == len(spec_outputs), ( |
| 136 | f"{context}: output count differs: " |
| 137 | f"reference={len(ref_outputs)}, speculative={len(spec_outputs)}" |
| 138 | ) |
| 139 | assert 0 <= required_matches <= len(ref_outputs), ( |
| 140 | f"{context}: invalid required_matches={required_matches} for " |
| 141 | f"{len(ref_outputs)} outputs" |
| 142 | ) |
| 143 | |
| 144 | mismatches: list[str] = [] |
| 145 | matches = 0 |
| 146 | for index, (ref_output, spec_output) in enumerate(zip(ref_outputs, spec_outputs)): |
| 147 | assert ref_output.outputs, ( |
| 148 | f"{context}: reference output {index} has no candidate" |
| 149 | ) |
| 150 | assert spec_output.outputs, ( |
| 151 | f"{context}: speculative output {index} has no candidate" |
| 152 | ) |
| 153 | ref_candidate = ref_output.outputs[0] |
| 154 | spec_candidate = spec_output.outputs[0] |
| 155 | if ref_candidate.text == spec_candidate.text: |
| 156 | matches += 1 |
| 157 | elif len(mismatches) < max_mismatches: |
| 158 | mismatches.append( |
| 159 | f"[{index}] ref_text={ref_candidate.text[:240]!r}, " |
| 160 | f"spec_text={spec_candidate.text[:240]!r}\n" |
| 161 | f" ref_token_ids={list(ref_candidate.token_ids)[:64]}\n" |
| 162 | f" spec_token_ids={list(spec_candidate.token_ids)[:64]}" |
| 163 | ) |
| 164 | |
| 165 | print( |
| 166 | f"{context}: exact text matches={matches}/{len(ref_outputs)} " |
| 167 | f"(required={required_matches})" |
| 168 | ) |
| 169 | mismatch_summary = "\n".join(mismatches) or "no mismatches captured" |
| 170 | assert matches >= required_matches, ( |
| 171 | f"{context}: only {matches}/{len(ref_outputs)} outputs matched; " |
| 172 | f"required at least {required_matches}. First mismatches:\n" |
| 173 | f"{mismatch_summary}" |
| 174 | ) |
| 175 | |
| 176 | |
| 177 | def get_spec_decode_metric_value(metrics: Sequence[Metric], metric_name: str) -> float: |
| 178 | """Get a spec-decode metric with an actionable error when stats are absent.""" |
| 179 | name2metric = {metric.name: metric for metric in metrics} |
| 180 | metric = name2metric.get(metric_name) |
| 181 | assert metric is not None, ( |
| 182 | f"Missing metric {metric_name!r}. Ensure disable_log_stats=False. " |
| 183 | "Available spec-decode metrics: " |
| 184 | f"{sorted(name for name in name2metric if 'spec_decode' in name) or ['<none>']}" |
| 185 | ) |
| 186 | return float(metric.value) |
| 187 | |
| 188 | |
| 189 | def compute_acceptance_rate( |
| 190 | metrics: list[Metric], prev_metrics: list[Metric] | None = None |
| 191 | ) -> float: |
| 192 | n_draft_toks = get_spec_decode_metric_value( |
| 193 | metrics, "vllm:spec_decode_num_draft_tokens" |
| 194 | ) |
| 195 | if n_draft_toks == 0: |
| 196 | return float("nan") |
| 197 | n_accepted_toks = get_spec_decode_metric_value( |
| 198 | metrics, "vllm:spec_decode_num_accepted_tokens" |
| 199 | ) |
| 200 | if prev_metrics is not None: |
| 201 | n_draft_toks -= get_spec_decode_metric_value( |
| 202 | prev_metrics, "vllm:spec_decode_num_draft_tokens" |
| 203 | ) |
| 204 | n_accepted_toks -= get_spec_decode_metric_value( |
| 205 | prev_metrics, "vllm:spec_decode_num_accepted_tokens" |
| 206 | ) |
| 207 | if n_draft_toks <= 0: |
| 208 | return float("nan") |
| 209 | return n_accepted_toks / n_draft_toks |
| 210 | |
| 211 | |
| 212 | def compute_acceptance_len( |
| 213 | metrics: list[Metric], prev_metrics: list[Metric] | None = None |
| 214 | ) -> float: |
| 215 | n_drafts = get_spec_decode_metric_value(metrics, "vllm:spec_decode_num_drafts") |
| 216 | n_accepted_toks = get_spec_decode_metric_value( |
| 217 | metrics, "vllm:spec_decode_num_accepted_tokens" |
| 218 | ) |
| 219 | if n_drafts == 0: |
| 220 | return 1 |
| 221 | if prev_metrics is not None: |
| 222 | n_drafts -= get_spec_decode_metric_value( |
| 223 | prev_metrics, "vllm:spec_decode_num_drafts" |
| 224 | ) |
| 225 | n_accepted_toks -= get_spec_decode_metric_value( |
| 226 | prev_metrics, "vllm:spec_decode_num_accepted_tokens" |
| 227 | ) |
| 228 | if n_drafts <= 0: |
| 229 | return 1 |
| 230 | return 1 + (n_accepted_toks / n_drafts) |