@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1298c67795
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

[FEAT] [Performance] Enable DP for ViT in Qwen2.5VL (#22742)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

TJian committed Aug 19, 2025 at 08:25 UTC 1298c67795b5955dc6c3ffc28a172b465046d8c9
5 files changed +633 -48
tests/multimodal/test_utils.py
+324 -2
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import base64
5 +import math
6 import mimetypes
7 import os
8 from tempfile import NamedTemporaryFile, TemporaryDirectory
@@ -20,6 +21,8 @@ from vllm.distributed.parallel_state import (init_distributed_environment,
21 from vllm.multimodal.image import convert_image_mode
22 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
@@ -425,8 +428,8 @@ def run_dp_sharded_vision_model_vs_direct(local_rank: int, world_size: int,
428 # Set random seed for reproducibility
429 current_platform.seed_everything(0)
430
428 - device = torch.device(f"cuda:{local_rank}")
429 - torch.cuda.set_device(device)
431 + device = f"{current_platform.device_name}:{local_rank}"
432 + current_platform.set_device(device)
433 torch.set_default_device(device)
434
435 update_environment_variables({
@@ -463,3 +466,322 @@ def run_dp_sharded_vision_model_vs_direct(local_rank: int, world_size: int,
466
467 # Check that the outputs are close (they should be identical)
468 assert torch.allclose(direct_output, sharded_output, rtol=1e-5, atol=1e-5)
469 +
470 +
471 +@pytest.mark.parametrize(
472 + "sizes,num_gpus,expected_shuffle_indices,expected_gpu_sample_counts,"
473 + "expected_grouped_sizes_per_gpu,test_description",
474 + [
475 + # Empty input
476 + ([], 2, [], [0, 0], [0, 0], "empty input"),
477 +
478 + # Fewer samples than GPUs
479 + ([100, 200], 4, [1, 0], [1, 1, 0, 0], [200, 100, 0, 0
480 + ], "fewer samples than GPUs"),
481 +
482 + # Single GPU
483 + ([100, 200, 300], 1, [2, 1, 0], [3], [600], "single GPU"),
484 +
485 + # Balanced assignment
486 + ([100, 100, 100, 100
487 + ], 2, [0, 2, 1, 3], [2, 2], [200, 200], "balanced assignment"),
488 +
489 + # Unbalanced sizes - this one is trickier since the algorithm is greedy
490 + ([1000, 100, 200, 50], 2, [0, 2, 1, 3
491 + ], [1, 3], [1000, 350], "unbalanced sizes"),
492 + ],
493 +)
494 +def test_get_load_balance_assignment_cases(sizes, num_gpus,
495 + expected_shuffle_indices,
496 + expected_gpu_sample_counts,
497 + expected_grouped_sizes_per_gpu,
498 + test_description):
499 + """Test get_load_balance_assignment with various input cases."""
500 + result = get_load_balance_assignment(sizes, num_gpus=num_gpus)
501 + (shuffle_indices, gpu_sample_counts, grouped_sizes_per_gpu) = result
502 +
503 + # Common assertions for all cases
504 + assert len(shuffle_indices) == len(sizes)
505 + assert len(gpu_sample_counts) == num_gpus
506 + assert len(grouped_sizes_per_gpu) == num_gpus
507 + assert sum(gpu_sample_counts) == len(sizes)
508 +
509 + assert shuffle_indices == expected_shuffle_indices
510 +
511 + assert gpu_sample_counts == expected_gpu_sample_counts
512 + assert grouped_sizes_per_gpu == expected_grouped_sizes_per_gpu
513 +
514 +
515 +class SimpleMRopeVisionModel(torch.nn.Module):
516 + """A simple vision model for testing mrope functionality."""
517 +
518 + def __init__(self, spatial_merge_size: int = 2, out_hidden_size: int = 64):
519 + super().__init__()
520 + self.spatial_merge_size = spatial_merge_size
521 + self.out_hidden_size = out_hidden_size
522 + self.linear = torch.nn.Linear(768, out_hidden_size)
523 +
524 + def forward(self, pixel_values: torch.Tensor,
525 + grid_thw_list: list[list[int]]):
526 + """Simple forward pass that simulates spatial merging."""
527 + # Apply linear transformation
528 + embeddings = self.linear(pixel_values)
529 +
530 + # Simulate spatial merging by reducing the number of patches
531 + merge_factor = self.spatial_merge_size * self.spatial_merge_size
532 +
533 + # Group patches and merge spatially
534 + merged_embeddings = []
535 + start_idx = 0
536 +
537 + for grid_thw in grid_thw_list:
538 + num_patches = math.prod(grid_thw)
539 + end_idx = start_idx + num_patches
540 +
541 + # Get patches for this image
542 + image_patches = embeddings[start_idx:end_idx]
543 +
544 + # Simulate spatial merging by averaging groups of patches
545 + merged_patches = num_patches // merge_factor
546 + if merged_patches > 0:
547 + # Reshape and average to simulate merging
548 + reshaped = image_patches[:merged_patches * merge_factor].view(
549 + merged_patches, merge_factor, -1)
550 + merged = reshaped.mean(dim=1)
551 + merged_embeddings.append(merged)
552 +
553 + start_idx = end_idx
554 +
555 + if merged_embeddings:
556 + return torch.cat(merged_embeddings, dim=0)
557 + else:
558 + return torch.empty((0, self.out_hidden_size),
559 + device=pixel_values.device,
560 + dtype=pixel_values.dtype)
561 +
562 +
563 +@multi_gpu_test(num_gpus=2)
564 +@pytest.mark.parametrize(
565 + "batch_size",
566 + [
567 + 1, # Single image
568 + 3, # Small batch
569 + 5, # Odd batch size (for testing padding)
570 + ],
571 +)
572 +def test_run_dp_sharded_mrope_vision_model(batch_size: int):
573 + world_size = 2
574 + # Launch processes
575 + mp.spawn(
576 + run_dp_sharded_mrope_vision_model_vs_direct,
577 + args=(
578 + world_size,
579 + batch_size,
580 + get_open_port(),
581 + ),
582 + nprocs=world_size,
583 + )
584 +
585 +
586 +def run_dp_sharded_mrope_vision_model_vs_direct(local_rank: int,
587 + world_size: int,
588 + batch_size: int,
589 + master_port: int):
590 + """
591 + Test that run_dp_sharded_mrope_vision_model produces the same results as
592 + calling the model directly.
593 + """
594 + # Set random seed for reproducibility
595 + current_platform.seed_everything(0)
596 + device = f"{current_platform.device_name}:{local_rank}"
597 + current_platform.set_device(device)
598 + torch.set_default_device(device)
599 +
600 + update_environment_variables({
601 + 'RANK': str(local_rank),
602 + 'LOCAL_RANK': str(local_rank),
603 + 'WORLD_SIZE': str(world_size),
604 + 'MASTER_ADDR': 'localhost',
605 + 'MASTER_PORT': str(master_port),
606 + })
607 +
608 + # initialize distributed
609 + init_distributed_environment()
610 + initialize_model_parallel(tensor_model_parallel_size=world_size)
611 +
612 + # Create test data
613 + grid_thw_list = []
614 + pixel_values_list = []
615 +
616 + for i in range(batch_size):
617 + # Varying image sizes for better testing
618 + t, h, w = 1, 4 + i, 4 + i
619 + grid_thw_list.append([t, h, w])
620 +
621 + num_patches = t * h * w
622 + # Create random pixel values for this image
623 + image_pixels = torch.randn(num_patches, 768)
624 + pixel_values_list.append(image_pixels)
625 +
626 + # Concatenate all pixel values
627 + pixel_values = torch.cat(pixel_values_list, dim=0)
628 +
629 + # Create a simple mrope vision model
630 + vision_model = SimpleMRopeVisionModel()
631 +
632 + # Run the model directly on the full input (only on rank 0)
633 + if local_rank == 0:
634 + with torch.inference_mode():
635 + direct_output = vision_model(pixel_values, grid_thw_list)
636 +
637 + # Run the model through the sharded function
638 + with torch.inference_mode():
639 + sharded_output = run_dp_sharded_mrope_vision_model(
640 + vision_model, pixel_values, grid_thw_list)
641 + sharded_output = torch.cat(sharded_output, dim=0)
642 +
643 + # Check that the world size is setup correctly
644 + assert get_tensor_model_parallel_world_size() == world_size
645 +
646 + # Compare outputs (only on rank 0)
647 + if local_rank == 0:
648 + # Check that the outputs have the same shape
649 + assert direct_output.shape == sharded_output.shape
650 + # Check that the outputs are close (they should be identical)
651 + assert torch.allclose(direct_output,
652 + sharded_output,
653 + rtol=1e-5,
654 + atol=1e-5)
655 +
656 +
657 +@multi_gpu_test(num_gpus=2)
658 +def test_run_dp_sharded_mrope_vision_model_empty_input():
659 + world_size = 2
660 + mp.spawn(
661 + run_dp_sharded_mrope_vision_model_empty_input_worker,
662 + args=(world_size, get_open_port()),
663 + nprocs=world_size,
664 + )
665 +
666 +
667 +def run_dp_sharded_mrope_vision_model_empty_input_worker(
668 + local_rank: int, world_size: int, master_port: int):
669 + """Test run_dp_sharded_mrope_vision_model with empty input."""
670 + # Set up distributed environment
671 + device = f"{current_platform.device_name}:{local_rank}"
672 + current_platform.set_device(device)
673 + torch.set_default_device(device)
674 +
675 + update_environment_variables({
676 + 'RANK': str(local_rank),
677 + 'LOCAL_RANK': str(local_rank),
678 + 'WORLD_SIZE': str(world_size),
679 + 'MASTER_ADDR': 'localhost',
680 + 'MASTER_PORT': str(master_port),
681 + })
682 +
683 + init_distributed_environment()
684 + initialize_model_parallel(tensor_model_parallel_size=world_size)
685 +
686 + # Create empty inputs
687 + pixel_values = torch.empty((0, 768))
688 + grid_thw_list: list[list[int]] = []
689 +
690 + vision_model = SimpleMRopeVisionModel()
691 +
692 + # Should handle empty input gracefully
693 + with torch.inference_mode():
694 + output = run_dp_sharded_mrope_vision_model(vision_model, pixel_values,
695 + grid_thw_list)
696 +
697 + assert len(output) == 0
698 +
699 +
700 +@multi_gpu_test(num_gpus=4)
701 +def test_run_dp_sharded_mrope_vision_model_uneven_load():
702 + world_size = 4
703 + mp.spawn(
704 + run_dp_sharded_mrope_vision_model_uneven_load_worker,
705 + args=(world_size, get_open_port()),
706 + nprocs=world_size,
707 + )
708 +
709 +
710 +def run_dp_sharded_mrope_vision_model_uneven_load_worker(
711 + local_rank: int, world_size: int, master_port: int):
712 + """Test run_dp_sharded_mrope_vision_model with uneven load distribution."""
713 + # Set up distributed environment
714 + current_platform.seed_everything(123)
715 + device = f"{current_platform.device_name}:{local_rank}"
716 + current_platform.set_device(device)
717 + torch.set_default_device(device)
718 +
719 + update_environment_variables({
720 + 'RANK': str(local_rank),
721 + 'LOCAL_RANK': str(local_rank),
722 + 'WORLD_SIZE': str(world_size),
723 + 'MASTER_ADDR': 'localhost',
724 + 'MASTER_PORT': str(master_port),
725 + })
726 +
727 + init_distributed_environment()
728 + initialize_model_parallel(tensor_model_parallel_size=world_size)
729 +
730 + # Create images with very different sizes
731 + grid_thw_list = [
732 + [1, 2, 2], # Small: 4 patches
733 + [1, 8, 8], # Large: 64 patches
734 + [1, 3, 3], # Medium: 9 patches
735 + ]
736 +
737 + pixel_values_list = []
738 + for grid_thw in grid_thw_list:
739 + num_patches = math.prod(grid_thw)
740 + image_pixels = torch.randn(num_patches, 768)
741 + pixel_values_list.append(image_pixels)
742 +
743 + pixel_values = torch.cat(pixel_values_list, dim=0)
744 + vision_model = SimpleMRopeVisionModel()
745 +
746 + # Should handle uneven distribution without errors
747 + with torch.inference_mode():
748 + output_tuple = run_dp_sharded_mrope_vision_model(
749 + vision_model, pixel_values, grid_thw_list)
750 +
751 + # Verify output shape is reasonable
752 + merge_factor = vision_model.spatial_merge_size**2
753 + expected_output_patches = list(
754 + math.prod(grid_thw) // merge_factor for grid_thw in grid_thw_list)
755 +
756 + for i, output in enumerate(output_tuple):
757 + assert output.shape[0] == expected_output_patches[i]
758 + assert output.shape[1] == vision_model.out_hidden_size
759 +
760 +
761 +@pytest.mark.parametrize("spatial_merge_size", [2, 4])
762 +def test_simple_mrope_vision_model_spatial_merge(spatial_merge_size: int):
763 + """Test SimpleMRopeVisionModel with different spatial merge sizes."""
764 + device = current_platform.device_type
765 +
766 + grid_thw_list = [[1, 4, 4], [1, 6, 6]] # Two images
767 + pixel_values_list = []
768 +
769 + for grid_thw in grid_thw_list:
770 + num_patches = math.prod(grid_thw)
771 + image_pixels = torch.randn(num_patches, 768, device=device)
772 + pixel_values_list.append(image_pixels)
773 +
774 + pixel_values = torch.cat(pixel_values_list, dim=0)
775 + vision_model = SimpleMRopeVisionModel(
776 + spatial_merge_size=spatial_merge_size).to(device)
777 +
778 + with torch.inference_mode():
779 + output = vision_model(pixel_values, grid_thw_list)
780 +
781 + # Verify output dimensions based on spatial merging
782 + total_patches = sum(math.prod(grid_thw) for grid_thw in grid_thw_list)
783 + merge_factor = spatial_merge_size**2
784 + expected_output_patches = total_patches // merge_factor
785 +
786 + assert output.shape[0] == expected_output_patches
787 + assert output.shape[1] == vision_model.out_hidden_size
vllm/model_executor/layers/linear.py
+1 -1
@@ -437,7 +437,7 @@ class MergedReplicatedLinear(ReplicatedLinear):
437 shard_offset = sum(self.output_sizes[:loaded_shard_id])
438 shard_size = self.output_sizes[loaded_shard_id]
439
440 - param[shard_offset:shard_offset + shard_size] = loaded_weight
440 + param.data[shard_offset:shard_offset + shard_size] = loaded_weight
441
442
443 @CustomOp.register("column_parallel_linear")
vllm/model_executor/models/qwen2_5_vl.py
+93 -43
@@ -45,10 +45,14 @@ from vllm.logger import init_logger
45 from vllm.model_executor import SamplingMetadata
46 from vllm.model_executor.layers.activation import get_act_and_mul_fn
47 from vllm.model_executor.layers.layernorm import RMSNorm
48 +# yapf: disable
49 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
50 MergedColumnParallelLinear,
51 + MergedReplicatedLinear,
52 QKVParallelLinear,
53 + ReplicatedLinear,
54 RowParallelLinear)
55 +# yapf: enable
56 from vllm.model_executor.layers.quantization import QuantizationConfig
57 from vllm.model_executor.layers.quantization.gptq import GPTQConfig
58 from vllm.model_executor.layers.quantization.gptq_marlin import (
@@ -57,6 +61,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
61 from vllm.model_executor.models.module_mapping import MultiModelKeys
62 from vllm.multimodal import MULTIMODAL_REGISTRY
63 from vllm.multimodal.inputs import MultiModalFieldConfig
64 +from vllm.multimodal.utils import run_dp_sharded_mrope_vision_model
65 from vllm.platforms import _Backend
66 from vllm.sequence import IntermediateTensors
67 from vllm.transformers_utils.config import uses_mrope
@@ -170,19 +175,25 @@ class Qwen2_5_VisionMLP(nn.Module):
175 bias: bool = False,
176 act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,
177 quant_config: Optional[QuantizationConfig] = None,
173 - prefix: str = ""):
178 + prefix: str = "",
179 + use_data_parallel: bool = False):
180 super().__init__()
175 - self.gate_up_proj = MergedColumnParallelLinear(
181 + cls_gate_up_proj = (MergedReplicatedLinear if use_data_parallel else
182 + MergedColumnParallelLinear)
183 + self.gate_up_proj = cls_gate_up_proj(
184 input_size=in_features,
185 output_sizes=[hidden_features] * 2, # [gate_proj, up_proj]
186 bias=bias,
187 quant_config=quant_config,
188 prefix=f"{prefix}.gate_up_proj")
181 - self.down_proj = RowParallelLinear(hidden_features,
182 - in_features,
183 - bias=bias,
184 - quant_config=quant_config,
185 - prefix=f"{prefix}.down_proj")
189 +
190 + cls_down_proj = (ReplicatedLinear
191 + if use_data_parallel else RowParallelLinear)
192 + self.down_proj = cls_down_proj(hidden_features,
193 + in_features,
194 + bias=bias,
195 + quant_config=quant_config,
196 + prefix=f"{prefix}.down_proj")
197 self.act_fn = act_fn
198
199 def forward(self, x: torch.Tensor):
@@ -220,28 +231,42 @@ class Qwen2_5_VisionAttention(nn.Module):
231 projection_size: int,
232 quant_config: Optional[QuantizationConfig] = None,
233 prefix: str = "",
234 + use_data_parallel: bool = False,
235 ) -> None:
236 super().__init__()
237 # Per attention head and per partition values.
226 - self.tp_size = parallel_state.get_tensor_model_parallel_world_size()
238 + self.tp_size = (1 if use_data_parallel else
239 + parallel_state.get_tensor_model_parallel_world_size())
240 self.tp_rank = parallel_state.get_tensor_model_parallel_rank()
241 self.hidden_size_per_attention_head = dist_utils.divide(
242 projection_size, num_heads)
243 self.num_attention_heads_per_partition = dist_utils.divide(
244 num_heads, self.tp_size)
245
233 - self.qkv = QKVParallelLinear(
234 - hidden_size=embed_dim,
235 - head_size=self.hidden_size_per_attention_head,
236 - total_num_heads=num_heads,
237 - total_num_kv_heads=num_heads,
238 - bias=True,
239 - quant_config=quant_config,
240 - prefix=f"{prefix}.qkv")
241 - self.proj = RowParallelLinear(input_size=projection_size,
242 - output_size=embed_dim,
243 - quant_config=quant_config,
244 - prefix=f"{prefix}.proj")
246 + if use_data_parallel:
247 + self.qkv = ReplicatedLinear(embed_dim,
248 + self.hidden_size_per_attention_head *
249 + 3 * num_heads,
250 + bias=True,
251 + quant_config=quant_config,
252 + prefix=f"{prefix}.qkv")
253 +
254 + else:
255 + self.qkv = QKVParallelLinear(
256 + hidden_size=embed_dim,
257 + head_size=self.hidden_size_per_attention_head,
258 + total_num_heads=num_heads,
259 + total_num_kv_heads=num_heads,
260 + bias=True,
261 + quant_config=quant_config,
262 + prefix=f"{prefix}.qkv")
263 +
264 + cls_proj = (ReplicatedLinear
265 + if use_data_parallel else RowParallelLinear)
266 + self.proj = cls_proj(input_size=projection_size,
267 + output_size=embed_dim,
268 + quant_config=quant_config,
269 + prefix=f"{prefix}.proj")
270
271 # Detect attention implementation.
272 self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
@@ -302,8 +327,6 @@ class Qwen2_5_VisionAttention(nn.Module):
327 k = apply_rotary_pos_emb_vision(k, rotary_pos_emb)
328
329 if self.is_flash_attn_backend:
305 - # from vllm_flash_attn.flash_attn_interface import (
306 - # flash_attn_varlen_func)
330 if self.attn_backend == _Backend.ROCM_AITER_FA:
331 from aiter import flash_attn_varlen_func
332 else:
@@ -370,23 +393,27 @@ class Qwen2_5_VisionBlock(nn.Module):
393 norm_layer: Optional[Callable[[int], nn.Module]] = None,
394 quant_config: Optional[QuantizationConfig] = None,
395 prefix: str = "",
396 + use_data_parallel: bool = False,
397 ) -> None:
398 super().__init__()
399 if norm_layer is None:
400 norm_layer = partial(nn.LayerNorm, eps=1e-6)
401 self.norm1 = norm_layer(dim)
402 self.norm2 = norm_layer(dim)
379 - self.attn = Qwen2_5_VisionAttention(embed_dim=dim,
380 - num_heads=num_heads,
381 - projection_size=dim,
382 - quant_config=quant_config,
383 - prefix=f"{prefix}.attn")
403 + self.attn = Qwen2_5_VisionAttention(
404 + embed_dim=dim,
405 + num_heads=num_heads,
406 + projection_size=dim,
407 + quant_config=quant_config,
408 + prefix=f"{prefix}.attn",
409 + use_data_parallel=use_data_parallel)
410 self.mlp = Qwen2_5_VisionMLP(dim,
411 mlp_hidden_dim,
412 act_fn=act_fn,
413 bias=True,
414 quant_config=quant_config,
389 - prefix=f"{prefix}.mlp")
415 + prefix=f"{prefix}.mlp",
416 + use_data_parallel=use_data_parallel)
417
418 def forward(
419 self,
@@ -445,24 +472,30 @@ class Qwen2_5_VisionPatchMerger(nn.Module):
472 spatial_merge_size: int = 2,
473 quant_config: Optional[QuantizationConfig] = None,
474 prefix: str = "",
475 + use_data_parallel: bool = False,
476 ) -> None:
477 super().__init__()
478 self.hidden_size = context_dim * (spatial_merge_size**2)
479 if norm_layer is None:
480 norm_layer = partial(nn.LayerNorm, eps=1e-6)
481 self.ln_q = norm_layer(context_dim)
482 +
483 + cls_fc1 = (ReplicatedLinear
484 + if use_data_parallel else ColumnParallelLinear)
485 + cls_fc2 = (ReplicatedLinear
486 + if use_data_parallel else RowParallelLinear)
487 self.mlp = nn.ModuleList([
455 - ColumnParallelLinear(self.hidden_size,
456 - self.hidden_size,
457 - bias=True,
458 - quant_config=quant_config,
459 - prefix=f"{prefix}.mlp.0"),
488 + cls_fc1(self.hidden_size,
489 + self.hidden_size,
490 + bias=True,
491 + quant_config=quant_config,
492 + prefix=f"{prefix}.mlp.0"),
493 nn.GELU(),
461 - RowParallelLinear(self.hidden_size,
462 - d_model,
463 - bias=True,
464 - quant_config=quant_config,
465 - prefix=f"{prefix}.mlp.2"),
494 + cls_fc2(self.hidden_size,
495 + d_model,
496 + bias=True,
497 + quant_config=quant_config,
498 + prefix=f"{prefix}.mlp.2"),
499 ])
500
501 def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -514,6 +547,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
547 norm_eps: float = 1e-6,
548 quant_config: Optional[QuantizationConfig] = None,
549 prefix: str = "",
550 + use_data_parallel: bool = False,
551 ) -> None:
552 super().__init__()
553
@@ -523,6 +557,8 @@ class Qwen2_5_VisionTransformer(nn.Module):
557 depth = vision_config.depth
558 self.hidden_size = vision_config.hidden_size
559 self.num_heads = vision_config.num_heads
560 + self.use_data_parallel = use_data_parallel
561 + self.out_hidden_size = vision_config.out_hidden_size
562
563 # args for get_window_index_thw
564 self.window_size = vision_config.window_size
@@ -550,7 +586,8 @@ class Qwen2_5_VisionTransformer(nn.Module):
586 vision_config.hidden_act),
587 norm_layer=norm_layer,
588 quant_config=quant_config,
553 - prefix=f"{prefix}.blocks.{layer_idx}")
589 + prefix=f"{prefix}.blocks.{layer_idx}",
590 + use_data_parallel=use_data_parallel)
591 for layer_idx in range(depth)
592 ])
593 self.merger = Qwen2_5_VisionPatchMerger(
@@ -560,6 +597,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
597 spatial_merge_size=self.spatial_merge_size,
598 quant_config=quant_config,
599 prefix=f"{prefix}.merger",
600 + use_data_parallel=use_data_parallel,
601 )
602 self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
603
@@ -767,7 +805,6 @@ class Qwen2_5_VisionTransformer(nn.Module):
805 if weight_name not in name:
806 continue
807 name = name.replace(weight_name, param_name)
770 -
808 param = params_dict[name]
809 weight_loader = param.weight_loader
810 weight_loader(param, loaded_weight, shard_id)
@@ -840,6 +877,8 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,
877 config: Qwen2_5_VLConfig = vllm_config.model_config.hf_config
878 multimodal_config = vllm_config.model_config.multimodal_config
879
880 + self.use_data_parallel = (vllm_config.parallel_config.
881 + enable_multimodal_encoder_data_parallel)
882 self.config = config
883 self.multimodal_config = multimodal_config
884
@@ -851,6 +890,7 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,
890 quant_config=self._maybe_ignore_quant_config(
891 self.quant_config),
892 prefix=maybe_prefix(prefix, "visual"),
893 + use_data_parallel=self.use_data_parallel,
894 )
895 else:
896 self.visual = None
@@ -973,7 +1013,13 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,
1013 image_embeds = image_input["image_embeds"].type(self.visual.dtype)
1014 else:
1015 pixel_values = image_input["pixel_values"]
976 - image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)
1016 +
1017 + if self.use_data_parallel:
1018 + return run_dp_sharded_mrope_vision_model(
1019 + self.visual, pixel_values, grid_thw_list)
1020 + else:
1021 + image_embeds = self.visual(pixel_values,
1022 + grid_thw=grid_thw_list)
1023
1024 # Split concatenated embeddings for each image item.
1025 # Using prod on grid_thw_list instead of grid_thw.prod avoids CUDA sync
@@ -995,8 +1041,12 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,
1041 video_embeds = video_input["video_embeds"].type(self.visual.dtype)
1042 else:
1043 pixel_values_videos = video_input["pixel_values_videos"]
998 - video_embeds = self.visual(pixel_values_videos,
999 - grid_thw=grid_thw_list)
1044 + if self.use_data_parallel:
1045 + return run_dp_sharded_mrope_vision_model(
1046 + self.visual, pixel_values_videos, grid_thw_list)
1047 + else:
1048 + video_embeds = self.visual(pixel_values_videos,
1049 + grid_thw=grid_thw_list)
1050
1051 # Split concatenated embeddings for each video item.
1052 merge_size = self.visual.spatial_merge_size
vllm/model_executor/models/qwen2_vl.py
-2
@@ -329,8 +329,6 @@ class Qwen2VisionAttention(nn.Module):
329 k = apply_rotary_pos_emb_vision(k, rotary_pos_emb)
330
331 if self.is_flash_attn_backend:
332 - # from vllm_flash_attn.flash_attn_interface import (
333 - # flash_attn_varlen_func)
332 if self.attn_backend == _Backend.ROCM_AITER_FA:
333 from aiter import flash_attn_varlen_func
334 else:
vllm/multimodal/utils.py
+215
@@ -3,6 +3,8 @@
3
4 import asyncio
5 import atexit
6 +import itertools
7 +import math
8 from collections.abc import Iterable
9 from concurrent.futures import ThreadPoolExecutor
10 from itertools import groupby
@@ -465,6 +467,219 @@ def run_dp_sharded_vision_model(image_input: torch.Tensor,
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 +) -> tuple[torch.Tensor, ...]:
544 + """Run a vision model with data parallelism (DP) sharding.
545 + The function will shard the input image tensor on the
546 + first dimension and run the vision model.
547 + This function is used to run the vision model with mrope.
548 +
549 + Args:
550 + vision_model (torch.nn.Module): Vision model.
551 + pixel_values (torch.Tensor): Image/Video input tensor.
552 + grid_thw_list: List of grid dimensions for each image
553 + Returns:
554 + torch.Tensor: Output image embeddings
555 +
556 + Example:
557 + ```
558 + vision_model.out_hidden_size = 64
559 + vision_model.spatial_merge_size = 2
560 + pixel_values.shape = (1350, channel)
561 + grid_thw_list = [[1, 10, 100], [1, 10, 10], [1, 10, 20], [1, 50]]
562 + tp_size=2
563 + ```
564 +
565 + """
566 + tp_size = get_tensor_model_parallel_world_size()
567 +
568 + # GPU_0 tp_rank_local = 0
569 + # GPU_1 tp_rank_local = 1
570 + tp_rank_local = get_tensor_model_parallel_rank()
571 +
572 + # patches_per_image = [1000, 100, 200, 50]
573 + patches_per_image = [math.prod(grid_thw) for grid_thw in grid_thw_list]
574 + # patches_per_image = [0, 1000, 1100, 1300, 1350]
575 + cum_patches_per_image = [0, *itertools.accumulate(patches_per_image)]
576 +
577 + # Get load balancing assignment with all metadata
578 + # image_to_tp_rank = [0, 2, 1, 3]
579 + # gpu_sample_counts = [1, 3]
580 + # grouped_pixel_values_len = [1000, 350]
581 + (image_to_tp_rank, gpu_sample_counts,
582 + grouped_pixel_values_len) = get_load_balance_assignment(
583 + patches_per_image, tp_size)
584 +
585 + # cu_gpu_sample_counts = [0, 1, 4]
586 + cum_gpu_sample_counts = [0, *itertools.accumulate(gpu_sample_counts)]
587 +
588 + # GPU_0 image_idxs_local = [0]
589 + # GPU_1 image_idxs_local = [2, 1, 3]
590 + image_idxs_local = image_to_tp_rank[cum_gpu_sample_counts[tp_rank_local]:
591 + cum_gpu_sample_counts[tp_rank_local +
592 + 1]]
593 +
594 + # Get the pixel values for the local images based on the image_idxs_local
595 + if len(image_idxs_local) > 0:
596 + pixel_values_local = torch.cat([
597 + pixel_values[cum_patches_per_image[i]:cum_patches_per_image[i + 1]]
598 + for i in image_idxs_local
599 + ])
600 + else:
601 + # Handle case where this rank has no images
602 + pixel_values_local = torch.empty((0, pixel_values.shape[1]),
603 + device=pixel_values.device,
604 + dtype=pixel_values.dtype)
605 + # embed_dim_reduction_factor = 2 * 2
606 + embed_dim_reduction_factor = (vision_model.spatial_merge_size *
607 + vision_model.spatial_merge_size)
608 +
609 + # Find the max length across all ranks
610 + # The output embedding of every DP rank has to be
611 + # padded to this length for tensor_model_parallel_all_gather
612 + # to work
613 + max_len_per_rank = max(
614 + grouped_pixel_values_len) // embed_dim_reduction_factor
615 + local_grid_thw_list = [grid_thw_list[i] for i in image_idxs_local]
616 +
617 + # Run the vision model on the local pixel_values_local
618 + if pixel_values_local.shape[0] > 0:
619 + image_embeds_local = vision_model(pixel_values_local,
620 + local_grid_thw_list)
621 + else:
622 + # Handle empty case
623 + image_embeds_local = torch.empty((0, vision_model.out_hidden_size),
624 + device=pixel_values.device,
625 + dtype=pixel_values.dtype)
626 +
627 + # Pad the output based on max_len_per_rank
628 + # for tensor_model_parallel_all_gather to work
629 + current_len = image_embeds_local.shape[0]
630 + if current_len < max_len_per_rank:
631 + padding_size = max_len_per_rank - current_len
632 + padding = torch.empty((padding_size, image_embeds_local.shape[1]),
633 + dtype=image_embeds_local.dtype,
634 + device=image_embeds_local.device)
635 + image_embeds_local_padded = torch.cat([image_embeds_local, padding],
636 + dim=0)
637 + else:
638 + image_embeds_local_padded = image_embeds_local
639 +
640 + # Do all_gather to collect embeddings from all ranks
641 + gathered_embeds = tensor_model_parallel_all_gather(
642 + image_embeds_local_padded, dim=0)
643 +
644 + # Remove padding and reconstruct per-rank embeddings
645 + rank_embeddings = list[torch.Tensor]()
646 + for rank in range(tp_size):
647 + start_idx = rank * max_len_per_rank
648 + end_idx = start_idx + (grouped_pixel_values_len[rank] //
649 + embed_dim_reduction_factor)
650 + rank_embeddings.append(gathered_embeds[start_idx:end_idx])
651 +
652 + patches_per_output_image = [(patch_size // embed_dim_reduction_factor)
653 + for patch_size in patches_per_image]
654 +
655 + # Reconstruct embeddings in the original order
656 + original_order_embeddings = [None] * len(grid_thw_list)
657 + current_idx = 0
658 + for rank in range(tp_size):
659 + count = gpu_sample_counts[rank]
660 + if count > 0:
661 + # Get images assigned to this rank in shuffled order
662 + # GPU_0 = image_idxs_local [0]
663 + # GPU_1 = image_idxs_local [2, 1, 3]
664 + rank_images = image_to_tp_rank[current_idx:current_idx + count]
665 +
666 + rank_embed = rank_embeddings[rank]
667 + # Split rank embeddings back to individual images
668 + embed_start = 0
669 + for img_idx in rank_images:
670 + img_patches = patches_per_output_image[img_idx]
671 + original_order_embeddings[img_idx] = rank_embed[
672 + embed_start:embed_start + img_patches]
673 + embed_start += img_patches
674 + current_idx += count
675 +
676 + out_embeddings = tuple(embed for embed in original_order_embeddings
677 + if embed is not None)
678 + assert len(out_embeddings) == len(
679 + original_order_embeddings), "Found unassigned embeddings"
680 + return out_embeddings
681 +
682 +
683 def fetch_audio(
684 audio_url: str,
685 audio_io_kwargs: Optional[dict[str, Any]] = None,