c98be0a232 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -66,6 +66,7 @@ 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

70 from vllm.platforms import _Backend, current_platform

71 from vllm.sequence import IntermediateTensors

72 from vllm.transformers_utils.config import uses_mrope

@@ -217,17 +218,20 @@ class Qwen2VisionMLP(nn.Module):

218 act_layer: type[nn.Module] = QuickGELU,

219 quant_config: Optional[QuantizationConfig] = None,

220 prefix: str = "",

221 + use_data_parallel: bool = False,

222 ):

223 super().__init__()

224 self.fc1 = ColumnParallelLinear(in_features,

225 hidden_features,

226 quant_config=quant_config,

225 - prefix=f"{prefix}.fc1")

227 + prefix=f"{prefix}.fc1",

228 + disable_tp=use_data_parallel)

229 self.act = act_layer()

230 self.fc2 = RowParallelLinear(hidden_features,

231 in_features,

232 quant_config=quant_config,

230 - prefix=f"{prefix}.fc2")

233 + prefix=f"{prefix}.fc2",

234 + disable_tp=use_data_parallel)

235

236 def forward(self, x: torch.Tensor) -> torch.Tensor:

237 x_parallel, _ = self.fc1(x)

@@ -293,25 +297,28 @@ class Qwen2VisionAttention(nn.Module):

297 projection_size: int,

298 quant_config: Optional[QuantizationConfig] = None,

299 prefix: str = "",

300 + use_data_parallel: bool = False,

301 ) -> None:

302 super().__init__()

303 # Per attention head and per partition values.

299 - world_size = parallel_state.get_tensor_model_parallel_world_size()

300 - self.tp_size = world_size

304 + self.tp_size = (1 if use_data_parallel else

305 + parallel_state.get_tensor_model_parallel_world_size())

306 self.tp_rank = parallel_state.get_tensor_model_parallel_rank()

307 self.hidden_size_per_attention_head = dist_utils.divide(

308 projection_size, num_heads)

309 self.num_attention_heads_per_partition = dist_utils.divide(

305 - num_heads, world_size)

310 + num_heads, self.tp_size)

311

312 self.qkv = ColumnParallelLinear(input_size=embed_dim,

313 output_size=3 * projection_size,

314 quant_config=quant_config,

310 - prefix=f"{prefix}.qkv")

315 + prefix=f"{prefix}.qkv",

316 + disable_tp=use_data_parallel)

317 self.proj = RowParallelLinear(input_size=projection_size,

318 output_size=embed_dim,

319 quant_config=quant_config,

314 - prefix=f"{prefix}.proj")

320 + prefix=f"{prefix}.proj",

321 + disable_tp=use_data_parallel)

322

323 # Detect attention implementation.

324 self.attn_backend = get_vit_attn_backend(

@@ -453,6 +460,7 @@ class Qwen2VisionBlock(nn.Module):

460 norm_layer: Optional[Callable[[int], nn.Module]] = None,

461 quant_config: Optional[QuantizationConfig] = None,

462 prefix: str = "",

463 + use_data_parallel: bool = False,

464 ) -> None:

465 super().__init__()

466 if norm_layer is None:

@@ -465,12 +473,14 @@ class Qwen2VisionBlock(nn.Module):

473 num_heads=num_heads,

474 projection_size=dim,

475 quant_config=quant_config,

468 - prefix=f"{prefix}.attn")

476 + prefix=f"{prefix}.attn",

477 + use_data_parallel=use_data_parallel)

478 self.mlp = Qwen2VisionMLP(dim,

479 mlp_hidden_dim,

480 act_layer=act_layer,

481 quant_config=quant_config,

473 - prefix=f"{prefix}.mlp")

482 + prefix=f"{prefix}.mlp",

483 + use_data_parallel=use_data_parallel)

484

485 def forward(

486 self,

@@ -531,6 +541,7 @@ class Qwen2VisionPatchMerger(nn.Module):

541 spatial_merge_size: int = 2,

542 quant_config: Optional[QuantizationConfig] = None,

543 prefix: str = "",

544 + use_data_parallel: bool = False,

545 ) -> None:

546 super().__init__()

547 self.hidden_size = context_dim * (spatial_merge_size**2)

@@ -542,13 +553,15 @@ class Qwen2VisionPatchMerger(nn.Module):

553 self.hidden_size,

554 bias=True,

555 quant_config=quant_config,

545 - prefix=f"{prefix}.mlp.0"),

556 + prefix=f"{prefix}.mlp.0",

557 + disable_tp=use_data_parallel),

558 nn.GELU(),

559 RowParallelLinear(self.hidden_size,

560 d_model,

561 bias=True,

562 quant_config=quant_config,

551 - prefix=f"{prefix}.mlp.2"),

563 + prefix=f"{prefix}.mlp.2",

564 + disable_tp=use_data_parallel),

565 ])

566

567 def forward(self, x: torch.Tensor) -> torch.Tensor:

@@ -600,6 +613,7 @@ class Qwen2VisionTransformer(nn.Module):

613 norm_eps: float = 1e-6,

614 quant_config: Optional[QuantizationConfig] = None,

615 prefix: str = "",

616 + use_data_parallel: bool = False,

617 ) -> None:

618 super().__init__()

619

@@ -613,6 +627,9 @@ class Qwen2VisionTransformer(nn.Module):

627 num_heads = vision_config.num_heads

628 mlp_ratio = vision_config.mlp_ratio

629

630 + self.use_data_parallel = use_data_parallel

631 + self.out_hidden_size = vision_config.hidden_size

632 +

633 self.spatial_merge_size = spatial_merge_size

634 self.num_heads = num_heads

635 self.embed_dim = embed_dim

@@ -634,7 +651,8 @@ class Qwen2VisionTransformer(nn.Module):

651 mlp_ratio=mlp_ratio,

652 norm_layer=norm_layer,

653 quant_config=quant_config,

637 - prefix=f"{prefix}.blocks.{layer_idx}")

654 + prefix=f"{prefix}.blocks.{layer_idx}",

655 + use_data_parallel=use_data_parallel)

656 for layer_idx in range(depth)

657 ])

658 self.merger = Qwen2VisionPatchMerger(

@@ -643,6 +661,7 @@ class Qwen2VisionTransformer(nn.Module):

661 norm_layer=norm_layer,

662 quant_config=quant_config,

663 prefix=f"{prefix}.merger",

664 + use_data_parallel=use_data_parallel,

665 )

666 self.attn_backend = get_vit_attn_backend(

667 head_size=head_dim, dtype=torch.get_default_dtype())

@@ -659,8 +678,9 @@ class Qwen2VisionTransformer(nn.Module):

678 def device(self) -> torch.device:

679 return self.patch_embed.proj.weight.device

680

662 - def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:

681 + def rot_pos_emb(self, grid_thw: list[list[int]]) -> torch.Tensor:

682 pos_ids = []

683 + max_grid_size = 0

684 for t, h, w in grid_thw:

685 hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)

686 wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)

@@ -678,8 +698,8 @@ class Qwen2VisionTransformer(nn.Module):

698 ).permute(0, 2, 1, 3).flatten()

699 pos_ids.append(

700 torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))

701 + max_grid_size = max(max_grid_size, h, w)

702 pos_ids = torch.cat(pos_ids, dim=0)

682 - max_grid_size = grid_thw[:, 1:].max()

703 rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)

704 rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)

705 return rotary_pos_emb

@@ -698,7 +718,7 @@ class Qwen2VisionTransformer(nn.Module):

718 def forward(

719 self,

720 x: torch.Tensor,

701 - grid_thw: torch.Tensor,

721 + grid_thw: list[list[int]],

722 ) -> torch.Tensor:

723 # patchify

724 x = x.to(device=self.device, dtype=self.dtype)

@@ -708,8 +728,9 @@ class Qwen2VisionTransformer(nn.Module):

728 rotary_pos_emb = self.rot_pos_emb(grid_thw)

729

730 # compute cu_seqlens

711 - cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2],

712 - grid_thw[:, 0]).cumsum(

731 + grid_thw_ = torch.tensor(grid_thw)

732 + cu_seqlens = torch.repeat_interleave(grid_thw_[:, 1] * grid_thw_[:, 2],

733 + grid_thw_[:, 0]).cumsum(

734 dim=0, dtype=torch.int32)

735 cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0)

736

@@ -1112,6 +1133,8 @@ class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1133 "model.": "language_model.model.",

1134 })

1135

1136 + supports_encoder_tp_data = True

1137 +

1138 def get_mrope_input_positions(

1139 self,

1140 input_tokens: list[int],

@@ -1239,6 +1262,7 @@ class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1262 quant_config = vllm_config.quant_config

1263 multimodal_config = vllm_config.model_config.multimodal_config

1264

1265 + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

1266 self.config = config

1267 self.multimodal_config = multimodal_config

1268

@@ -1249,6 +1273,7 @@ class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1273 norm_eps=getattr(config, "rms_norm_eps", 1e-6),

1274 quant_config=self._maybe_ignore_quant_config(quant_config),

1275 prefix=maybe_prefix(prefix, "visual"),

1276 + use_data_parallel=self.use_data_parallel,

1277 )

1278 else:

1279 self.visual = None

@@ -1357,7 +1382,15 @@ class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1382 image_embeds = image_input["image_embeds"]

1383 else:

1384 pixel_values = image_input["pixel_values"]

1360 - image_embeds = self.visual(pixel_values, grid_thw=grid_thw)

1385 +

1386 + if self.use_data_parallel:

1387 + return run_dp_sharded_mrope_vision_model(self.visual,

1388 + pixel_values,

1389 + grid_thw_list,

1390 + rope_type="rope_3d")

1391 + else:

1392 + image_embeds = self.visual(pixel_values,

1393 + grid_thw=grid_thw_list)

1394

1395 # Split concatenated embeddings for each image item.

1396 merge_size = self.visual.spatial_merge_size

@@ -1377,7 +1410,14 @@ class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1410 video_embeds = video_input["video_embeds"]

1411 else:

1412 pixel_values_videos = video_input["pixel_values_videos"]

1380 - video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)

1413 + if self.use_data_parallel:

1414 + return run_dp_sharded_mrope_vision_model(self.visual,

1415 + pixel_values_videos,

1416 + grid_thw_list,

1417 + rope_type="rope_3d")

1418 + else:

1419 + video_embeds = self.visual(pixel_values_videos,

1420 + grid_thw=grid_thw_list)

1421

1422 # Split concatenated embeddings for each video item.

1423 merge_size = self.visual.spatial_merge_size

[Model] Enable DP for ViT in Qwen2-VL (#25445)