@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 155 lines 4.65 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5
6 from vllm.assets.audio import AudioAsset
7 from vllm.model_executor.models.moss_audio import MOSS_AUDIO_PLACEHOLDER
8 from vllm.platforms import current_platform
9
10 from ...registry import HF_EXAMPLE_MODELS
11 from ...utils import check_logprobs_close
12
13 CORE_MODEL = pytest.param(
14 "OpenMOSS-Team/MOSS-Audio-4B-Instruct",
15 marks=pytest.mark.core_model,
16 id="4b-instruct",
17 )
18
19 EXTENDED_MODELS = [
20 "OpenMOSS-Team/MOSS-Audio-4B-Thinking",
21 "OpenMOSS-Team/MOSS-Audio-8B-Instruct",
22 "OpenMOSS-Team/MOSS-Audio-8B-Thinking",
23 ]
24
25 ACCURACY_MODELS = [CORE_MODEL, *EXTENDED_MODELS]
26
27 PARALLEL_SMOKE_CASES = [
28 pytest.param({"tensor_parallel_size": 2}, id="tp2"),
29 pytest.param({"pipeline_parallel_size": 2}, id="pp2"),
30 pytest.param(
31 {"tensor_parallel_size": 2, "pipeline_parallel_size": 2},
32 id="tp2_pp2",
33 ),
34 ]
35
36 HF_ACCURACY_SKIP_REASON = (
37 "HF AutoModelForCausalLM cannot load remote MOSS-Audio configs; "
38 "vLLM generation coverage is provided by the smoke tests below."
39 )
40
41
42 @pytest.mark.core_model
43 def test_moss_audio_generation_smoke(vllm_runner) -> None:
44 model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
45 model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
46 model_info.check_available_online(on_fail="skip")
47 model_info.check_transformers_version(on_fail="skip")
48
49 prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nBriefly describe this audio."]
50 audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
51
52 with vllm_runner(
53 model,
54 dtype="half",
55 enforce_eager=True,
56 max_model_len=1024,
57 limit_mm_per_prompt={"audio": 1},
58 trust_remote_code=True,
59 ) as vllm_model:
60 outputs = vllm_model.generate_greedy(
61 prompts,
62 max_tokens=4,
63 audios=audios,
64 )
65
66 assert len(outputs) == 1
67 assert len(outputs[0][1]) > 0
68
69
70 @pytest.mark.skip(reason=HF_ACCURACY_SKIP_REASON)
71 @pytest.mark.parametrize("model", ACCURACY_MODELS)
72 @pytest.mark.parametrize("dtype", ["half"])
73 @pytest.mark.parametrize("max_tokens", [8])
74 @pytest.mark.parametrize("num_logprobs", [5])
75 def test_moss_audio_hf_vllm_accuracy(
76 hf_runner,
77 vllm_runner,
78 model: str,
79 dtype: str,
80 max_tokens: int,
81 num_logprobs: int,
82 ) -> None:
83 model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
84 model_info.check_available_online(on_fail="skip")
85 model_info.check_transformers_version(on_fail="skip")
86
87 prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nTranscribe this audio."]
88 audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
89
90 with vllm_runner(
91 model,
92 dtype=dtype,
93 enforce_eager=True,
94 max_model_len=1024,
95 limit_mm_per_prompt={"audio": 1},
96 trust_remote_code=True,
97 ) as vllm_model:
98 vllm_outputs = vllm_model.generate_greedy_logprobs(
99 prompts,
100 max_tokens,
101 num_logprobs=num_logprobs,
102 audios=audios,
103 )
104
105 with hf_runner(model, dtype=dtype, trust_remote_code=True) as hf_model:
106 hf_outputs = hf_model.generate_greedy_logprobs_limit(
107 prompts,
108 max_tokens,
109 num_logprobs=num_logprobs,
110 audios=audios,
111 )
112
113 check_logprobs_close(
114 outputs_0_lst=hf_outputs,
115 outputs_1_lst=vllm_outputs,
116 name_0="hf",
117 name_1="vllm",
118 )
119
120
121 @pytest.mark.core_model
122 @pytest.mark.parametrize("parallel_kwargs", PARALLEL_SMOKE_CASES)
123 def test_moss_audio_parallel_smoke(vllm_runner, parallel_kwargs) -> None:
124 model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
125 required_gpus = parallel_kwargs.get(
126 "tensor_parallel_size", 1
127 ) * parallel_kwargs.get("pipeline_parallel_size", 1)
128 if current_platform.device_count() < required_gpus:
129 # TP/PP integration smoke runs on local or multi-GPU CI only.
130 pytest.skip(f"Requires at least {required_gpus} GPUs")
131
132 model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
133 model_info.check_available_online(on_fail="skip")
134 model_info.check_transformers_version(on_fail="skip")
135
136 prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nBriefly describe this audio."]
137 audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
138
139 with vllm_runner(
140 model,
141 dtype="half",
142 enforce_eager=True,
143 max_model_len=1024,
144 limit_mm_per_prompt={"audio": 1},
145 trust_remote_code=True,
146 **parallel_kwargs,
147 ) as vllm_model:
148 outputs = vllm_model.generate_greedy(
149 prompts,
150 max_tokens=4,
151 audios=audios,
152 )
153
154 assert len(outputs) == 1
155 assert len(outputs[0][1]) > 0