main
py 457 lines 17.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit tests for Qwen2.5-Omni embed_input_ids to verify embeddings are
4 correctly assigned to audio/image/video token positions.
5
6 Regression test for: https://github.com/vllm-project/vllm/issues/34506
7 - Non-interleaved mixed modalities (audio + image + video) should correctly
8 assign audio embeddings to audio positions, image to image, video to video.
9 - Interleaved (use_audio_in_video) should also work correctly.
10 """
11
12 from unittest.mock import Mock
13
14 import pytest
15 import torch
16
17 from vllm.model_executor.models.qwen2_5_omni_thinker import (
18 check_interleaved_audio_video,
19 merge_interleaved_embeddings,
20 )
21 from vllm.multimodal.utils import set_mm_embedding_modality
22
23 # Fake token IDs
24 AUDIO_TOKEN_ID = 1001
25 IMAGE_TOKEN_ID = 1002
26 VIDEO_TOKEN_ID = 1003
27 TEXT_TOKEN_ID = 0
28
29
30 def _mm_embed(shape: tuple[int, ...], value: float, modality: str) -> torch.Tensor:
31 return set_mm_embedding_modality(torch.full(shape, value), modality)
32
33
34 # ---------------------------------------------------------------------------
35 # Helpers
36 # ---------------------------------------------------------------------------
37
38
39 def make_token_seq(
40 audio_n: int, image_n: int, video_n: int, text_prefix: int = 3, text_sep: int = 2
41 ):
42 """Build a flat token sequence:
43 [text_prefix] [AUDIO * audio_n] [text_sep] [IMAGE * image_n]
44 [text_sep] [VIDEO * video_n] [text_sep]
45 Returns (input_ids tensor, is_multimodal mask, positions dict).
46 """
47 tokens = (
48 [TEXT_TOKEN_ID] * text_prefix
49 + [AUDIO_TOKEN_ID] * audio_n
50 + [TEXT_TOKEN_ID] * text_sep
51 + [IMAGE_TOKEN_ID] * image_n
52 + [TEXT_TOKEN_ID] * text_sep
53 + [VIDEO_TOKEN_ID] * video_n
54 + [TEXT_TOKEN_ID] * text_sep
55 )
56 input_ids = torch.tensor(tokens)
57 is_multimodal = (
58 (input_ids == AUDIO_TOKEN_ID)
59 | (input_ids == IMAGE_TOKEN_ID)
60 | (input_ids == VIDEO_TOKEN_ID)
61 )
62 return input_ids, is_multimodal
63
64
65 def make_interleaved_seq(
66 video_chunks: list[int], audio_chunks: list[int], text_prefix: int = 2
67 ):
68 """Build an interleaved sequence like use_audio_in_video:
69 [text] [V*v0] [A*a0] [V*v1] [A*a1] ...
70 """
71 tokens = [TEXT_TOKEN_ID] * text_prefix
72 for v, a in zip(video_chunks, audio_chunks):
73 tokens += [VIDEO_TOKEN_ID] * v + [AUDIO_TOKEN_ID] * a
74 input_ids = torch.tensor(tokens)
75 is_multimodal = (input_ids == VIDEO_TOKEN_ID) | (input_ids == AUDIO_TOKEN_ID)
76 return input_ids, is_multimodal
77
78
79 # ---------------------------------------------------------------------------
80 # Tests for check_interleaved_audio_video
81 # ---------------------------------------------------------------------------
82
83
84 class TestCheckInterleavedAudioVideo:
85 def test_non_interleaved_audio_then_video(self):
86 """Audio entirely before video → not interleaved."""
87 input_ids, is_multimodal = make_token_seq(5, 0, 4)
88 is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
89 is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
90 assert not check_interleaved_audio_video(
91 is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
92 )
93
94 def test_non_interleaved_with_image(self):
95 """Audio + image + video (the mixed_modalities case) → not interleaved."""
96 input_ids, is_multimodal = make_token_seq(5, 4, 6)
97 is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
98 is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
99 assert not check_interleaved_audio_video(
100 is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
101 )
102
103 def test_no_audio(self):
104 """Video only → not interleaved."""
105 input_ids, is_multimodal = make_token_seq(0, 0, 6)
106 is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
107 is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
108 assert not check_interleaved_audio_video(
109 is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
110 )
111
112 def test_interleaved(self):
113 """V A V A interleaved → True."""
114 input_ids, is_multimodal = make_interleaved_seq([4, 4], [3, 3])
115 is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
116 is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
117 assert check_interleaved_audio_video(
118 is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
119 )
120
121 def test_multi_video_with_boundary_tokens(self):
122 """Two interleaved videos separated by boundary tokens → still True.
123
124 use_audio_in_video expands each video into a local V/A span bounded by
125 non-pad tokens. A global density check would fail across those gaps.
126 """
127 # [text][V A V A][text boundary][V A V A]
128 first_ids, _ = make_interleaved_seq([2, 2], [2, 2], text_prefix=1)
129 second_ids, _ = make_interleaved_seq([2, 2], [2, 2], text_prefix=0)
130 boundary = torch.tensor([TEXT_TOKEN_ID, TEXT_TOKEN_ID])
131 input_ids = torch.cat([first_ids, boundary, second_ids])
132 is_multimodal = (input_ids == VIDEO_TOKEN_ID) | (input_ids == AUDIO_TOKEN_ID)
133 is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
134 is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
135 assert check_interleaved_audio_video(
136 is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
137 )
138
139 def test_batched_non_interleaved_no_false_positive(self):
140 """Regression test for https://github.com/vllm-project/vllm/issues/35394.
141
142 5 identical non-interleaved mixed-modality requests batched together:
143 each has [audio][image][video] in separate blocks with text between them.
144 Across the batch, audio from request N falls between video blocks of
145 request N and request N+1, causing the global ranges to overlap.
146 check_interleaved_audio_video must return False (not a false positive).
147 """
148 # Build one request: [text][audio*5][text][image*4][text][video*6][text]
149 single_ids, _ = make_token_seq(5, 4, 6)
150 # Batch 5 identical requests (separated by text tokens to simulate padding)
151 sep = torch.tensor([TEXT_TOKEN_ID] * 3)
152 batched_ids = torch.cat([single_ids, sep] * 5)
153 is_multimodal = (
154 (batched_ids == AUDIO_TOKEN_ID)
155 | (batched_ids == IMAGE_TOKEN_ID)
156 | (batched_ids == VIDEO_TOKEN_ID)
157 )
158 is_video = is_multimodal & (batched_ids == VIDEO_TOKEN_ID)
159 is_audio = is_multimodal & (batched_ids == AUDIO_TOKEN_ID)
160 assert not check_interleaved_audio_video(
161 is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
162 ), "Batched non-interleaved requests should not be detected as interleaved"
163
164
165 # ---------------------------------------------------------------------------
166 # Tests for embed_input_ids via a minimal mock
167 # ---------------------------------------------------------------------------
168
169
170 def make_mock_model(hidden: int = 8):
171 """Return a minimal mock of Qwen2_5OmniThinkerForConditionalGeneration
172 that has enough structure to run embed_input_ids.
173 """
174 from vllm.model_executor.models.qwen2_5_omni_thinker import (
175 Qwen2_5OmniThinkerForConditionalGeneration,
176 )
177
178 model = Mock(spec=Qwen2_5OmniThinkerForConditionalGeneration)
179
180 # Config with token IDs
181 cfg = Mock()
182 cfg.video_token_index = VIDEO_TOKEN_ID
183 cfg.audio_token_index = AUDIO_TOKEN_ID
184 model.config = cfg
185
186 # embed_input_ids: simply embed each token as a one-hot-like vector
187 # token_id * ones so we can verify which embedding ends up where.
188 def fake_lm_embed(ids: torch.Tensor) -> torch.Tensor:
189 # Use .clone() so the tensor is contiguous (expand() creates a strided
190 # view with shared memory, which masked_scatter_ cannot handle).
191 return ids.float().unsqueeze(-1).expand(-1, hidden).clone()
192
193 lang_model = Mock()
194 lang_model.embed_input_ids = fake_lm_embed
195 model.get_language_model = Mock(return_value=lang_model)
196
197 # _embed_text_input_ids: delegate to SupportsMultiModal's implementation
198 from vllm.model_executor.models.interfaces import SupportsMultiModal
199
200 model._embed_text_input_ids = (
201 lambda *a, **kw: SupportsMultiModal._embed_text_input_ids(model, *a, **kw)
202 )
203
204 # super().embed_input_ids → use SupportsMultiModal.embed_input_ids
205 def fake_super_embed(
206 ids,
207 mm_embs=None,
208 *,
209 is_multimodal=None,
210 ):
211 return SupportsMultiModal.embed_input_ids(
212 model,
213 ids,
214 mm_embs,
215 is_multimodal=is_multimodal,
216 )
217
218 # Bind embed_input_ids as the real method
219 model.embed_input_ids = (
220 lambda *a, **kw: Qwen2_5OmniThinkerForConditionalGeneration.embed_input_ids(
221 model, *a, **kw
222 )
223 )
224
225 # Store super-embed for use inside the method
226 model._super_embed_input_ids = fake_super_embed
227
228 return model, hidden
229
230
231 def build_mm_embeds(
232 audio_n, image_n, video_n, hidden, audio_val=10.0, image_val=20.0, video_val=30.0
233 ):
234 """Build multimodal_embeddings list in position order (audio, image, video).
235 Each embedding is filled with a distinct constant so we can verify placement.
236 """
237 embs = []
238 if audio_n:
239 embs.append(torch.full((audio_n, hidden), audio_val))
240 if image_n:
241 embs.append(torch.full((image_n, hidden), image_val))
242 if video_n:
243 embs.append(torch.full((video_n, hidden), video_val))
244 return embs
245
246
247 class TestEmbedInputIds:
248 def _run(self, audio_n, image_n, video_n, hidden=8):
249 """Run embed_input_ids for a non-interleaved mixed-modality sequence.
250 Returns (result_embeds, input_ids, is_multimodal).
251 """
252 input_ids, is_multimodal = make_token_seq(audio_n, image_n, video_n)
253 mm_embeds = build_mm_embeds(audio_n, image_n, video_n, hidden)
254
255 model, _ = make_mock_model(hidden)
256 result = model.embed_input_ids(
257 input_ids, mm_embeds, is_multimodal=is_multimodal
258 )
259 return result, input_ids, is_multimodal
260
261 def test_audio_only(self):
262 """Audio-only: audio positions get audio embeddings."""
263 audio_n, hidden = 5, 8
264 audio_val = 10.0
265 result, input_ids, is_multimodal = self._run(audio_n, 0, 0, hidden)
266
267 audio_pos = (input_ids == AUDIO_TOKEN_ID).nonzero(as_tuple=True)[0]
268 assert result[audio_pos].allclose(torch.full((audio_n, hidden), audio_val)), (
269 "Audio positions should get audio embeddings"
270 )
271
272 def test_video_only(self):
273 """Video-only: video positions get video embeddings."""
274 video_n, hidden = 6, 8
275 video_val = 30.0
276 result, input_ids, is_multimodal = self._run(0, 0, video_n, hidden)
277
278 video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0]
279 assert result[video_pos].allclose(torch.full((video_n, hidden), video_val)), (
280 "Video positions should get video embeddings"
281 )
282
283 def test_mixed_modalities_audio_goes_to_audio_pos(self):
284 """Regression test for GitHub issue #34506:
285 With audio + image + video (non-interleaved), audio positions must
286 receive audio embeddings (not image or video embeddings).
287 """
288 audio_n, image_n, video_n, hidden = 5, 4, 6, 8
289 audio_val, image_val, video_val = 10.0, 20.0, 30.0
290
291 result, input_ids, is_multimodal = self._run(audio_n, image_n, video_n, hidden)
292
293 audio_pos = (input_ids == AUDIO_TOKEN_ID).nonzero(as_tuple=True)[0]
294 image_pos = (input_ids == IMAGE_TOKEN_ID).nonzero(as_tuple=True)[0]
295 video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0]
296
297 mean_a = result[audio_pos].mean().item()
298 assert result[audio_pos].allclose(torch.full((audio_n, hidden), audio_val)), (
299 f"Audio emb wrong: expected {audio_val}, got mean={mean_a:.1f}"
300 )
301
302 mean_i = result[image_pos].mean().item()
303 assert result[image_pos].allclose(torch.full((image_n, hidden), image_val)), (
304 f"Image emb wrong: expected {image_val}, got mean={mean_i:.1f}"
305 )
306
307 mean_v = result[video_pos].mean().item()
308 assert result[video_pos].allclose(torch.full((video_n, hidden), video_val)), (
309 f"Video emb wrong: expected {video_val}, got mean={mean_v:.1f}"
310 )
311
312 def test_text_positions_unchanged(self):
313 """Text positions should keep their text embeddings."""
314 audio_n, image_n, video_n, hidden = 3, 2, 4, 8
315 result, input_ids, is_multimodal = self._run(audio_n, image_n, video_n, hidden)
316
317 text_pos = (~is_multimodal).nonzero(as_tuple=True)[0]
318 # Text tokens have value TEXT_TOKEN_ID=0, so embed → 0.0
319 assert result[text_pos].allclose(torch.zeros(len(text_pos), hidden)), (
320 "Text positions should keep text embeddings"
321 )
322
323 def test_interleaved_use_audio_in_video(self):
324 """Interleaved (use_audio_in_video): video chunks interleaved with audio.
325 Video embeddings must go to video positions, audio to audio positions.
326 """
327 hidden = 8
328 audio_val, video_val = 10.0, 30.0
329 # Two video chunks of 4, two audio chunks of 3
330 video_chunks = [4, 4]
331 audio_chunks = [3, 3]
332 input_ids, is_multimodal = make_interleaved_seq(video_chunks, audio_chunks)
333
334 video_n = sum(video_chunks) # 8
335 audio_n = sum(audio_chunks) # 6
336
337 # mm_embeds come in [video, audio] order (video feature first in
338 # mm_features when positions are the same for use_audio_in_video)
339 mm_embeds = [
340 _mm_embed((video_n, hidden), video_val, "video"),
341 _mm_embed((audio_n, hidden), audio_val, "audio"),
342 ]
343
344 model, _ = make_mock_model(hidden)
345 # Modalities are attached on the embedding tensors (as in encoder gather).
346 result = model.embed_input_ids(
347 input_ids,
348 mm_embeds,
349 is_multimodal=is_multimodal,
350 )
351
352 video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0]
353 audio_pos = (input_ids == AUDIO_TOKEN_ID).nonzero(as_tuple=True)[0]
354
355 assert result[video_pos].allclose(torch.full((video_n, hidden), video_val)), (
356 "Interleaved: video positions should get video embeddings"
357 )
358
359 assert result[audio_pos].allclose(torch.full((audio_n, hidden), audio_val)), (
360 "Interleaved: audio positions should get audio embeddings"
361 )
362
363
364 # ---------------------------------------------------------------------------
365 # Tests for merge_interleaved_embeddings helper
366 # ---------------------------------------------------------------------------
367
368
369 class TestMergeInterleavedEmbeddings:
370 def test_basic_interleaved(self):
371 """Video chunks + audio chunks scattered to correct positions."""
372 hidden = 4
373 input_ids, is_multimodal = make_interleaved_seq([3, 3], [2, 2])
374
375 is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
376 is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
377 num_video = is_video.sum().item() # 6
378 num_audio = is_audio.sum().item() # 4
379
380 inputs_embeds = torch.zeros(len(input_ids), hidden)
381 mm_embeds = [
382 _mm_embed((num_video, hidden), 30.0, "video"),
383 _mm_embed((num_audio, hidden), 10.0, "audio"),
384 ]
385
386 result = merge_interleaved_embeddings(
387 inputs_embeds,
388 mm_embeds,
389 is_video,
390 is_audio,
391 is_multimodal,
392 )
393
394 video_pos = is_video.nonzero(as_tuple=True)[0]
395 audio_pos = is_audio.nonzero(as_tuple=True)[0]
396 assert result[video_pos].allclose(torch.full((num_video, hidden), 30.0))
397 assert result[audio_pos].allclose(torch.full((num_audio, hidden), 10.0))
398
399 def test_image_and_video_mixed(self):
400 """Image embeddings must not be misclassified as video."""
401 hidden = 4
402 # [text][I I][V A V A]
403 tokens = (
404 [TEXT_TOKEN_ID] * 2
405 + [IMAGE_TOKEN_ID] * 2
406 + [VIDEO_TOKEN_ID, AUDIO_TOKEN_ID, VIDEO_TOKEN_ID, AUDIO_TOKEN_ID]
407 )
408 input_ids = torch.tensor(tokens)
409 is_multimodal = (
410 (input_ids == IMAGE_TOKEN_ID)
411 | (input_ids == VIDEO_TOKEN_ID)
412 | (input_ids == AUDIO_TOKEN_ID)
413 )
414 is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
415 is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
416 is_image = is_multimodal & (input_ids == IMAGE_TOKEN_ID)
417
418 inputs_embeds = torch.zeros(len(input_ids), hidden)
419 mm_embeds = [
420 _mm_embed((2, hidden), 20.0, "image"),
421 _mm_embed((2, hidden), 30.0, "video"),
422 _mm_embed((2, hidden), 10.0, "audio"),
423 ]
424 result = merge_interleaved_embeddings(
425 inputs_embeds,
426 mm_embeds,
427 is_video,
428 is_audio,
429 is_multimodal,
430 )
431 assert result[is_image.nonzero(as_tuple=True)[0]].allclose(
432 torch.full((2, hidden), 20.0)
433 )
434 assert result[is_video.nonzero(as_tuple=True)[0]].allclose(
435 torch.full((2, hidden), 30.0)
436 )
437 assert result[is_audio.nonzero(as_tuple=True)[0]].allclose(
438 torch.full((2, hidden), 10.0)
439 )
440
441 def test_missing_modality_raises(self):
442 hidden = 2
443 input_ids, is_multimodal = make_interleaved_seq([2], [2])
444 is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
445 is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
446 with pytest.raises(ValueError, match="Missing modality"):
447 merge_interleaved_embeddings(
448 torch.zeros(len(input_ids), hidden),
449 [torch.zeros(2, hidden), torch.zeros(2, hidden)],
450 is_video,
451 is_audio,
452 is_multimodal,
453 )
454
455
456 if __name__ == "__main__":
457 pytest.main([__file__, "-v"])