@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / babad6e5dd
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Misc] Move DP for ViT code inside model executor dir (#25459)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Sep 23, 2025 at 17:20 UTC babad6e5dde380b453be716c7ebc2b6928658c18
13 files changed +721 -730
tests/models/test_vision.py
+423 -1
@@ -1,10 +1,20 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import math
4
5 import pytest
6 import torch
7 +import torch.multiprocessing as mp
8
7 -from vllm.model_executor.models.vision import resolve_visual_encoder_outputs
9 +from tests.utils import multi_gpu_test
10 +from vllm.distributed import get_tensor_model_parallel_world_size
11 +from vllm.distributed.parallel_state import (init_distributed_environment,
12 + initialize_model_parallel)
13 +from vllm.model_executor.models.vision import (
14 + get_load_balance_assignment, resolve_visual_encoder_outputs,
15 + run_dp_sharded_mrope_vision_model, run_dp_sharded_vision_model)
16 +from vllm.platforms import current_platform
17 +from vllm.utils import get_open_port, update_environment_variables
18
19
20 @pytest.mark.parametrize(
@@ -33,3 +43,415 @@ def test_resolve_visual_encoder_outputs(feature_sample_layers,
43 post_layer_norm=None,
44 max_possible_layers=max_possible_layers)
45 assert torch.equal(torch.tensor(expected_features), output_tensor)
46 +
47 +
48 +class SimpleLinearModel(torch.nn.Module):
49 + """A simple linear vision model for testing."""
50 +
51 + def __init__(self, input_dim: int = 3 * 224 * 224, output_dim: int = 32):
52 + super().__init__()
53 + self.flatten = torch.nn.Flatten()
54 + self.linear = torch.nn.Linear(input_dim, output_dim)
55 +
56 + def forward(self, x: torch.Tensor):
57 + # Flatten the input and apply linear transformation
58 + x = self.flatten(x)
59 + return self.linear(x)
60 +
61 +
62 +@multi_gpu_test(num_gpus=2)
63 +@pytest.mark.parametrize(
64 + "batch_size",
65 + [
66 + 1, # Single image
67 + 4, # Small batch
68 + 5, # Odd batch size (for testing padding)
69 + ],
70 +)
71 +def test_run_dp_sharded_vision_model(batch_size: int):
72 + world_size = 2
73 + # Launch processes
74 + mp.spawn(
75 + run_dp_sharded_vision_model_vs_direct,
76 + args=(
77 + world_size,
78 + batch_size,
79 + get_open_port(),
80 + ),
81 + nprocs=world_size,
82 + )
83 +
84 +
85 +def run_dp_sharded_vision_model_vs_direct(local_rank: int, world_size: int,
86 + batch_size: int, master_port: int):
87 + """
88 + Test that run_dp_sharded_vision_model produces the same results as
89 + calling the model directly.
90 + """
91 +
92 + # Set random seed for reproducibility
93 + current_platform.seed_everything(0)
94 +
95 + device = f"{current_platform.device_name}:{local_rank}"
96 + current_platform.set_device(device)
97 + torch.set_default_device(device)
98 +
99 + update_environment_variables({
100 + 'RANK': str(local_rank),
101 + 'LOCAL_RANK': str(local_rank),
102 + 'WORLD_SIZE': str(world_size),
103 + 'MASTER_ADDR': 'localhost',
104 + 'MASTER_PORT': str(master_port),
105 + })
106 +
107 + # initialize distributed
108 + init_distributed_environment()
109 + initialize_model_parallel(tensor_model_parallel_size=world_size)
110 +
111 + # Create a test input tensor
112 + image_input = torch.randn(batch_size, 3, 224, 224)
113 +
114 + # Create a simple linear model
115 + vision_model = SimpleLinearModel()
116 +
117 + # Run the model directly on the full input
118 + with torch.inference_mode():
119 + direct_output = vision_model(image_input)
120 +
121 + # Run the model through the sharded function
122 + with torch.inference_mode():
123 + sharded_output = run_dp_sharded_vision_model(image_input, vision_model)
124 +
125 + # Check that the world size is set up correctly
126 + assert get_tensor_model_parallel_world_size() == world_size
127 +
128 + # Check that the outputs have the same shape
129 + assert direct_output.shape == sharded_output.shape
130 +
131 + # Check that the outputs are close (they should be identical)
132 + assert torch.allclose(direct_output, sharded_output, rtol=1e-5, atol=1e-5)
133 +
134 +
135 +@pytest.mark.parametrize(
136 + "sizes,num_gpus,expected_shuffle_indices,expected_gpu_sample_counts,"
137 + "expected_grouped_sizes_per_gpu,test_description",
138 + [
139 + # Empty input
140 + ([], 2, [], [0, 0], [0, 0], "empty input"),
141 +
142 + # Fewer samples than GPUs
143 + ([100, 200], 4, [1, 0], [1, 1, 0, 0], [200, 100, 0, 0
144 + ], "fewer samples than GPUs"),
145 +
146 + # Single GPU
147 + ([100, 200, 300], 1, [2, 1, 0], [3], [600], "single GPU"),
148 +
149 + # Balanced assignment
150 + ([100, 100, 100, 100
151 + ], 2, [0, 2, 1, 3], [2, 2], [200, 200], "balanced assignment"),
152 +
153 + # Unbalanced sizes - this one is trickier since the algorithm is greedy
154 + ([1000, 100, 200, 50], 2, [0, 2, 1, 3
155 + ], [1, 3], [1000, 350], "unbalanced sizes"),
156 + ],
157 +)
158 +def test_get_load_balance_assignment_cases(sizes, num_gpus,
159 + expected_shuffle_indices,
160 + expected_gpu_sample_counts,
161 + expected_grouped_sizes_per_gpu,
162 + test_description):
163 + """Test get_load_balance_assignment with various input cases."""
164 + result = get_load_balance_assignment(sizes, num_gpus=num_gpus)
165 + (shuffle_indices, gpu_sample_counts, grouped_sizes_per_gpu) = result
166 +
167 + # Common assertions for all cases
168 + assert len(shuffle_indices) == len(sizes)
169 + assert len(gpu_sample_counts) == num_gpus
170 + assert len(grouped_sizes_per_gpu) == num_gpus
171 + assert sum(gpu_sample_counts) == len(sizes)
172 +
173 + assert shuffle_indices == expected_shuffle_indices
174 +
175 + assert gpu_sample_counts == expected_gpu_sample_counts
176 + assert grouped_sizes_per_gpu == expected_grouped_sizes_per_gpu
177 +
178 +
179 +class SimpleMRopeVisionModel(torch.nn.Module):
180 + """A simple vision model for testing mrope functionality."""
181 +
182 + def __init__(self, spatial_merge_size: int = 2, out_hidden_size: int = 64):
183 + super().__init__()
184 + self.spatial_merge_size = spatial_merge_size
185 + self.out_hidden_size = out_hidden_size
186 + self.linear = torch.nn.Linear(768, out_hidden_size)
187 +
188 + def forward(self, pixel_values: torch.Tensor,
189 + grid_thw_list: list[list[int]]):
190 + """Simple forward pass that simulates spatial merging."""
191 + # Apply linear transformation
192 + embeddings = self.linear(pixel_values)
193 +
194 + # Simulate spatial merging by reducing the number of patches
195 + merge_factor = self.spatial_merge_size * self.spatial_merge_size
196 +
197 + # Group patches and merge spatially
198 + merged_embeddings = []
199 + start_idx = 0
200 +
201 + for grid_thw in grid_thw_list:
202 + num_patches = math.prod(grid_thw)
203 + end_idx = start_idx + num_patches
204 +
205 + # Get patches for this image
206 + image_patches = embeddings[start_idx:end_idx]
207 +
208 + # Simulate spatial merging by averaging groups of patches
209 + merged_patches = num_patches // merge_factor
210 + if merged_patches > 0:
211 + # Reshape and average to simulate merging
212 + reshaped = image_patches[:merged_patches * merge_factor].view(
213 + merged_patches, merge_factor, -1)
214 + merged = reshaped.mean(dim=1)
215 + merged_embeddings.append(merged)
216 +
217 + start_idx = end_idx
218 +
219 + if merged_embeddings:
220 + return torch.cat(merged_embeddings, dim=0)
221 + else:
222 + return torch.empty((0, self.out_hidden_size),
223 + device=pixel_values.device,
224 + dtype=pixel_values.dtype)
225 +
226 +
227 +@multi_gpu_test(num_gpus=2)
228 +@pytest.mark.parametrize(
229 + "batch_size",
230 + [
231 + 1, # Single image
232 + 3, # Small batch
233 + 5, # Odd batch size (for testing padding)
234 + ],
235 +)
236 +def test_run_dp_sharded_mrope_vision_model(batch_size: int):
237 + world_size = 2
238 + # Launch processes
239 + mp.spawn(
240 + run_dp_sharded_mrope_vision_model_vs_direct,
241 + args=(
242 + world_size,
243 + batch_size,
244 + get_open_port(),
245 + ),
246 + nprocs=world_size,
247 + )
248 +
249 +
250 +def run_dp_sharded_mrope_vision_model_vs_direct(local_rank: int,
251 + world_size: int,
252 + batch_size: int,
253 + master_port: int):
254 + """
255 + Test that run_dp_sharded_mrope_vision_model produces the same results as
256 + calling the model directly.
257 + """
258 + # Set random seed for reproducibility
259 + current_platform.seed_everything(0)
260 + device = f"{current_platform.device_name}:{local_rank}"
261 + current_platform.set_device(device)
262 + torch.set_default_device(device)
263 +
264 + update_environment_variables({
265 + 'RANK': str(local_rank),
266 + 'LOCAL_RANK': str(local_rank),
267 + 'WORLD_SIZE': str(world_size),
268 + 'MASTER_ADDR': 'localhost',
269 + 'MASTER_PORT': str(master_port),
270 + })
271 +
272 + # initialize distributed
273 + init_distributed_environment()
274 + initialize_model_parallel(tensor_model_parallel_size=world_size)
275 +
276 + # Create test data
277 + grid_thw_list = []
278 + pixel_values_list = []
279 +
280 + for i in range(batch_size):
281 + # Varying image sizes for better testing
282 + t, h, w = 1, 4 + i, 4 + i
283 + grid_thw_list.append([t, h, w])
284 +
285 + num_patches = t * h * w
286 + # Create random pixel values for this image
287 + image_pixels = torch.randn(num_patches, 768)
288 + pixel_values_list.append(image_pixels)
289 +
290 + # Concatenate all pixel values
291 + pixel_values = torch.cat(pixel_values_list, dim=0)
292 +
293 + # Create a simple mrope vision model
294 + vision_model = SimpleMRopeVisionModel()
295 +
296 + # Run the model directly on the full input (only on rank 0)
297 + if local_rank == 0:
298 + with torch.inference_mode():
299 + direct_output = vision_model(pixel_values, grid_thw_list)
300 +
301 + # Run the model through the sharded function
302 + with torch.inference_mode():
303 + sharded_output = run_dp_sharded_mrope_vision_model(vision_model,
304 + pixel_values,
305 + grid_thw_list,
306 + rope_type="rope_3d")
307 + sharded_output = torch.cat(sharded_output, dim=0)
308 +
309 + # Check that the world size is set up correctly
310 + assert get_tensor_model_parallel_world_size() == world_size
311 +
312 + # Compare outputs (only on rank 0)
313 + if local_rank == 0:
314 + # Check that the outputs have the same shape
315 + assert direct_output.shape == sharded_output.shape
316 + # Check that the outputs are close (they should be identical)
317 + assert torch.allclose(direct_output,
318 + sharded_output,
319 + rtol=1e-5,
320 + atol=1e-5)
321 +
322 +
323 +@multi_gpu_test(num_gpus=2)
324 +def test_run_dp_sharded_mrope_vision_model_empty_input():
325 + world_size = 2
326 + mp.spawn(
327 + run_dp_sharded_mrope_vision_model_empty_input_worker,
328 + args=(world_size, get_open_port()),
329 + nprocs=world_size,
330 + )
331 +
332 +
333 +def run_dp_sharded_mrope_vision_model_empty_input_worker(
334 + local_rank: int, world_size: int, master_port: int):
335 + """Test run_dp_sharded_mrope_vision_model with empty input."""
336 + # Set up distributed environment
337 + device = f"{current_platform.device_name}:{local_rank}"
338 + current_platform.set_device(device)
339 + torch.set_default_device(device)
340 +
341 + update_environment_variables({
342 + 'RANK': str(local_rank),
343 + 'LOCAL_RANK': str(local_rank),
344 + 'WORLD_SIZE': str(world_size),
345 + 'MASTER_ADDR': 'localhost',
346 + 'MASTER_PORT': str(master_port),
347 + })
348 +
349 + init_distributed_environment()
350 + initialize_model_parallel(tensor_model_parallel_size=world_size)
351 +
352 + # Create empty inputs
353 + pixel_values = torch.empty((0, 768))
354 + grid_thw_list: list[list[int]] = []
355 +
356 + vision_model = SimpleMRopeVisionModel()
357 +
358 + # Should handle empty input gracefully
359 + with torch.inference_mode():
360 + output = run_dp_sharded_mrope_vision_model(vision_model,
361 + pixel_values,
362 + grid_thw_list,
363 + rope_type="rope_3d")
364 +
365 + assert len(output) == 0
366 +
367 +
368 +@multi_gpu_test(num_gpus=4)
369 +def test_run_dp_sharded_mrope_vision_model_uneven_load():
370 + world_size = 4
371 + mp.spawn(
372 + run_dp_sharded_mrope_vision_model_uneven_load_worker,
373 + args=(world_size, get_open_port()),
374 + nprocs=world_size,
375 + )
376 +
377 +
378 +def run_dp_sharded_mrope_vision_model_uneven_load_worker(
379 + local_rank: int, world_size: int, master_port: int):
380 + """Test run_dp_sharded_mrope_vision_model with uneven load distribution."""
381 + # Set up distributed environment
382 + current_platform.seed_everything(123)
383 + device = f"{current_platform.device_name}:{local_rank}"
384 + current_platform.set_device(device)
385 + torch.set_default_device(device)
386 +
387 + update_environment_variables({
388 + 'RANK': str(local_rank),
389 + 'LOCAL_RANK': str(local_rank),
390 + 'WORLD_SIZE': str(world_size),
391 + 'MASTER_ADDR': 'localhost',
392 + 'MASTER_PORT': str(master_port),
393 + })
394 +
395 + init_distributed_environment()
396 + initialize_model_parallel(tensor_model_parallel_size=world_size)
397 +
398 + # Create images with very different sizes
399 + grid_thw_list = [
400 + [1, 2, 2], # Small: 4 patches
401 + [1, 8, 8], # Large: 64 patches
402 + [1, 3, 3], # Medium: 9 patches
403 + ]
404 +
405 + pixel_values_list = []
406 + for grid_thw in grid_thw_list:
407 + num_patches = math.prod(grid_thw)
408 + image_pixels = torch.randn(num_patches, 768)
409 + pixel_values_list.append(image_pixels)
410 +
411 + pixel_values = torch.cat(pixel_values_list, dim=0)
412 + vision_model = SimpleMRopeVisionModel()
413 +
414 + # Should handle uneven distribution without errors
415 + with torch.inference_mode():
416 + output_tuple = run_dp_sharded_mrope_vision_model(vision_model,
417 + pixel_values,
418 + grid_thw_list,
419 + rope_type="rope_3d")
420 +
421 + # Verify output shape is reasonable
422 + merge_factor = vision_model.spatial_merge_size**2
423 + expected_output_patches = list(
424 + math.prod(grid_thw) // merge_factor for grid_thw in grid_thw_list)
425 +
426 + for i, output in enumerate(output_tuple):
427 + assert output.shape[0] == expected_output_patches[i]
428 + assert output.shape[1] == vision_model.out_hidden_size
429 +
430 +
431 +@pytest.mark.parametrize("spatial_merge_size", [2, 4])
432 +def test_simple_mrope_vision_model_spatial_merge(spatial_merge_size: int):
433 + """Test SimpleMRopeVisionModel with different spatial merge sizes."""
434 + device = current_platform.device_type
435 +
436 + grid_thw_list = [[1, 4, 4], [1, 6, 6]] # Two images
437 + pixel_values_list = []
438 +
439 + for grid_thw in grid_thw_list:
440 + num_patches = math.prod(grid_thw)
441 + image_pixels = torch.randn(num_patches, 768, device=device)
442 + pixel_values_list.append(image_pixels)
443 +
444 + pixel_values = torch.cat(pixel_values_list, dim=0)
445 + vision_model = SimpleMRopeVisionModel(
446 + spatial_merge_size=spatial_merge_size).to(device)
447 +
448 + with torch.inference_mode():
449 + output = vision_model(pixel_values, grid_thw_list)
450 +
451 + # Verify output dimensions based on spatial merging
452 + total_patches = sum(math.prod(grid_thw) for grid_thw in grid_thw_list)
453 + merge_factor = spatial_merge_size**2
454 + expected_output_patches = total_patches // merge_factor
455 +
456 + assert output.shape[0] == expected_output_patches
457 + assert output.shape[1] == vision_model.out_hidden_size
tests/multimodal/test_utils.py
+1 -425
@@ -2,7 +2,6 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import base64
5 -import math
5 import mimetypes
6 import os
7 from tempfile import NamedTemporaryFile, TemporaryDirectory
@@ -10,22 +9,11 @@ from typing import TYPE_CHECKING, NamedTuple
9
10 import numpy as np
11 import pytest
13 -import torch
14 -import torch.multiprocessing as mp
12 from PIL import Image, ImageChops
13
17 -from tests.utils import multi_gpu_test
18 -from vllm.distributed import get_tensor_model_parallel_world_size
19 -from vllm.distributed.parallel_state import (init_distributed_environment,
20 - initialize_model_parallel)
14 from vllm.multimodal.image import convert_image_mode
15 from vllm.multimodal.inputs import PlaceholderRange
23 -from vllm.multimodal.utils import (MediaConnector, argsort_mm_positions,
24 - get_load_balance_assignment,
25 - run_dp_sharded_mrope_vision_model,
26 - run_dp_sharded_vision_model)
27 -from vllm.platforms import current_platform
28 -from vllm.utils import get_open_port, update_environment_variables
16 +from vllm.multimodal.utils import MediaConnector, argsort_mm_positions
17
18 if TYPE_CHECKING:
19 from vllm.multimodal.inputs import MultiModalPlaceholderDict
@@ -404,415 +392,3 @@ def test_argsort_mm_positions():
392 modality_idxs = argsort_mm_positions(mm_positions)
393
394 assert modality_idxs == expected_modality_idxs
407 -
408 -
409 -class SimpleLinearModel(torch.nn.Module):
410 - """A simple linear vision model for testing."""
411 -
412 - def __init__(self, input_dim: int = 3 * 224 * 224, output_dim: int = 32):
413 - super().__init__()
414 - self.flatten = torch.nn.Flatten()
415 - self.linear = torch.nn.Linear(input_dim, output_dim)
416 -
417 - def forward(self, x: torch.Tensor):
418 - # Flatten the input and apply linear transformation
419 - x = self.flatten(x)
420 - return self.linear(x)
421 -
422 -
423 -@multi_gpu_test(num_gpus=2)
424 -@pytest.mark.parametrize(
425 - "batch_size",
426 - [
427 - 1, # Single image
428 - 4, # Small batch
429 - 5, # Odd batch size (for testing padding)
430 - ],
431 -)
432 -def test_run_dp_sharded_vision_model(batch_size: int):
433 - world_size = 2
434 - # Launch processes
435 - mp.spawn(
436 - run_dp_sharded_vision_model_vs_direct,
437 - args=(
438 - world_size,
439 - batch_size,
440 - get_open_port(),
441 - ),
442 - nprocs=world_size,
443 - )
444 -
445 -
446 -def run_dp_sharded_vision_model_vs_direct(local_rank: int, world_size: int,
447 - batch_size: int, master_port: int):
448 - """
449 - Test that run_dp_sharded_vision_model produces the same results as
450 - calling the model directly.
451 - """
452 -
453 - # Set random seed for reproducibility
454 - current_platform.seed_everything(0)
455 -
456 - device = f"{current_platform.device_name}:{local_rank}"
457 - current_platform.set_device(device)
458 - torch.set_default_device(device)
459 -
460 - update_environment_variables({
461 - 'RANK': str(local_rank),
462 - 'LOCAL_RANK': str(local_rank),
463 - 'WORLD_SIZE': str(world_size),
464 - 'MASTER_ADDR': 'localhost',
465 - 'MASTER_PORT': str(master_port),
466 - })
467 -
468 - # initialize distributed
469 - init_distributed_environment()
470 - initialize_model_parallel(tensor_model_parallel_size=world_size)
471 -
472 - # Create a test input tensor
473 - image_input = torch.randn(batch_size, 3, 224, 224)
474 -
475 - # Create a simple linear model
476 - vision_model = SimpleLinearModel()
477 -
478 - # Run the model directly on the full input
479 - with torch.inference_mode():
480 - direct_output = vision_model(image_input)
481 -
482 - # Run the model through the sharded function
483 - with torch.inference_mode():
484 - sharded_output = run_dp_sharded_vision_model(image_input, vision_model)
485 -
486 - # Check that the world size is set up correctly
487 - assert get_tensor_model_parallel_world_size() == world_size
488 -
489 - # Check that the outputs have the same shape
490 - assert direct_output.shape == sharded_output.shape
491 -
492 - # Check that the outputs are close (they should be identical)
493 - assert torch.allclose(direct_output, sharded_output, rtol=1e-5, atol=1e-5)
494 -
495 -
496 -@pytest.mark.parametrize(
497 - "sizes,num_gpus,expected_shuffle_indices,expected_gpu_sample_counts,"
498 - "expected_grouped_sizes_per_gpu,test_description",
499 - [
500 - # Empty input
501 - ([], 2, [], [0, 0], [0, 0], "empty input"),
502 -
503 - # Fewer samples than GPUs
504 - ([100, 200], 4, [1, 0], [1, 1, 0, 0], [200, 100, 0, 0
505 - ], "fewer samples than GPUs"),
506 -
507 - # Single GPU
508 - ([100, 200, 300], 1, [2, 1, 0], [3], [600], "single GPU"),
509 -
510 - # Balanced assignment
511 - ([100, 100, 100, 100
512 - ], 2, [0, 2, 1, 3], [2, 2], [200, 200], "balanced assignment"),
513 -
514 - # Unbalanced sizes - this one is trickier since the algorithm is greedy
515 - ([1000, 100, 200, 50], 2, [0, 2, 1, 3
516 - ], [1, 3], [1000, 350], "unbalanced sizes"),
517 - ],
518 -)
519 -def test_get_load_balance_assignment_cases(sizes, num_gpus,
520 - expected_shuffle_indices,
521 - expected_gpu_sample_counts,
522 - expected_grouped_sizes_per_gpu,
523 - test_description):
524 - """Test get_load_balance_assignment with various input cases."""
525 - result = get_load_balance_assignment(sizes, num_gpus=num_gpus)
526 - (shuffle_indices, gpu_sample_counts, grouped_sizes_per_gpu) = result
527 -
528 - # Common assertions for all cases
529 - assert len(shuffle_indices) == len(sizes)
530 - assert len(gpu_sample_counts) == num_gpus
531 - assert len(grouped_sizes_per_gpu) == num_gpus
532 - assert sum(gpu_sample_counts) == len(sizes)
533 -
534 - assert shuffle_indices == expected_shuffle_indices
535 -
536 - assert gpu_sample_counts == expected_gpu_sample_counts
537 - assert grouped_sizes_per_gpu == expected_grouped_sizes_per_gpu
538 -
539 -
540 -class SimpleMRopeVisionModel(torch.nn.Module):
541 - """A simple vision model for testing mrope functionality."""
542 -
543 - def __init__(self, spatial_merge_size: int = 2, out_hidden_size: int = 64):
544 - super().__init__()
545 - self.spatial_merge_size = spatial_merge_size
546 - self.out_hidden_size = out_hidden_size
547 - self.linear = torch.nn.Linear(768, out_hidden_size)
548 -
549 - def forward(self, pixel_values: torch.Tensor,
550 - grid_thw_list: list[list[int]]):
551 - """Simple forward pass that simulates spatial merging."""
552 - # Apply linear transformation
553 - embeddings = self.linear(pixel_values)
554 -
555 - # Simulate spatial merging by reducing the number of patches
556 - merge_factor = self.spatial_merge_size * self.spatial_merge_size
557 -
558 - # Group patches and merge spatially
559 - merged_embeddings = []
560 - start_idx = 0
561 -
562 - for grid_thw in grid_thw_list:
563 - num_patches = math.prod(grid_thw)
564 - end_idx = start_idx + num_patches
565 -
566 - # Get patches for this image
567 - image_patches = embeddings[start_idx:end_idx]
568 -
569 - # Simulate spatial merging by averaging groups of patches
570 - merged_patches = num_patches // merge_factor
571 - if merged_patches > 0:
572 - # Reshape and average to simulate merging
573 - reshaped = image_patches[:merged_patches * merge_factor].view(
574 - merged_patches, merge_factor, -1)
575 - merged = reshaped.mean(dim=1)
576 - merged_embeddings.append(merged)
577 -
578 - start_idx = end_idx
579 -
580 - if merged_embeddings:
581 - return torch.cat(merged_embeddings, dim=0)
582 - else:
583 - return torch.empty((0, self.out_hidden_size),
584 - device=pixel_values.device,
585 - dtype=pixel_values.dtype)
586 -
587 -
588 -@multi_gpu_test(num_gpus=2)
589 -@pytest.mark.parametrize(
590 - "batch_size",
591 - [
592 - 1, # Single image
593 - 3, # Small batch
594 - 5, # Odd batch size (for testing padding)
595 - ],
596 -)
597 -def test_run_dp_sharded_mrope_vision_model(batch_size: int):
598 - world_size = 2
599 - # Launch processes
600 - mp.spawn(
601 - run_dp_sharded_mrope_vision_model_vs_direct,
602 - args=(
603 - world_size,
604 - batch_size,
605 - get_open_port(),
606 - ),
607 - nprocs=world_size,
608 - )
609 -
610 -
611 -def run_dp_sharded_mrope_vision_model_vs_direct(local_rank: int,
612 - world_size: int,
613 - batch_size: int,
614 - master_port: int):
615 - """
616 - Test that run_dp_sharded_mrope_vision_model produces the same results as
617 - calling the model directly.
618 - """
619 - # Set random seed for reproducibility
620 - current_platform.seed_everything(0)
621 - device = f"{current_platform.device_name}:{local_rank}"
622 - current_platform.set_device(device)
623 - torch.set_default_device(device)
624 -
625 - update_environment_variables({
626 - 'RANK': str(local_rank),
627 - 'LOCAL_RANK': str(local_rank),
628 - 'WORLD_SIZE': str(world_size),
629 - 'MASTER_ADDR': 'localhost',
630 - 'MASTER_PORT': str(master_port),
631 - })
632 -
633 - # initialize distributed
634 - init_distributed_environment()
635 - initialize_model_parallel(tensor_model_parallel_size=world_size)
636 -
637 - # Create test data
638 - grid_thw_list = []
639 - pixel_values_list = []
640 -
641 - for i in range(batch_size):
642 - # Varying image sizes for better testing
643 - t, h, w = 1, 4 + i, 4 + i
644 - grid_thw_list.append([t, h, w])
645 -
646 - num_patches = t * h * w
647 - # Create random pixel values for this image
648 - image_pixels = torch.randn(num_patches, 768)
649 - pixel_values_list.append(image_pixels)
650 -
651 - # Concatenate all pixel values
652 - pixel_values = torch.cat(pixel_values_list, dim=0)
653 -
654 - # Create a simple mrope vision model
655 - vision_model = SimpleMRopeVisionModel()
656 -
657 - # Run the model directly on the full input (only on rank 0)
658 - if local_rank == 0:
659 - with torch.inference_mode():
660 - direct_output = vision_model(pixel_values, grid_thw_list)
661 -
662 - # Run the model through the sharded function
663 - with torch.inference_mode():
664 - sharded_output = run_dp_sharded_mrope_vision_model(vision_model,
665 - pixel_values,
666 - grid_thw_list,
667 - rope_type="rope_3d")
668 - sharded_output = torch.cat(sharded_output, dim=0)
669 -
670 - # Check that the world size is set up correctly
671 - assert get_tensor_model_parallel_world_size() == world_size
672 -
673 - # Compare outputs (only on rank 0)
674 - if local_rank == 0:
675 - # Check that the outputs have the same shape
676 - assert direct_output.shape == sharded_output.shape
677 - # Check that the outputs are close (they should be identical)
678 - assert torch.allclose(direct_output,
679 - sharded_output,
680 - rtol=1e-5,
681 - atol=1e-5)
682 -
683 -
684 -@multi_gpu_test(num_gpus=2)
685 -def test_run_dp_sharded_mrope_vision_model_empty_input():
686 - world_size = 2
687 - mp.spawn(
688 - run_dp_sharded_mrope_vision_model_empty_input_worker,
689 - args=(world_size, get_open_port()),
690 - nprocs=world_size,
691 - )
692 -
693 -
694 -def run_dp_sharded_mrope_vision_model_empty_input_worker(
695 - local_rank: int, world_size: int, master_port: int):
696 - """Test run_dp_sharded_mrope_vision_model with empty input."""
697 - # Set up distributed environment
698 - device = f"{current_platform.device_name}:{local_rank}"
699 - current_platform.set_device(device)
700 - torch.set_default_device(device)
701 -
702 - update_environment_variables({
703 - 'RANK': str(local_rank),
704 - 'LOCAL_RANK': str(local_rank),
705 - 'WORLD_SIZE': str(world_size),
706 - 'MASTER_ADDR': 'localhost',
707 - 'MASTER_PORT': str(master_port),
708 - })
709 -
710 - init_distributed_environment()
711 - initialize_model_parallel(tensor_model_parallel_size=world_size)
712 -
713 - # Create empty inputs
714 - pixel_values = torch.empty((0, 768))
715 - grid_thw_list: list[list[int]] = []
716 -
717 - vision_model = SimpleMRopeVisionModel()
718 -
719 - # Should handle empty input gracefully
720 - with torch.inference_mode():
721 - output = run_dp_sharded_mrope_vision_model(vision_model,
722 - pixel_values,
723 - grid_thw_list,
724 - rope_type="rope_3d")
725 -
726 - assert len(output) == 0
727 -
728 -
729 -@multi_gpu_test(num_gpus=4)
730 -def test_run_dp_sharded_mrope_vision_model_uneven_load():
731 - world_size = 4
732 - mp.spawn(
733 - run_dp_sharded_mrope_vision_model_uneven_load_worker,
734 - args=(world_size, get_open_port()),
735 - nprocs=world_size,
736 - )
737 -
738 -
739 -def run_dp_sharded_mrope_vision_model_uneven_load_worker(
740 - local_rank: int, world_size: int, master_port: int):
741 - """Test run_dp_sharded_mrope_vision_model with uneven load distribution."""
742 - # Set up distributed environment
743 - current_platform.seed_everything(123)
744 - device = f"{current_platform.device_name}:{local_rank}"
745 - current_platform.set_device(device)
746 - torch.set_default_device(device)
747 -
748 - update_environment_variables({
749 - 'RANK': str(local_rank),
750 - 'LOCAL_RANK': str(local_rank),
751 - 'WORLD_SIZE': str(world_size),
752 - 'MASTER_ADDR': 'localhost',
753 - 'MASTER_PORT': str(master_port),
754 - })
755 -
756 - init_distributed_environment()
757 - initialize_model_parallel(tensor_model_parallel_size=world_size)
758 -
759 - # Create images with very different sizes
760 - grid_thw_list = [
761 - [1, 2, 2], # Small: 4 patches
762 - [1, 8, 8], # Large: 64 patches
763 - [1, 3, 3], # Medium: 9 patches
764 - ]
765 -
766 - pixel_values_list = []
767 - for grid_thw in grid_thw_list:
768 - num_patches = math.prod(grid_thw)
769 - image_pixels = torch.randn(num_patches, 768)
770 - pixel_values_list.append(image_pixels)
771 -
772 - pixel_values = torch.cat(pixel_values_list, dim=0)
773 - vision_model = SimpleMRopeVisionModel()
774 -
775 - # Should handle uneven distribution without errors
776 - with torch.inference_mode():
777 - output_tuple = run_dp_sharded_mrope_vision_model(vision_model,
778 - pixel_values,
779 - grid_thw_list,
780 - rope_type="rope_3d")
781 -
782 - # Verify output shape is reasonable
783 - merge_factor = vision_model.spatial_merge_size**2
784 - expected_output_patches = list(
785 - math.prod(grid_thw) // merge_factor for grid_thw in grid_thw_list)
786 -
787 - for i, output in enumerate(output_tuple):
788 - assert output.shape[0] == expected_output_patches[i]
789 - assert output.shape[1] == vision_model.out_hidden_size
790 -
791 -
792 -@pytest.mark.parametrize("spatial_merge_size", [2, 4])
793 -def test_simple_mrope_vision_model_spatial_merge(spatial_merge_size: int):
794 - """Test SimpleMRopeVisionModel with different spatial merge sizes."""
795 - device = current_platform.device_type
796 -
797 - grid_thw_list = [[1, 4, 4], [1, 6, 6]] # Two images
798 - pixel_values_list = []
799 -
800 - for grid_thw in grid_thw_list:
801 - num_patches = math.prod(grid_thw)
802 - image_pixels = torch.randn(num_patches, 768, device=device)
803 - pixel_values_list.append(image_pixels)
804 -
805 - pixel_values = torch.cat(pixel_values_list, dim=0)
806 - vision_model = SimpleMRopeVisionModel(
807 - spatial_merge_size=spatial_merge_size).to(device)
808 -
809 - with torch.inference_mode():
810 - output = vision_model(pixel_values, grid_thw_list)
811 -
812 - # Verify output dimensions based on spatial merging
813 - total_patches = sum(math.prod(grid_thw) for grid_thw in grid_thw_list)
814 - merge_factor = spatial_merge_size**2
815 - expected_output_patches = total_patches // merge_factor
816 -
817 - assert output.shape[0] == expected_output_patches
818 - assert output.shape[1] == vision_model.out_hidden_size
vllm/model_executor/models/glm4_1v.py
+1 -2
@@ -69,7 +69,6 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
69 BaseProcessingInfo, PromptReplacement,
70 PromptUpdate, PromptUpdateDetails)
71 from vllm.multimodal.profiling import BaseDummyInputsBuilder
72 -from vllm.multimodal.utils import run_dp_sharded_mrope_vision_model
72 from vllm.platforms import _Backend
73 from vllm.sequence import IntermediateTensors
74 from vllm.transformers_utils.config import uses_mrope
@@ -83,7 +82,7 @@ from .qwen2_vl import (_create_qwen2vl_field_factory,
82 from .utils import (AutoWeightsLoader, WeightsMapper,
83 init_vllm_registered_model, maybe_prefix,
84 merge_multimodal_embeddings)
86 -from .vision import get_vit_attn_backend
85 +from .vision import get_vit_attn_backend, run_dp_sharded_mrope_vision_model
86
87 logger = init_logger(__name__)
88
vllm/model_executor/models/idefics2_vision_model.py
+2 -1
@@ -34,7 +34,8 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
34 RowParallelLinear)
35 from vllm.model_executor.layers.quantization import QuantizationConfig
36 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
37 -from vllm.multimodal.utils import run_dp_sharded_vision_model
37 +
38 +from .vision import run_dp_sharded_vision_model
39
40
41 class Idefics2VisionEmbeddings(nn.Module):
vllm/model_executor/models/intern_vit.py
+2 -1
@@ -28,7 +28,8 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
28 RowParallelLinear)
29 from vllm.model_executor.layers.quantization import QuantizationConfig
30 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
31 -from vllm.multimodal.utils import run_dp_sharded_vision_model
31 +
32 +from .vision import run_dp_sharded_vision_model
33
34 NORM2FN = {
35 'rms_norm': RMSNorm,
vllm/model_executor/models/kimi_vl.py
+1 -1
@@ -76,13 +76,13 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
76 BaseProcessingInfo, PromptReplacement,
77 PromptUpdate)
78 from vllm.multimodal.profiling import BaseDummyInputsBuilder
79 -from vllm.multimodal.utils import run_dp_sharded_mrope_vision_model
79 from vllm.sequence import IntermediateTensors
80 from vllm.transformers_utils.configs import KimiVLConfig, MoonViTConfig
81 from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekV2Config
82 from vllm.utils.tensor_schema import TensorSchema, TensorShape
83
84 from .utils import PPMissingLayer, is_pp_missing_parameter, maybe_prefix
85 +from .vision import run_dp_sharded_mrope_vision_model
86
87
88 # For dummy input only
vllm/model_executor/models/mllama4.py
+1 -1
@@ -50,7 +50,6 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
50 BaseProcessingInfo, PromptReplacement,
51 PromptUpdate, PromptUpdateDetails)
52 from vllm.multimodal.profiling import BaseDummyInputsBuilder
53 -from vllm.multimodal.utils import run_dp_sharded_vision_model
53 from vllm.sequence import IntermediateTensors
54 from vllm.utils.tensor_schema import TensorSchema, TensorShape
55
@@ -58,6 +57,7 @@ from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
57 from .llama4 import Llama4ForCausalLM
58 from .utils import (AutoWeightsLoader, flatten_bn, maybe_prefix,
59 merge_multimodal_embeddings)
60 +from .vision import run_dp_sharded_vision_model
61
62
63 class Llama4ImagePatchInputs(TensorSchema):
vllm/model_executor/models/qwen2_5_vl.py
+1 -2
@@ -59,7 +59,6 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
59 from vllm.model_executor.models.module_mapping import MultiModelKeys
60 from vllm.multimodal import MULTIMODAL_REGISTRY
61 from vllm.multimodal.inputs import MultiModalFieldConfig
62 -from vllm.multimodal.utils import run_dp_sharded_mrope_vision_model
62 from vllm.platforms import _Backend
63 from vllm.sequence import IntermediateTensors
64 from vllm.transformers_utils.config import uses_mrope
@@ -74,7 +73,7 @@ from .qwen2_vl import (Qwen2VLMultiModalProcessor, Qwen2VLProcessingInfo,
73 from .utils import (AutoWeightsLoader, WeightsMapper, cast_overflow_tensors,
74 init_vllm_registered_model, maybe_prefix,
75 merge_multimodal_embeddings)
77 -from .vision import get_vit_attn_backend
76 +from .vision import get_vit_attn_backend, run_dp_sharded_mrope_vision_model
77
78 logger = init_logger(__name__)
79
vllm/model_executor/models/qwen2_vl.py
+1 -2
@@ -66,7 +66,6 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
66 BaseProcessingInfo, PromptReplacement,
67 PromptUpdate)
68 from vllm.multimodal.profiling import BaseDummyInputsBuilder
69 -from vllm.multimodal.utils import run_dp_sharded_mrope_vision_model
69 from vllm.platforms import _Backend, current_platform
70 from vllm.sequence import IntermediateTensors
71 from vllm.transformers_utils.config import uses_mrope
@@ -78,7 +77,7 @@ from .interfaces import (MultiModalEmbeddings, SupportsLoRA, SupportsMRoPE,
77 from .utils import (AutoWeightsLoader, WeightsMapper,
78 init_vllm_registered_model, maybe_prefix,
79 merge_multimodal_embeddings)
81 -from .vision import get_vit_attn_backend
80 +from .vision import get_vit_attn_backend, run_dp_sharded_mrope_vision_model
81
82 logger = init_logger(__name__)
83
vllm/model_executor/models/qwen3_vl.py
+1 -5
@@ -83,7 +83,7 @@ from .qwen2_vl import Qwen2VLProcessingInfo
83 from .qwen3 import Qwen3ForCausalLM, Qwen3Model
84 from .utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper,
85 maybe_prefix, merge_multimodal_embeddings)
86 -from .vision import get_vit_attn_backend
86 +from .vision import get_vit_attn_backend, run_dp_sharded_mrope_vision_model
87
88 logger = init_logger(__name__)
89
@@ -1214,8 +1214,6 @@ class Qwen3VLForConditionalGeneration(nn.Module, SupportsMultiModal,
1214 else:
1215 pixel_values = image_input["pixel_values"].type(self.visual.dtype)
1216 if self.use_data_parallel:
1217 - from vllm.multimodal.utils import (
1218 - run_dp_sharded_mrope_vision_model)
1217 return run_dp_sharded_mrope_vision_model(self.visual,
1218 pixel_values,
1219 grid_thw_list,
@@ -1245,8 +1243,6 @@ class Qwen3VLForConditionalGeneration(nn.Module, SupportsMultiModal,
1243 pixel_values_videos = video_input["pixel_values_videos"].type(
1244 self.visual.dtype)
1245 if self.use_data_parallel:
1248 - from vllm.multimodal.utils import (
1249 - run_dp_sharded_mrope_vision_model)
1246 return run_dp_sharded_mrope_vision_model(self.visual,
1247 pixel_values_videos,
1248 grid_thw_list,
vllm/model_executor/models/step3_vl.py
+1 -1
@@ -31,7 +31,6 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
31 BaseProcessingInfo, PromptReplacement,
32 PromptUpdate, PromptUpdateDetails)
33 from vllm.multimodal.profiling import BaseDummyInputsBuilder
34 -from vllm.multimodal.utils import run_dp_sharded_vision_model
34 from vllm.sequence import IntermediateTensors
35 from vllm.transformers_utils.configs import Step3VisionEncoderConfig
36 from vllm.transformers_utils.tokenizer import AnyTokenizer
@@ -40,6 +39,7 @@ from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
39 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
40 init_vllm_registered_model, maybe_prefix,
41 merge_multimodal_embeddings)
42 +from .vision import run_dp_sharded_vision_model
43
44
45 class Step3VLImagePixelInputs(TypedDict):
vllm/model_executor/models/vision.py
+280 -1
@@ -1,12 +1,17 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import itertools
5 +import math
6 from abc import ABC, abstractmethod
5 -from typing import Final, Generic, Optional, Protocol, TypeVar, Union
7 +from typing import Final, Generic, Literal, Optional, Protocol, TypeVar, Union
8
9 import torch
10 from transformers import PretrainedConfig
11
12 +from vllm.distributed import (get_tensor_model_parallel_rank,
13 + get_tensor_model_parallel_world_size,
14 + tensor_model_parallel_all_gather)
15 from vllm.logger import init_logger
16 from vllm.platforms import _Backend, current_platform
17
@@ -123,3 +128,277 @@ def resolve_visual_encoder_outputs(
128 if post_layer_norm is not None and uses_last_layer:
129 hs_pool[-1] = post_layer_norm(encoder_outputs)
130 return torch.cat(hs_pool, dim=-1)
131 +
132 +
133 +def run_dp_sharded_vision_model(image_input: torch.Tensor,
134 + vision_model: torch.nn.Module) -> torch.Tensor:
135 + """Run a vision model with data parallelism (DP) sharding. The function
136 + will shard the input image tensor on the first dimension and run the vision
137 + model
138 +
139 + Args:
140 + image_input (torch.Tensor): Image input tensor.
141 + vision_model (torch.nn.Module): Vision model.
142 + Returns:
143 + torch.Tensor: Output image embeddings
144 + """
145 +
146 + num_chunks = image_input.shape[0]
147 + mp_world_size = get_tensor_model_parallel_world_size()
148 + num_chunks_per_rank = (num_chunks + mp_world_size - 1) // mp_world_size
149 + num_padded_chunks = num_chunks_per_rank * mp_world_size - num_chunks
150 + pad = (0, ) * (2 * (image_input.dim() - 1)) + (0, num_padded_chunks)
151 + image_input_padded = torch.nn.functional.pad(image_input, pad)
152 + rank = get_tensor_model_parallel_rank()
153 + image_input_per_rank = image_input_padded[rank *
154 + num_chunks_per_rank:(rank + 1) *
155 + num_chunks_per_rank, ...]
156 +
157 + vision_embeddings = vision_model(image_input_per_rank)
158 + # Ensure tensor is contiguous before all_gather
159 + vision_embeddings = vision_embeddings.contiguous()
160 + vision_embeddings = tensor_model_parallel_all_gather(vision_embeddings,
161 + dim=0)
162 + vision_embeddings = vision_embeddings[:num_chunks, ...]
163 + return vision_embeddings
164 +
165 +
166 +def get_load_balance_assignment(
167 + sizes: list[int],
168 + num_gpus: int = 2,
169 +) -> tuple[list[int], list[int], list[int]]:
170 + """
171 + Generate load balancing assignment and metadata
172 + for distributing data across GPUs.
173 + The load is determined by the total image sizes,
174 + not the number of images.
175 +
176 + Args:
177 + sizes: The size of each image
178 + num_gpus: Number of GPUs to balance across
179 +
180 + Returns:
181 + shuffle_indices:
182 + Indices to reorder data for balanced loading
183 + gpu_sample_counts:
184 + Number of samples assigned to each GPU
185 + grouped_sizes_per_gpu:
186 + Total size assigned to each GPU
187 +
188 + Example:
189 + ```
190 + sizes = [1000, 100, 200, 50]
191 + num_gpus=2
192 + ```
193 +
194 + """
195 +
196 + n_samples = len(sizes)
197 +
198 + # Handle edge cases
199 + if n_samples == 0:
200 + return [], [0] * num_gpus, [0] * num_gpus
201 +
202 + # Use greedy algorithm - balance by total size, not sample count
203 + gpu_assignments = [list[int]() for _ in range(num_gpus)]
204 + gpu_loads = [0] * num_gpus # This tracks total SIZE, not sample count
205 +
206 + # Sort indices by size (largest first for better load balancing)
207 + # sizes = [1000, 100, 200, 50]
208 + # large_to_small_indices = [0, 2, 1, 3]
209 + large_to_small_indices = sorted(range(n_samples),
210 + key=lambda i: sizes[i],
211 + reverse=True)
212 +
213 + for idx in large_to_small_indices:
214 + # Find GPU with minimum current load (by total size)
215 + min_gpu = min(range(num_gpus), key=lambda i: gpu_loads[i])
216 + gpu_assignments[min_gpu].append(idx)
217 + gpu_loads[min_gpu] += sizes[idx]
218 +
219 + # Create shuffle indices and counts
220 + shuffle_indices = list[int]()
221 + gpu_sample_counts = list[int]()
222 + for gpu_id in range(num_gpus):
223 + # GPU_0 = [1000] = [0]
224 + # GPU_1 = [200, 100, 50] = [2, 1, 3]
225 + # shuffle_indices = [0, 2, 1, 3]
226 + shuffle_indices.extend(gpu_assignments[gpu_id])
227 + # GPU_0 = [1]
228 + # GPU_1 = [3]
229 + # gpu_sample_counts = [1, 3]
230 + gpu_sample_counts.append(len(gpu_assignments[gpu_id]))
231 +
232 + return (shuffle_indices, gpu_sample_counts, gpu_loads)
233 +
234 +
235 +def run_dp_sharded_mrope_vision_model(
236 + vision_model: torch.nn.Module,
237 + pixel_values: torch.Tensor,
238 + grid_thw_list: list[list[int]],
239 + *,
240 + rope_type: Literal["rope_3d", "rope_2d"],
241 +) -> tuple[torch.Tensor, ...]:
242 + """Run a vision model with data parallelism (DP) sharding.
243 + The function will shard the input image tensor on the
244 + first dimension and run the vision model.
245 + This function is used to run the vision model with mrope.
246 +
247 + Args:
248 + vision_model (torch.nn.Module): Vision model.
249 + pixel_values (torch.Tensor): Image/Video input tensor.
250 + grid_thw_list: List of grid dimensions for each image
251 + rope_type: Type of rope used in the vision model.
252 + Different rope types have different dimension to do ViT.
253 + "rope_3d" for 3D rope (e.g., Qwen2.5-VL)
254 + "rope_2d" for 2D rope (e.g., Kimi-VL)
255 + Returns:
256 + torch.Tensor: Output image embeddings
257 +
258 + Example:
259 + ```
260 + vision_model.out_hidden_size = 64
261 + vision_model.spatial_merge_size = 2
262 + pixel_values.shape = (1350, channel)
263 + grid_thw_list = [[1, 10, 100], [1, 10, 10], [1, 10, 20], [1, 50]]
264 + tp_size=2
265 + ```
266 +
267 + """
268 + tp_size = get_tensor_model_parallel_world_size()
269 +
270 + # GPU_0 tp_rank_local = 0
271 + # GPU_1 tp_rank_local = 1
272 + tp_rank_local = get_tensor_model_parallel_rank()
273 +
274 + # patches_per_image = [1000, 100, 200, 50]
275 + patches_per_image = [math.prod(grid_thw) for grid_thw in grid_thw_list]
276 + # patches_per_image = [0, 1000, 1100, 1300, 1350]
277 + cum_patches_per_image = [0, *itertools.accumulate(patches_per_image)]
278 +
279 + # Get load balancing assignment with all metadata
280 + # image_to_tp_rank = [0, 2, 1, 3]
281 + # gpu_sample_counts = [1, 3]
282 + # grouped_pixel_values_len = [1000, 350]
283 + (image_to_tp_rank, gpu_sample_counts,
284 + grouped_pixel_values_len) = get_load_balance_assignment(
285 + patches_per_image, tp_size)
286 +
287 + # cu_gpu_sample_counts = [0, 1, 4]
288 + cum_gpu_sample_counts = [0, *itertools.accumulate(gpu_sample_counts)]
289 +
290 + # GPU_0 image_idxs_local = [0]
291 + # GPU_1 image_idxs_local = [2, 1, 3]
292 + image_idxs_local = image_to_tp_rank[cum_gpu_sample_counts[tp_rank_local]:
293 + cum_gpu_sample_counts[tp_rank_local +
294 + 1]]
295 +
296 + # Get the pixel values for the local images based on the image_idxs_local
297 + if len(image_idxs_local) > 0:
298 + pixel_values_local = torch.cat([
299 + pixel_values[cum_patches_per_image[i]:cum_patches_per_image[i + 1]]
300 + for i in image_idxs_local
301 + ])
302 + else:
303 + # Handle case where this rank has no images
304 + pixel_values_local = torch.empty((0, pixel_values.shape[1]),
305 + device=pixel_values.device,
306 + dtype=pixel_values.dtype)
307 + # embed_dim_reduction_factor = 2 * 2
308 + if rope_type == "rope_2d":
309 + embed_dim_reduction_factor = (vision_model.merge_kernel_size[0] *
310 + vision_model.merge_kernel_size[1])
311 + else:
312 + embed_dim_reduction_factor = (vision_model.spatial_merge_size *
313 + vision_model.spatial_merge_size)
314 +
315 + # Find the max length across all ranks
316 + # The output embedding of every DP rank has to be
317 + # padded to this length for tensor_model_parallel_all_gather
318 + # to work
319 + max_len_per_rank = max(
320 + grouped_pixel_values_len) // embed_dim_reduction_factor
321 + local_grid_thw_list = [grid_thw_list[i] for i in image_idxs_local]
322 +
323 + # Run the vision model on the local pixel_values_local
324 + if rope_type == "rope_2d":
325 + if pixel_values_local.shape[0] > 0:
326 + image_embeds_local = vision_model(
327 + pixel_values_local, torch.tensor(local_grid_thw_list))
328 + if isinstance(image_embeds_local, list):
329 + image_embeds_local = torch.cat(image_embeds_local, dim=0)
330 + else:
331 + out_dim = getattr(vision_model.config, "hidden_size", None)
332 + image_embeds_local = torch.empty(
333 + (0, embed_dim_reduction_factor, out_dim),
334 + device=pixel_values.device,
335 + dtype=pixel_values.dtype)
336 + else:
337 + if pixel_values_local.shape[0] > 0:
338 + image_embeds_local = vision_model(pixel_values_local,
339 + local_grid_thw_list)
340 + else:
341 + # Handle empty case
342 + image_embeds_local = torch.empty((0, vision_model.out_hidden_size),
343 + device=pixel_values.device,
344 + dtype=pixel_values.dtype)
345 +
346 + # Pad the output based on max_len_per_rank
347 + # for tensor_model_parallel_all_gather to work
348 + current_len = image_embeds_local.shape[0]
349 + if current_len < max_len_per_rank:
350 + padding_size = max_len_per_rank - current_len
351 + if rope_type == "rope_2d":
352 + padding = torch.empty((padding_size, image_embeds_local.shape[1],
353 + image_embeds_local.shape[2]),
354 + dtype=image_embeds_local.dtype,
355 + device=image_embeds_local.device)
356 + else:
357 + padding = torch.empty((padding_size, image_embeds_local.shape[1]),
358 + dtype=image_embeds_local.dtype,
359 + device=image_embeds_local.device)
360 + image_embeds_local_padded = torch.cat([image_embeds_local, padding],
361 + dim=0)
362 + else:
363 + image_embeds_local_padded = image_embeds_local
364 +
365 + # Do all_gather to collect embeddings from all ranks
366 + gathered_embeds = tensor_model_parallel_all_gather(
367 + image_embeds_local_padded, dim=0)
368 +
369 + # Remove padding and reconstruct per-rank embeddings
370 + rank_embeddings = list[torch.Tensor]()
371 + for rank in range(tp_size):
372 + start_idx = rank * max_len_per_rank
373 + end_idx = start_idx + (grouped_pixel_values_len[rank] //
374 + embed_dim_reduction_factor)
375 + rank_embeddings.append(gathered_embeds[start_idx:end_idx])
376 +
377 + patches_per_output_image = [(patch_size // embed_dim_reduction_factor)
378 + for patch_size in patches_per_image]
379 +
380 + # Reconstruct embeddings in the original order
381 + original_order_embeddings = [None] * len(grid_thw_list)
382 + current_idx = 0
383 + for rank in range(tp_size):
384 + count = gpu_sample_counts[rank]
385 + if count > 0:
386 + # Get images assigned to this rank in shuffled order
387 + # GPU_0 = image_idxs_local [0]
388 + # GPU_1 = image_idxs_local [2, 1, 3]
389 + rank_images = image_to_tp_rank[current_idx:current_idx + count]
390 +
391 + rank_embed = rank_embeddings[rank]
392 + # Split rank embeddings back to individual images
393 + embed_start = 0
394 + for img_idx in rank_images:
395 + img_patches = patches_per_output_image[img_idx]
396 + original_order_embeddings[img_idx] = rank_embed[
397 + embed_start:embed_start + img_patches]
398 + embed_start += img_patches
399 + current_idx += count
400 + out_embeddings = tuple(embed for embed in original_order_embeddings
401 + if embed is not None)
402 + assert len(out_embeddings) == len(
403 + original_order_embeddings), "Found unassigned embeddings"
404 + return out_embeddings
vllm/multimodal/utils.py
+6 -287
@@ -3,13 +3,11 @@
3
4 import asyncio
5 import atexit
6 -import itertools
7 -import math
6 from collections.abc import Iterable
7 from concurrent.futures import ThreadPoolExecutor
8 from itertools import groupby
9 from pathlib import Path
12 -from typing import TYPE_CHECKING, Any, Literal, Optional, TypeVar, Union
10 +from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
11 from urllib.parse import ParseResult, urlparse
12 from urllib.request import url2pathname
13
@@ -21,9 +19,6 @@ from typing_extensions import deprecated
19
20 import vllm.envs as envs
21 from vllm.connections import HTTPConnection, global_http_connection
24 -from vllm.distributed import (get_tensor_model_parallel_rank,
25 - get_tensor_model_parallel_world_size,
26 - tensor_model_parallel_all_gather)
22
23 from .audio import AudioMediaIO
24 from .base import MediaIO
@@ -33,12 +28,10 @@ from .video import VideoMediaIO
28 _M = TypeVar("_M")
29
30 if TYPE_CHECKING:
36 - from .inputs import (BatchedTensorInputs, MultiModalKwargs,
37 - MultiModalKwargsItem, MultiModalKwargsItems,
38 - MultiModalPlaceholderDict)
31 + from .inputs import (BatchedTensorInputs, MultiModalKwargsItem,
32 + MultiModalKwargsItems, MultiModalPlaceholderDict)
33 else:
34 BatchedTensorInputs = Any
41 - MultiModalKwargs = Any
35 MultiModalKwargsItem = Any
36 MultiModalKwargsItems = Any
37 MultiModalPlaceholderDict = Any
@@ -93,7 +86,7 @@ class MediaConnector:
86 self,
87 url_spec: ParseResult,
88 media_io: MediaIO[_M],
96 - ) -> _M:
89 + ) -> _M: # type: ignore[type-var]
90 data_spec, data = url_spec.path.split(",", 1)
91 media_type, data_type = data_spec.split(";", 1)
92
@@ -107,7 +100,7 @@ class MediaConnector:
100 self,
101 url_spec: ParseResult,
102 media_io: MediaIO[_M],
110 - ) -> _M:
103 + ) -> _M: # type: ignore[type-var]
104 allowed_local_media_path = self.allowed_local_media_path
105 if allowed_local_media_path is None:
106 raise RuntimeError("Cannot load local files without "
@@ -127,7 +120,7 @@ class MediaConnector:
120 media_io: MediaIO[_M],
121 *,
122 fetch_timeout: Optional[int] = None,
130 - ) -> _M:
123 + ) -> _M: # type: ignore[type-var]
124 url_spec = urlparse(url)
125
126 if url_spec.scheme.startswith("http"):
@@ -434,280 +427,6 @@ def group_mm_kwargs_by_modality(
427 yield modality, len(items_lst), mm_kwargs_group
428
429
437 -def run_dp_sharded_vision_model(image_input: torch.Tensor,
438 - vision_model: torch.nn.Module) -> torch.Tensor:
439 - """Run a vision model with data parallelism (DP) sharding. The function
440 - will shard the input image tensor on the first dimension and run the vision
441 - model
442 -
443 - Args:
444 - image_input (torch.Tensor): Image input tensor.
445 - vision_model (torch.nn.Module): Vision model.
446 - Returns:
447 - torch.Tensor: Output image embeddings
448 - """
449 -
450 - num_chunks = image_input.shape[0]
451 - mp_world_size = get_tensor_model_parallel_world_size()
452 - num_chunks_per_rank = (num_chunks + mp_world_size - 1) // mp_world_size
453 - num_padded_chunks = num_chunks_per_rank * mp_world_size - num_chunks
454 - pad = (0, ) * (2 * (image_input.dim() - 1)) + (0, num_padded_chunks)
455 - image_input_padded = torch.nn.functional.pad(image_input, pad)
456 - rank = get_tensor_model_parallel_rank()
457 - image_input_per_rank = image_input_padded[rank *
458 - num_chunks_per_rank:(rank + 1) *
459 - num_chunks_per_rank, ...]
460 -
461 - vision_embeddings = vision_model(image_input_per_rank)
462 - # Ensure tensor is contiguous before all_gather
463 - vision_embeddings = vision_embeddings.contiguous()
464 - vision_embeddings = tensor_model_parallel_all_gather(vision_embeddings,
465 - dim=0)
466 - vision_embeddings = vision_embeddings[:num_chunks, ...]
467 - return vision_embeddings
468 -
469 -
470 -def get_load_balance_assignment(
471 - sizes: list[int],
472 - num_gpus: int = 2,
473 -) -> tuple[list[int], list[int], list[int]]:
474 - """
475 - Generate load balancing assignment and metadata
476 - for distributing data across GPUs.
477 - The load is determined by the total image sizes,
478 - not the number of images.
479 -
480 - Args:
481 - sizes: The size of each image
482 - num_gpus: Number of GPUs to balance across
483 -
484 - Returns:
485 - shuffle_indices:
486 - Indices to reorder data for balanced loading
487 - gpu_sample_counts:
488 - Number of samples assigned to each GPU
489 - grouped_sizes_per_gpu:
490 - Total size assigned to each GPU
491 -
492 - Example:
493 - ```
494 - sizes = [1000, 100, 200, 50]
495 - num_gpus=2
496 - ```
497 -
498 - """
499 -
500 - n_samples = len(sizes)
501 -
502 - # Handle edge cases
503 - if n_samples == 0:
504 - return [], [0] * num_gpus, [0] * num_gpus
505 -
506 - # Use greedy algorithm - balance by total size, not sample count
507 - gpu_assignments = [list[int]() for _ in range(num_gpus)]
508 - gpu_loads = [0] * num_gpus # This tracks total SIZE, not sample count
509 -
510 - # Sort indices by size (largest first for better load balancing)
511 - # sizes = [1000, 100, 200, 50]
512 - # large_to_small_indices = [0, 2, 1, 3]
513 - large_to_small_indices = sorted(range(n_samples),
514 - key=lambda i: sizes[i],
515 - reverse=True)
516 -
517 - for idx in large_to_small_indices:
518 - # Find GPU with minimum current load (by total size)
519 - min_gpu = min(range(num_gpus), key=lambda i: gpu_loads[i])
520 - gpu_assignments[min_gpu].append(idx)
521 - gpu_loads[min_gpu] += sizes[idx]
522 -
523 - # Create shuffle indices and counts
524 - shuffle_indices = list[int]()
525 - gpu_sample_counts = list[int]()
526 - for gpu_id in range(num_gpus):
527 - # GPU_0 = [1000] = [0]
528 - # GPU_1 = [200, 100, 50] = [2, 1, 3]
529 - # shuffle_indices = [0, 2, 1, 3]
530 - shuffle_indices.extend(gpu_assignments[gpu_id])
531 - # GPU_0 = [1]
532 - # GPU_1 = [3]
533 - # gpu_sample_counts = [1, 3]
534 - gpu_sample_counts.append(len(gpu_assignments[gpu_id]))
535 -
536 - return (shuffle_indices, gpu_sample_counts, gpu_loads)
537 -
538 -
539 -def run_dp_sharded_mrope_vision_model(
540 - vision_model: torch.nn.Module,
541 - pixel_values: torch.Tensor,
542 - grid_thw_list: list[list[int]],
543 - *,
544 - rope_type: Literal["rope_3d", "rope_2d"],
545 -) -> tuple[torch.Tensor, ...]:
546 - """Run a vision model with data parallelism (DP) sharding.
547 - The function will shard the input image tensor on the
548 - first dimension and run the vision model.
549 - This function is used to run the vision model with mrope.
550 -
551 - Args:
552 - vision_model (torch.nn.Module): Vision model.
553 - pixel_values (torch.Tensor): Image/Video input tensor.
554 - grid_thw_list: List of grid dimensions for each image
555 - rope_type: Type of rope used in the vision model.
556 - Different rope types have different dimension to do ViT.
557 - "rope_3d" for 3D rope (e.g., Qwen2.5-VL)
558 - "rope_2d" for 2D rope (e.g., Kimi-VL)
559 - Returns:
560 - torch.Tensor: Output image embeddings
561 -
562 - Example:
563 - ```
564 - vision_model.out_hidden_size = 64
565 - vision_model.spatial_merge_size = 2
566 - pixel_values.shape = (1350, channel)
567 - grid_thw_list = [[1, 10, 100], [1, 10, 10], [1, 10, 20], [1, 50]]
568 - tp_size=2
569 - ```
570 -
571 - """
572 - tp_size = get_tensor_model_parallel_world_size()
573 -
574 - # GPU_0 tp_rank_local = 0
575 - # GPU_1 tp_rank_local = 1
576 - tp_rank_local = get_tensor_model_parallel_rank()
577 -
578 - # patches_per_image = [1000, 100, 200, 50]
579 - patches_per_image = [math.prod(grid_thw) for grid_thw in grid_thw_list]
580 - # patches_per_image = [0, 1000, 1100, 1300, 1350]
581 - cum_patches_per_image = [0, *itertools.accumulate(patches_per_image)]
582 -
583 - # Get load balancing assignment with all metadata
584 - # image_to_tp_rank = [0, 2, 1, 3]
585 - # gpu_sample_counts = [1, 3]
586 - # grouped_pixel_values_len = [1000, 350]
587 - (image_to_tp_rank, gpu_sample_counts,
588 - grouped_pixel_values_len) = get_load_balance_assignment(
589 - patches_per_image, tp_size)
590 -
591 - # cu_gpu_sample_counts = [0, 1, 4]
592 - cum_gpu_sample_counts = [0, *itertools.accumulate(gpu_sample_counts)]
593 -
594 - # GPU_0 image_idxs_local = [0]
595 - # GPU_1 image_idxs_local = [2, 1, 3]
596 - image_idxs_local = image_to_tp_rank[cum_gpu_sample_counts[tp_rank_local]:
597 - cum_gpu_sample_counts[tp_rank_local +
598 - 1]]
599 -
600 - # Get the pixel values for the local images based on the image_idxs_local
601 - if len(image_idxs_local) > 0:
602 - pixel_values_local = torch.cat([
603 - pixel_values[cum_patches_per_image[i]:cum_patches_per_image[i + 1]]
604 - for i in image_idxs_local
605 - ])
606 - else:
607 - # Handle case where this rank has no images
608 - pixel_values_local = torch.empty((0, pixel_values.shape[1]),
609 - device=pixel_values.device,
610 - dtype=pixel_values.dtype)
611 - # embed_dim_reduction_factor = 2 * 2
612 - if rope_type == "rope_2d":
613 - embed_dim_reduction_factor = (vision_model.merge_kernel_size[0] *
614 - vision_model.merge_kernel_size[1])
615 - else:
616 - embed_dim_reduction_factor = (vision_model.spatial_merge_size *
617 - vision_model.spatial_merge_size)
618 -
619 - # Find the max length across all ranks
620 - # The output embedding of every DP rank has to be
621 - # padded to this length for tensor_model_parallel_all_gather
622 - # to work
623 - max_len_per_rank = max(
624 - grouped_pixel_values_len) // embed_dim_reduction_factor
625 - local_grid_thw_list = [grid_thw_list[i] for i in image_idxs_local]
626 -
627 - # Run the vision model on the local pixel_values_local
628 - if rope_type == "rope_2d":
629 - if pixel_values_local.shape[0] > 0:
630 - image_embeds_local = vision_model(
631 - pixel_values_local, torch.tensor(local_grid_thw_list))
632 - if isinstance(image_embeds_local, list):
633 - image_embeds_local = torch.cat(image_embeds_local, dim=0)
634 - else:
635 - out_dim = getattr(vision_model.config, "hidden_size", None)
636 - image_embeds_local = torch.empty(
637 - (0, embed_dim_reduction_factor, out_dim),
638 - device=pixel_values.device,
639 - dtype=pixel_values.dtype)
640 - else:
641 - if pixel_values_local.shape[0] > 0:
642 - image_embeds_local = vision_model(pixel_values_local,
643 - local_grid_thw_list)
644 - else:
645 - # Handle empty case
646 - image_embeds_local = torch.empty((0, vision_model.out_hidden_size),
647 - device=pixel_values.device,
648 - dtype=pixel_values.dtype)
649 -
650 - # Pad the output based on max_len_per_rank
651 - # for tensor_model_parallel_all_gather to work
652 - current_len = image_embeds_local.shape[0]
653 - if current_len < max_len_per_rank:
654 - padding_size = max_len_per_rank - current_len
655 - if rope_type == "rope_2d":
656 - padding = torch.empty((padding_size, image_embeds_local.shape[1],
657 - image_embeds_local.shape[2]),
658 - dtype=image_embeds_local.dtype,
659 - device=image_embeds_local.device)
660 - else:
661 - padding = torch.empty((padding_size, image_embeds_local.shape[1]),
662 - dtype=image_embeds_local.dtype,
663 - device=image_embeds_local.device)
664 - image_embeds_local_padded = torch.cat([image_embeds_local, padding],
665 - dim=0)
666 - else:
667 - image_embeds_local_padded = image_embeds_local
668 -
669 - # Do all_gather to collect embeddings from all ranks
670 - gathered_embeds = tensor_model_parallel_all_gather(
671 - image_embeds_local_padded, dim=0)
672 -
673 - # Remove padding and reconstruct per-rank embeddings
674 - rank_embeddings = list[torch.Tensor]()
675 - for rank in range(tp_size):
676 - start_idx = rank * max_len_per_rank
677 - end_idx = start_idx + (grouped_pixel_values_len[rank] //
678 - embed_dim_reduction_factor)
679 - rank_embeddings.append(gathered_embeds[start_idx:end_idx])
680 -
681 - patches_per_output_image = [(patch_size // embed_dim_reduction_factor)
682 - for patch_size in patches_per_image]
683 -
684 - # Reconstruct embeddings in the original order
685 - original_order_embeddings = [None] * len(grid_thw_list)
686 - current_idx = 0
687 - for rank in range(tp_size):
688 - count = gpu_sample_counts[rank]
689 - if count > 0:
690 - # Get images assigned to this rank in shuffled order
691 - # GPU_0 = image_idxs_local [0]
692 - # GPU_1 = image_idxs_local [2, 1, 3]
693 - rank_images = image_to_tp_rank[current_idx:current_idx + count]
694 -
695 - rank_embed = rank_embeddings[rank]
696 - # Split rank embeddings back to individual images
697 - embed_start = 0
698 - for img_idx in rank_images:
699 - img_patches = patches_per_output_image[img_idx]
700 - original_order_embeddings[img_idx] = rank_embed[
701 - embed_start:embed_start + img_patches]
702 - embed_start += img_patches
703 - current_idx += count
704 - out_embeddings = tuple(embed for embed in original_order_embeddings
705 - if embed is not None)
706 - assert len(out_embeddings) == len(
707 - original_order_embeddings), "Found unassigned embeddings"
708 - return out_embeddings
709 -
710 -
430 def fetch_audio(
431 audio_url: str,
432 audio_io_kwargs: Optional[dict[str, Any]] = None,