A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[Core][Bugfix] Fix Offline MM Beam Search (#16390)
Signed-off-by: Alex-Brooks <Alex.Brooks@ibm.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
Alex Brooks committed
Apr 14, 2025 at 20:33 UTC
6b40996ae8a5b065de5b3b650b5b1324f67f6334
4 files changed
+140
-30
tests/conftest.py
+20
-12
@@ -29,12 +29,11 @@ from vllm.distributed import (cleanup_dist_env_and_memory,
29
init_distributed_environment,
30
initialize_model_parallel)
31
from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
32
- TokensPrompt, to_enc_dec_tuple_list,
33
- zip_enc_dec_prompts)
32
+ to_enc_dec_tuple_list, zip_enc_dec_prompts)
33
from vllm.logger import init_logger
34
from vllm.outputs import RequestOutput
35
from vllm.sampling_params import BeamSearchParams
37
-from vllm.utils import cuda_device_count_stateless, is_list_of
36
+from vllm.utils import cuda_device_count_stateless
37
38
logger = init_logger(__name__)
39
@@ -469,12 +468,19 @@ class HfRunner:
468
prompts: list[str],
469
beam_width: int,
470
max_tokens: int,
471
+ images: Optional[PromptImageInput] = None,
472
+ videos: Optional[PromptVideoInput] = None,
473
+ audios: Optional[PromptAudioInput] = None,
474
) -> list[tuple[list[list[int]], list[str]]]:
475
outputs = self.generate(prompts,
476
do_sample=False,
477
max_new_tokens=max_tokens,
478
num_beams=beam_width,
477
- num_return_sequences=beam_width)
479
+ num_return_sequences=beam_width,
480
+ images=images,
481
+ videos=videos,
482
+ audios=audios)
483
+
484
for i in range(len(outputs)):
485
output_ids, output_str = outputs[i]
486
for j in range(len(output_ids)):
@@ -936,18 +942,20 @@ class VllmRunner:
942
943
def generate_beam_search(
944
self,
939
- prompts: Union[list[str], list[list[int]]],
945
+ prompts: list[str],
946
beam_width: int,
947
max_tokens: int,
948
+ images: Optional[PromptImageInput] = None,
949
+ videos: Optional[PromptVideoInput] = None,
950
+ audios: Optional[PromptAudioInput] = None,
951
) -> list[tuple[list[list[int]], list[str]]]:
943
- if is_list_of(prompts, str, check="all"):
944
- prompts = [TextPrompt(prompt=prompt) for prompt in prompts]
945
- else:
946
- prompts = [
947
- TokensPrompt(prompt_token_ids=tokens) for tokens in prompts
948
- ]
952
+ inputs = self.get_inputs(prompts,
953
+ images=images,
954
+ videos=videos,
955
+ audios=audios)
956
+
957
outputs = self.model.beam_search(
950
- prompts,
958
+ inputs,
959
BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens))
960
returned_outputs = []
961
for output in outputs:
tests/samplers/test_beam_search.py
+82
-3
@@ -5,6 +5,9 @@ Run `pytest tests/samplers/test_beam_search.py`.
5
"""
6
7
import pytest
8
+from transformers import AutoModelForSeq2SeqLM
9
+
10
+from vllm.assets.audio import AudioAsset
11
12
13
@pytest.fixture(autouse=True)
@@ -19,6 +22,7 @@ def v1(run_with_both_engines):
22
# 3. Use the model "huggyllama/llama-7b".
23
MAX_TOKENS = [64]
24
BEAM_WIDTHS = [4]
25
+MM_BEAM_WIDTHS = [2]
26
MODELS = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"]
27
28
@@ -48,15 +52,90 @@ def test_beam_search_single_input(
52
for i in range(len(example_prompts)):
53
hf_output_ids, hf_output_texts = hf_outputs[i]
54
vllm_output_ids, vllm_output_texts = vllm_outputs[i]
51
- for i, (hf_text,
55
+ for j, (hf_text,
56
vllm_text) in enumerate(zip(hf_output_texts,
57
vllm_output_texts)):
54
- print(f">>>{i}-th hf output:")
58
+ print(f">>>{j}-th hf output:")
59
print(hf_text)
56
- print(f">>>{i}-th vllm output:")
60
+ print(f">>>{j}-th vllm output:")
61
print(vllm_text)
62
assert len(hf_output_ids) == len(vllm_output_ids)
63
for j in range(len(hf_output_ids)):
64
assert hf_output_ids[j] == vllm_output_ids[j], (
65
f"Test{i} output{j}:\nHF: {hf_output_ids}\n"
66
f"vLLM: {vllm_output_ids}")
67
+
68
+
69
+@pytest.mark.parametrize("dtype", ["half"])
70
+@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
71
+@pytest.mark.parametrize("beam_width", MM_BEAM_WIDTHS)
72
+def test_beam_search_passes_multimodal_data(
73
+ hf_runner,
74
+ vllm_runner,
75
+ dtype: str,
76
+ max_tokens: int,
77
+ beam_width: int,
78
+) -> None:
79
+ """Ensure that beam search passes multimodal data through correctly."""
80
+ # NOTE - this test is primarily to check that mm data is passed to beams
81
+ # correctly. As such, we just need to check one extra modality to make
82
+ # sure things pass through properly.
83
+ audios = [AudioAsset("mary_had_lamb").audio_and_sample_rate]
84
+ model = "Qwen/Qwen2-Audio-7B-Instruct"
85
+ audio_seq = "<|audio_bos|><|AUDIO|><|audio_eos|>"
86
+ prompts = [
87
+ f"<|im_start|>user\n{audio_seq}Can you transcribe this?<|im_end|>\n<|im_start|>assistant\n" #noqa: E501
88
+ ]
89
+
90
+ with hf_runner(model, dtype=dtype,
91
+ auto_cls=AutoModelForSeq2SeqLM) as hf_model:
92
+ audio_token_id = hf_model.config.audio_token_index
93
+ eos_token_id = hf_model.tokenizer.eos_token_id # <|im_end|>
94
+ hf_outputs = hf_model.generate_beam_search(
95
+ prompts,
96
+ beam_width=beam_width,
97
+ max_tokens=max_tokens,
98
+ audios=audios,
99
+ )
100
+
101
+ with vllm_runner(model, dtype=dtype) as vllm_model:
102
+ vllm_outputs = vllm_model.generate_beam_search(
103
+ prompts,
104
+ beam_width=beam_width,
105
+ max_tokens=max_tokens,
106
+ audios=audios,
107
+ )
108
+
109
+ seq_with_no_audio_toks = lambda seq: [
110
+ tok for tok in seq if tok != audio_token_id
111
+ ]
112
+
113
+ for i in range(len(prompts)):
114
+ hf_output_ids, hf_output_texts = hf_outputs[i]
115
+ vllm_output_ids, vllm_output_texts = vllm_outputs[i]
116
+
117
+ for j, (hf_text,
118
+ vllm_text) in enumerate(zip(hf_output_texts,
119
+ vllm_output_texts)):
120
+ print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:")
121
+ print(hf_text)
122
+ print(f">>>{j}-th vllm output:")
123
+ print(vllm_text)
124
+ assert len(hf_output_ids) == len(vllm_output_ids)
125
+
126
+ for j in range(len(hf_output_ids)):
127
+ # Compare everything except for the audio tokens; we do this since
128
+ # the IDs returned from the transformers helper expands the audio
129
+ # token to match features, while the vLLM helper maintains the
130
+ # single audio token in the input text
131
+ filtered_hf_output_ids = seq_with_no_audio_toks(hf_output_ids[j])
132
+ filtered_vllm_output_ids = seq_with_no_audio_toks(
133
+ vllm_output_ids[j])
134
+
135
+ # HF output IDs may contain the end of sequence
136
+ if len(filtered_hf_output_ids
137
+ ) == len(filtered_vllm_output_ids) + 1:
138
+ assert filtered_hf_output_ids[-1] == eos_token_id
139
+ filtered_hf_output_ids = filtered_hf_output_ids[:-1]
140
+
141
+ assert filtered_hf_output_ids == filtered_vllm_output_ids
vllm/beam_search.py
+11
-2
@@ -38,9 +38,18 @@ class BeamSearchOutput:
38
39
class BeamSearchInstance:
40
41
- def __init__(self, prompt_tokens: list[int]):
41
+ def __init__(
42
+ self,
43
+ prompt_tokens: list[int],
44
+ logprobs: Optional[list[dict[int, Logprob]]] = None,
45
+ **kwargs,
46
+ ):
47
self.beams: list[BeamSearchSequence] = [
43
- BeamSearchSequence(tokens=prompt_tokens, logprobs=[])
48
+ BeamSearchSequence(
49
+ tokens=prompt_tokens,
50
+ logprobs=[] if logprobs is None else list(logprobs),
51
+ **kwargs,
52
+ )
53
]
54
self.completed: list[BeamSearchSequence] = []
55
vllm/entrypoints/llm.py
+27
-13
@@ -536,15 +536,18 @@ class LLM:
536
tokenizer.eos_token_id,
537
length_penalty)
538
539
- # TODO - fix handling of multimodal data for beam search; we pass it
540
- # through in the async version on the abstract EngineClient, but not
541
- # here.
542
- if any("multi_modal_data" in prompt
543
- and prompt["multi_modal_data"] is not None
544
- for prompt in prompts):
545
- logger.warning(
546
- "Multimodal data appears to have been provided, but is not"
547
- " currently being passed through in LLM.beam_search()!")
539
+ def create_tokens_prompt_from_beam(
540
+ beam: BeamSearchSequence) -> TokensPrompt:
541
+ token_prompt_kwargs: TokensPrompt = {
542
+ "prompt_token_ids": beam.tokens
543
+ }
544
+ if beam.multi_modal_data is not None:
545
+ token_prompt_kwargs["multi_modal_data"] = beam.multi_modal_data
546
+
547
+ if beam.mm_processor_kwargs is not None:
548
+ token_prompt_kwargs[
549
+ "mm_processor_kwargs"] = beam.mm_processor_kwargs
550
+ return TokensPrompt(**token_prompt_kwargs)
551
552
tokenizer = self.get_tokenizer()
553
# generate 2 * beam_width candidates at each step
@@ -556,11 +559,20 @@ class LLM:
559
instances: list[BeamSearchInstance] = []
560
561
for prompt in prompts:
562
+ # Add multimodal processor kwargs & data
563
+ mm_kwargs = {}
564
+ if "multi_modal_data" in prompt:
565
+ mm_kwargs["multi_modal_data"] = prompt["multi_modal_data"]
566
+ if "mm_processor_kwargs" in prompt:
567
+ mm_kwargs["mm_processor_kwargs"] = prompt[
568
+ "mm_processor_kwargs"]
569
+
570
if is_token_prompt(prompt):
571
prompt_tokens = prompt["prompt_token_ids"]
572
else:
573
prompt_tokens = tokenizer.encode(prompt["prompt"])
563
- instances.append(BeamSearchInstance(prompt_tokens))
574
+ instances.append(
575
+ BeamSearchInstance(prompt_tokens, logprobs=None, **mm_kwargs))
576
577
for _ in range(max_tokens):
578
all_beams: list[BeamSearchSequence] = list(
@@ -575,8 +587,7 @@ class LLM:
587
break
588
589
prompts_batch = [
578
- TokensPrompt(prompt_token_ids=beam.tokens)
579
- for beam in all_beams
590
+ create_tokens_prompt_from_beam(beam) for beam in all_beams
591
]
592
593
# only runs for one step
@@ -602,7 +613,10 @@ class LLM:
613
tokens=current_beam.tokens + [token_id],
614
logprobs=current_beam.logprobs + [logprobs],
615
cum_logprob=current_beam.cum_logprob +
605
- logprob_obj.logprob)
616
+ logprob_obj.logprob,
617
+ multi_modal_data=current_beam.multi_modal_data,
618
+ mm_processor_kwargs=current_beam.
619
+ mm_processor_kwargs)
620
621
if token_id == tokenizer.eos_token_id and \
622
not ignore_eos: