@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 201 lines 6.46 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from unittest.mock import patch
5
6 import numpy as np
7 import pytest
8 import torch
9
10 from vllm import envs
11 from vllm.model_executor.models.nano_nemotron_vl import (
12 NanoNemotronVLMultiModalProcessor,
13 NemotronH_Nano_VL_V2,
14 )
15 from vllm.multimodal.parse import (
16 MultiModalDataItems,
17 MultiModalDataParser,
18 VideoProcessorItems,
19 )
20
21
22 @pytest.mark.parametrize("input_key", ["image_embeds", "video_embeds"])
23 def test_precomputed_multimodal_embeddings(input_key: str):
24 model = object.__new__(NemotronH_Nano_VL_V2)
25 embeds = torch.randn(2, 4, 8)
26
27 outputs = model.embed_multimodal(**{input_key: embeds})
28
29 assert len(outputs) == len(embeds)
30 assert all(torch.equal(output, embed) for output, embed in zip(outputs, embeds))
31
32
33 class _TextOnlyMultiModalConfig:
34 def get_limit_per_prompt(self, modality: str) -> int:
35 return 0
36
37
38 class _ImageOnlyMultiModalConfig:
39 def get_limit_per_prompt(self, modality: str) -> int:
40 return 1 if modality == "image" else 0
41
42
43 class _ModelConfig:
44 multimodal_config = _TextOnlyMultiModalConfig()
45
46
47 class _ImageOnlyModelConfig:
48 multimodal_config = _ImageOnlyMultiModalConfig()
49
50
51 class _LanguageModel:
52 def __init__(self) -> None:
53 self.loaded_weights: list[tuple[str, object]] = []
54
55 def load_weights(self, weights):
56 self.loaded_weights = list(weights)
57
58
59 class _MissingMultiModalModule:
60 def named_parameters(self):
61 raise AssertionError("multimodal weights should not be inspected")
62
63 def load_weights(self, weights):
64 raise AssertionError("multimodal weights should not be loaded")
65
66
67 class _AdapterModule:
68 def named_parameters(self):
69 return []
70
71
72 class _VisionModel:
73 def __init__(self) -> None:
74 self.loaded_weights: list[tuple[str, object]] = []
75
76 def load_weights(self, weights):
77 self.loaded_weights = list(weights)
78
79
80 class _FakeTensor:
81 """Sentinel stand-in for torch.Tensor in load_weights tests. Supports the
82 .detach().clone() chain used by load_weights for buffered mm weights;
83 both methods return self so identity (and the existing equality
84 assertions) are preserved through cloning."""
85
86 def detach(self):
87 return self
88
89 def clone(self):
90 return self
91
92
93 def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode():
94 model = object.__new__(NemotronH_Nano_VL_V2)
95 language_model = _LanguageModel()
96 object.__setattr__(model, "model_config", _ModelConfig())
97 object.__setattr__(model, "language_model", language_model)
98 object.__setattr__(model, "mlp1", _AdapterModule())
99 object.__setattr__(model, "vision_model", _MissingMultiModalModule())
100 object.__setattr__(model, "sound_encoder", None)
101
102 language_weight = object()
103 model.load_weights(
104 [
105 ("language_model.layers.0.weight", language_weight),
106 ("mlp1.0.weight", object()),
107 ("vision_model.radio_model.encoder.weight", object()),
108 ("sound_encoder.encoder.weight", object()),
109 ]
110 )
111
112 assert language_model.loaded_weights == [("layers.0.weight", language_weight)]
113
114
115 def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder():
116 model = object.__new__(NemotronH_Nano_VL_V2)
117 language_model = _LanguageModel()
118 vision_model = _VisionModel()
119 object.__setattr__(model, "model_config", _ImageOnlyModelConfig())
120 object.__setattr__(model, "language_model", language_model)
121 object.__setattr__(model, "mlp1", _AdapterModule())
122 object.__setattr__(model, "vision_model", vision_model)
123 object.__setattr__(model, "sound_encoder", None)
124
125 language_weight = object()
126 vision_weight = _FakeTensor()
127 model.load_weights(
128 [
129 ("language_model.layers.0.weight", language_weight),
130 ("vision_model.radio_model.encoder.weight", vision_weight),
131 ]
132 )
133
134 assert language_model.loaded_weights == [("layers.0.weight", language_weight)]
135 assert vision_model.loaded_weights == [
136 ("radio_model.encoder.weight", vision_weight)
137 ]
138
139
140 def test_nano_nemotron_vl_requires_sound_encoder_for_sound_weights():
141 model = object.__new__(NemotronH_Nano_VL_V2)
142 language_model = _LanguageModel()
143 vision_model = _VisionModel()
144 object.__setattr__(model, "model_config", _ImageOnlyModelConfig())
145 object.__setattr__(model, "language_model", language_model)
146 object.__setattr__(model, "mlp1", _AdapterModule())
147 object.__setattr__(model, "vision_model", vision_model)
148 object.__setattr__(model, "sound_encoder", None)
149
150 with pytest.raises(AssertionError):
151 model.load_weights([("sound_encoder.encoder.weight", object())])
152
153
154 def _make_mm_items_with_video_bytes(
155 video_bytes: bytes,
156 ) -> MultiModalDataItems:
157 """Build a minimal MultiModalDataItems with one video entry."""
158 items = MultiModalDataItems()
159 items["video"] = VideoProcessorItems(
160 data=[None],
161 metadata=[{"original_video_bytes": video_bytes}],
162 )
163 return items
164
165
166 def test_extract_audio_from_videos_passes_max_duration():
167 """_extract_audio_from_videos must forward VLLM_MAX_AUDIO_DECODE_DURATION_S
168 to load_audio_pyav so decompression-bomb audio is rejected."""
169 dummy_audio = (np.zeros(16000, dtype=np.float32), 16000.0)
170 mm_items = _make_mm_items_with_video_bytes(b"\x00" * 64)
171
172 processor = object.__new__(NanoNemotronVLMultiModalProcessor)
173 processor.data_parser = MultiModalDataParser(
174 target_sr=dummy_audio[1], target_channels=1
175 )
176
177 target = "vllm.model_executor.models.nano_nemotron_vl.load_audio_pyav"
178 with patch(target, return_value=dummy_audio) as mock_load:
179 processor._extract_audio_from_videos(mm_items)
180
181 mock_load.assert_called_once()
182 _, kwargs = mock_load.call_args
183 assert kwargs["max_duration_s"] == envs.VLLM_MAX_AUDIO_DECODE_DURATION_S
184
185
186 def test_extract_audio_from_videos_rejects_oversized_audio():
187 """When load_audio_pyav raises due to duration limit the video is
188 marked as having no audio instead of crashing the server."""
189 mm_items = _make_mm_items_with_video_bytes(b"\x00" * 64)
190
191 processor = object.__new__(NanoNemotronVLMultiModalProcessor)
192
193 target = "vllm.model_executor.models.nano_nemotron_vl.load_audio_pyav"
194 with patch(
195 target,
196 side_effect=ValueError("Audio exceeds maximum allowed duration"),
197 ):
198 _, audio_items, has_audio = processor._extract_audio_from_videos(mm_items)
199
200 assert audio_items == []
201 assert has_audio == [False]