@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 210 lines 5.48 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass, field
4
5 import pytest
6 import torch
7
8 from vllm.model_executor.models.paddleocr_vl import (
9 PaddleOCRVLForConditionalGeneration,
10 )
11 from vllm.multimodal.inputs import (
12 MultiModalFeatureSpec,
13 MultiModalFieldElem,
14 MultiModalKwargsItem,
15 PlaceholderRange,
16 )
17
18 pytestmark = pytest.mark.skip_global_cleanup
19
20
21 @pytest.fixture(autouse=True, scope="module")
22 def _force_cpu_default_device():
23 original = torch.get_default_device()
24 torch.set_default_device("cpu")
25 yield
26 torch.set_default_device(original)
27
28
29 @dataclass
30 class DummyVisionConfig:
31 spatial_merge_size: int = 2
32 patch_size: int = 14
33
34
35 @dataclass
36 class DummyConfig:
37 image_token_id: int = 151655
38 video_token_id: int = 151654
39 vision_start_token_id: int = 151652
40 vision_end_token_id: int = 151653
41 vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
42
43
44 def make_model(config: DummyConfig) -> PaddleOCRVLForConditionalGeneration:
45 model = object.__new__(PaddleOCRVLForConditionalGeneration)
46 model.config = config
47 return model
48
49
50 def make_mm_feature(
51 *,
52 offset: int,
53 length: int,
54 image_grid_thw: tuple[int, int, int],
55 ) -> MultiModalFeatureSpec:
56 return MultiModalFeatureSpec(
57 data=MultiModalKwargsItem(
58 {
59 "image_grid_thw": MultiModalFieldElem(
60 data=torch.tensor(image_grid_thw),
61 field=None,
62 ),
63 }
64 ),
65 modality="image",
66 identifier="DUMMY",
67 mm_position=PlaceholderRange(offset=offset, length=length),
68 )
69
70
71 def test_get_mrope_input_positions_text_only():
72 model = make_model(DummyConfig())
73 input_tokens = [11, 12, 13, 14, 15]
74 positions, delta = model.get_mrope_input_positions(
75 input_tokens=input_tokens,
76 mm_features=[],
77 )
78 expected = torch.tensor(
79 [
80 [0, 1, 2, 3, 4],
81 [0, 1, 2, 3, 4],
82 [0, 1, 2, 3, 4],
83 ]
84 )
85 assert torch.equal(positions, expected)
86 assert delta == 0
87
88
89 def test_get_mrope_input_positions_single_image():
90 model = make_model(DummyConfig())
91 spatial_merge_size = model.config.vision_config.spatial_merge_size
92
93 t, h, w = 1, 2, 2
94 num_image_tokens = t * h * w
95
96 input_tokens = (
97 [10]
98 + [model.config.vision_start_token_id]
99 + [model.config.image_token_id] * num_image_tokens
100 + [model.config.vision_end_token_id]
101 + [30, 31]
102 )
103
104 mm_features = [
105 make_mm_feature(
106 offset=2, # 1 (text) + 1 (vision_start)
107 length=num_image_tokens,
108 image_grid_thw=(t, h * spatial_merge_size, w * spatial_merge_size),
109 )
110 ]
111
112 positions, delta = model.get_mrope_input_positions(
113 input_tokens=input_tokens,
114 mm_features=mm_features,
115 )
116
117 expected = torch.tensor(
118 [
119 [0, 1, 2, 2, 2, 2, 4, 5, 6],
120 [0, 1, 2, 2, 3, 3, 4, 5, 6],
121 [0, 1, 2, 3, 2, 3, 4, 5, 6],
122 ]
123 )
124
125 assert torch.equal(positions, expected)
126 expected_delta = (positions.max().item() + 1) - len(input_tokens)
127 assert delta == expected_delta
128
129
130 def test_get_mrope_input_positions_multiple_images():
131 model = make_model(DummyConfig())
132 spatial_merge_size = model.config.vision_config.spatial_merge_size
133
134 t1, h1, w1 = 1, 2, 2
135 num1 = t1 * h1 * w1
136
137 t2, h2, w2 = 1, 1, 3
138 num2 = t2 * h2 * w2
139
140 input_tokens = (
141 [10]
142 + [model.config.vision_start_token_id]
143 + [model.config.image_token_id] * num1
144 + [model.config.vision_end_token_id]
145 + [20, 21]
146 + [model.config.vision_start_token_id]
147 + [model.config.image_token_id] * num2
148 + [model.config.vision_end_token_id]
149 + [30]
150 )
151
152 mm_features = [
153 make_mm_feature(
154 offset=2,
155 length=num1,
156 image_grid_thw=(t1, h1 * spatial_merge_size, w1 * spatial_merge_size),
157 ),
158 make_mm_feature(
159 offset=2 + num1 + 1 + 2 + 1,
160 length=num2,
161 image_grid_thw=(t2, h2 * spatial_merge_size, w2 * spatial_merge_size),
162 ),
163 ]
164
165 positions, delta = model.get_mrope_input_positions(
166 input_tokens=input_tokens,
167 mm_features=mm_features,
168 )
169
170 assert positions.shape == (3, 15)
171 assert not torch.equal(positions[:, 2:6], torch.arange(4).expand(3, 4) + 2)
172 assert not torch.equal(positions[:, 10:13], torch.arange(3).expand(3, 3) + 10)
173
174
175 def test_get_mrope_input_positions_image_at_start():
176 model = make_model(DummyConfig())
177 spatial_merge_size = model.config.vision_config.spatial_merge_size
178
179 t, h, w = 1, 2, 2
180 num_tokens = t * h * w
181
182 input_tokens = (
183 [model.config.vision_start_token_id]
184 + [model.config.image_token_id] * num_tokens
185 + [model.config.vision_end_token_id]
186 + [10, 11]
187 )
188
189 mm_features = [
190 make_mm_feature(
191 offset=1, # start token at index 0
192 length=num_tokens,
193 image_grid_thw=(t, h * spatial_merge_size, w * spatial_merge_size),
194 )
195 ]
196
197 positions, delta = model.get_mrope_input_positions(
198 input_tokens=input_tokens,
199 mm_features=mm_features,
200 )
201
202 expected = torch.tensor(
203 [
204 [0, 1, 1, 1, 1, 3, 4, 5],
205 [0, 1, 1, 2, 2, 3, 4, 5],
206 [0, 1, 2, 1, 2, 3, 4, 5],
207 ]
208 )
209
210 assert torch.equal(positions, expected)