test_granite_speech.py 5.67 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
7
8

from collections.abc import Sequence

import pytest
from transformers import AutoModelForSpeechSeq2Seq

9
from vllm.logprobs import SampleLogprobs
10
from vllm.lora.request import LoRARequest
11
from vllm.platforms import current_platform
12

13
from ....conftest import AudioTestAssets, HfRunner, PromptAudioInput, VllmRunner
14
15
16
17
18
19
20
from ...registry import HF_EXAMPLE_MODELS
from ...utils import check_logprobs_close

HF_AUDIO_PROMPT = "<|start_of_role|>system<|end_of_role|>Knowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant<|end_of_text|>\n<|start_of_role|>user<|end_of_role|><|audio|>can you transcribe the speech into a written format?<|end_of_text|>\n<|start_of_role|>assistant<|end_of_role|>"  # noqa: E501


def vllm_to_hf_output(
21
22
    vllm_output: tuple[list[int], str, SampleLogprobs | None],
) -> tuple[list[int], str, SampleLogprobs | None]:
23
24
25
26
27
28
29
30
    """Sanitize hf output to be comparable with vllm output."""
    output_ids, output_str, out_logprobs = vllm_output

    hf_output_str = output_str + "<|end_of_text|>"

    return output_ids, hf_output_str, out_logprobs


31
MODEL_NAME = "ibm-granite/granite-speech-3.3-2b"
32
33
34
35
36
37
# Audio lora co-exists directly in the model directory, but
# currently still needs to be passed directly to vLLM.
audio_lora_path = MODEL_NAME
models = [MODEL_NAME]


38
39
40
@pytest.fixture
def granite_speech_attention_config():
    """Return attention config for Granite Speech tests on ROCm."""
41
    if current_platform.is_rocm():
42
43
        return {"backend": "TRITON_ATTN"}
    return None
44
45


46
47
48
49
50
51
52
53
54
55
56
def run_test(
    hf_runner: type[HfRunner],
    vllm_runner: type[VllmRunner],
    inputs: Sequence[tuple[list[str], PromptAudioInput]],
    model: str,
    *,
    max_model_len: int,
    dtype: str,
    max_tokens: int,
    num_logprobs: int,
    tensor_parallel_size: int,
57
    distributed_executor_backend: str | None = None,
58
    attention_config: dict | None = None,
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
):
    """Inference result should be the same between hf and vllm.

    All the audio fixtures for the test are from AUDIO_ASSETS.
    For huggingface runner, we provide the audio as input.
    For vllm runner, we provide MultiModalDataDict objects
    and corresponding MultiModalConfig as input.
    Note, the text input is also adjusted to abide by vllm contract.
    The text output is sanitized to be able to compare with hf.
    """
    # NOTE: take care of the order. run vLLM first, and then run HF.
    # vLLM needs a fresh new process without cuda initialization.
    # if we run HF first, the cuda initialization will be done and it
    # will hurt multiprocessing backend with fork method (the default method).
    # max_model_len should be greater than image_feature_size
    with vllm_runner(
75
76
77
78
79
80
81
82
83
84
85
        model,
        runner="generate",
        max_model_len=max_model_len,
        max_num_seqs=1,
        dtype=dtype,
        limit_mm_per_prompt={"audio": 1},
        tensor_parallel_size=tensor_parallel_size,
        distributed_executor_backend=distributed_executor_backend,
        enable_lora=True,
        max_lora_rank=64,
        enforce_eager=True,
86
        attention_config=attention_config,
87
88
89
    ) as vllm_model:
        lora_request = LoRARequest("audio", 1, audio_lora_path)
        vllm_outputs_per_case = [
90
91
92
93
94
95
96
            vllm_model.generate_greedy_logprobs(
                prompts,
                max_tokens,
                num_logprobs=num_logprobs,
                audios=audios,
                lora_request=lora_request,
            )
97
98
99
            for prompts, audios in inputs
        ]

100
    with hf_runner(model, dtype=dtype, auto_cls=AutoModelForSpeechSeq2Seq) as hf_model:
101
102
103
104
        hf_processor = hf_model.processor
        eos_token_id = hf_processor.tokenizer.eos_token_id

        hf_outputs_per_case = [
105
106
107
108
109
110
111
            hf_model.generate_greedy_logprobs_limit(
                prompts,
                max_tokens,
                num_logprobs=num_logprobs,
                audios=[audios],
                eos_token_id=eos_token_id,
            )
112
113
114
            for prompts, audios in inputs
        ]

115
    for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
116
117
        check_logprobs_close(
            outputs_0_lst=hf_outputs,
118
            outputs_1_lst=[vllm_to_hf_output(output) for output in vllm_outputs],
119
120
121
122
123
124
            name_0="hf",
            name_1="vllm",
        )


@pytest.mark.parametrize("model", models)
125
126
127
128
129
130
@pytest.mark.parametrize(
    "dtype", ["float16"] if current_platform.is_rocm() else ["bfloat16"]
)
@pytest.mark.parametrize(
    "max_model_len", [512] if current_platform.is_rocm() else [2048]
)
131
132
@pytest.mark.parametrize("max_tokens", [128])
@pytest.mark.parametrize("num_logprobs", [10])
133
134
135
136
137
def test_models(
    hf_runner,
    vllm_runner,
    model: str,
    audio_assets: AudioTestAssets,
138
    granite_speech_attention_config,
139
140
141
142
143
    dtype: str,
    max_model_len: int,
    max_tokens: int,
    num_logprobs: int,
) -> None:
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
    model_info.check_available_online(on_fail="skip")
    model_info.check_transformers_version(on_fail="skip")

    audio, sr = audio_assets[0].audio_and_sample_rate
    # This model expects 16k sample rate, which our test audio
    # already is; if this changes, it may break this test,
    # so we check it directly
    assert sr == 16000
    run_test(
        hf_runner,
        vllm_runner,
        [
            ([HF_AUDIO_PROMPT], [audio]),
        ],
        model,
        dtype=dtype,
        max_model_len=max_model_len,
        max_tokens=max_tokens,
        num_logprobs=num_logprobs,
        tensor_parallel_size=1,
165
        attention_config=granite_speech_attention_config,
166
    )