@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2f0bab3f26

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

[Model] Support dp on ViT on GLM-4.5V (#23168)

Signed-off-by: David Chen <530634352@qq.com>

WeiQing Chen committed Sep 2, 2025 at 18:48 UTC 2f0bab3f26678ebaacdbdd61e0bec581fa7e30b9

2 files changed +145 -59

docs/configuration/optimization.md

@@ -174,6 +174,7 @@ Regardless, you need to set `mm_encoder_tp_mode="data"` in engine arguments to u

174

175 Known supported models:

176

177 +- GLM-4.5V GLM-4.1V (<gh-pr:23168>)

178 - Kimi-VL (<gh-pr:23817>)

179 - Llama4 (<gh-pr:18368>)

180 - MiniCPM-V-2.5 or above (<gh-pr:23327>, <gh-pr:23948>)

vllm/model_executor/models/glm4_1v.py

+144 -59

@@ -45,15 +45,20 @@ from transformers.models.glm4v.video_processing_glm4v import (

45 from transformers.video_utils import VideoMetadata

47 from vllm.config import VllmConfig

48 -from vllm.distributed import parallel_state

48 +from vllm.distributed import (get_tensor_model_parallel_world_size,

49 + parallel_state)

50 from vllm.distributed import utils as dist_utils

51 from vllm.logger import init_logger

52 from vllm.model_executor import SamplingMetadata

53 from vllm.model_executor.layers.layernorm import RMSNorm

54 +# yapf: disable

55 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

56 MergedColumnParallelLinear,

57 + MergedReplicatedLinear,

58 QKVParallelLinear,

59 + ReplicatedLinear,

60 RowParallelLinear)

61 +# yapf: enable

62 from vllm.model_executor.layers.quantization import QuantizationConfig

63 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

64 from vllm.model_executor.models.module_mapping import MultiModelKeys

@@ -66,6 +71,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,

71 BaseProcessingInfo, PromptReplacement,

72 PromptUpdate, PromptUpdateDetails)

73 from vllm.multimodal.profiling import BaseDummyInputsBuilder

74 +from vllm.multimodal.utils import run_dp_sharded_mrope_vision_model

75 from vllm.platforms import _Backend

76 from vllm.sequence import IntermediateTensors

77 from vllm.transformers_utils.config import uses_mrope

@@ -153,7 +159,7 @@ class Glm4vVideoEmbeddingInputs(TensorSchema):

159

160 Glm4vVideoInputs = Union[Glm4vVideoPixelInputs, Glm4vVideoEmbeddingInputs]

161

156 -# === Vision Encoder === #

162 +# ==== Vision Encoder ==== #

163

164

165 class Glm4vVisionMLP(nn.Module):

@@ -165,19 +171,23 @@ class Glm4vVisionMLP(nn.Module):

171 bias: bool = False,

172 quant_config: Optional[QuantizationConfig] = None,

173 prefix: str = "",

174 + use_data_parallel: bool = False,

175 ):

176 super().__init__()

170 - self.gate_up_proj = MergedColumnParallelLinear(

171 - input_size=in_features,

172 - output_sizes=[hidden_features] * 2,

173 - bias=bias,

174 - quant_config=quant_config,

175 - prefix=f"{prefix}.gate_up_proj")

176 - self.down_proj = RowParallelLinear(hidden_features,

177 - in_features,

178 - bias=bias,

179 - quant_config=quant_config,

180 - prefix=f"{prefix}.down_proj")

177 + cls_gate_up = (MergedReplicatedLinear

178 + if use_data_parallel else MergedColumnParallelLinear)

179 + self.gate_up_proj = cls_gate_up(input_size=in_features,

180 + output_sizes=[hidden_features] * 2,

181 + bias=bias,

182 + quant_config=quant_config,

183 + prefix=f"{prefix}.gate_up_proj")

184 + cls_down = (ReplicatedLinear

185 + if use_data_parallel else RowParallelLinear)

186 + self.down_proj = cls_down(hidden_features,

187 + in_features,

188 + bias=bias,

189 + quant_config=quant_config,

190 + prefix=f"{prefix}.down_proj")

191 self.act_fn = SiluAndMul()

192

193 def forward(self, x: torch.Tensor):

@@ -218,33 +228,54 @@ class Glm4vVisionAttention(nn.Module):

228 projection_size: int,

229 quant_config: Optional[QuantizationConfig] = None,

230 prefix: str = "",

231 + use_data_parallel: bool = False,

232 ) -> None:

233 super().__init__()

234 # Per attention head and per partition values.

224 - self.tp_size = parallel_state.get_tensor_model_parallel_world_size()

235 + self.tp_size = (1 if use_data_parallel else

236 + get_tensor_model_parallel_world_size())

237 self.tp_rank = parallel_state.get_tensor_model_parallel_rank()

238 self.hidden_size_per_attention_head = dist_utils.divide(

239 projection_size, num_heads)

240 self.num_attention_heads_per_partition = dist_utils.divide(

241 num_heads, self.tp_size)

242

231 - self.qkv = QKVParallelLinear(

232 - hidden_size=embed_dim,

233 - head_size=self.hidden_size_per_attention_head,

234 - total_num_heads=num_heads,

235 - total_num_kv_heads=num_heads,

236 - bias=False,

237 - quant_config=quant_config,

238 - # Change qkv prefix to align with GLM-4.5V-FP8 quantization config

239 - prefix=f"{prefix}.qkv_proj" if quant_config else f"{prefix}.qkv",

240 - )

241 - self.proj = RowParallelLinear(

242 - input_size=projection_size,

243 - output_size=embed_dim,

244 - quant_config=quant_config,

245 - prefix=f"{prefix}.proj",

246 - bias=False,

247 - )

243 + if use_data_parallel:

244 + self.qkv = ReplicatedLinear(

245 + input_size=embed_dim,

246 + output_size=3 * projection_size,

247 + bias=False,

248 + quant_config=quant_config,

249 + # Change qkv prefix to align with GLM-4.5V-FP8 quantization cfg

250 + prefix=f"{prefix}.qkv_proj"

251 + if quant_config else f"{prefix}.qkv",

252 + )

253 + self.proj = ReplicatedLinear(

254 + input_size=projection_size,

255 + output_size=embed_dim,

256 + quant_config=quant_config,

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

258 + bias=False,

259 + )

260 + else:

261 + self.qkv = QKVParallelLinear(

262 + hidden_size=embed_dim,

263 + head_size=self.hidden_size_per_attention_head,

264 + total_num_heads=num_heads,

265 + total_num_kv_heads=num_heads,

266 + bias=False,

267 + quant_config=quant_config,

268 + # Change qkv prefix to align with GLM-4.5V-FP8 quantization cfg

269 + prefix=f"{prefix}.qkv_proj"

270 + if quant_config else f"{prefix}.qkv",

271 + )

272 + self.proj = RowParallelLinear(

273 + input_size=projection_size,

274 + output_size=embed_dim,

275 + quant_config=quant_config,

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

277 + bias=False,

278 + )

279

280 # Detect attention implementation.

281 self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)

@@ -375,6 +406,7 @@ class Glm4vVisionBlock(nn.Module):

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

407 quant_config: Optional[QuantizationConfig] = None,

408 prefix: str = "",

409 + use_data_parallel: bool = False,

410 ) -> None:

411 super().__init__()

412 if norm_layer is None:

@@ -387,6 +419,7 @@ class Glm4vVisionBlock(nn.Module):

419 projection_size=dim,

420 quant_config=quant_config,

421 prefix=f"{prefix}.attn",

422 + use_data_parallel=use_data_parallel,

423 )

424 self.mlp = Glm4vVisionMLP(

425 dim,

@@ -394,6 +427,7 @@ class Glm4vVisionBlock(nn.Module):

427 bias=False,

428 quant_config=quant_config,

429 prefix=f"{prefix}.mlp",

430 + use_data_parallel=use_data_parallel,

431 )

432

433 def forward(

@@ -456,24 +490,40 @@ class Glm4vPatchMerger(nn.Module):

490 quant_config: Optional[QuantizationConfig] = None,

491 bias: bool = False,

492 prefix: str = "",

493 + use_data_parallel: bool = False,

494 ) -> None:

495 super().__init__()

496 self.hidden_size = d_model

462 - self.proj = ColumnParallelLinear(self.hidden_size,

463 - self.hidden_size,

464 - bias=bias,

465 - gather_output=True,

466 - quant_config=quant_config,

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

497 + if use_data_parallel:

498 + self.proj = ReplicatedLinear(

499 + input_size=self.hidden_size,

500 + output_size=self.hidden_size,

501 + bias=bias,

502 + quant_config=quant_config,

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

504 + )

505 + else:

506 + self.proj = ColumnParallelLinear(

507 + self.hidden_size,

508 + self.hidden_size,

509 + bias=bias,

510 + gather_output=True,

511 + quant_config=quant_config,

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

513 + )

514 self.post_projection_norm = nn.LayerNorm(self.hidden_size)

469 - self.gate_up_proj = MergedColumnParallelLinear(

515 + cls_gate_up = (MergedReplicatedLinear

516 + if use_data_parallel else MergedColumnParallelLinear)

517 + self.gate_up_proj = cls_gate_up(

518 input_size=self.hidden_size,

519 output_sizes=[context_dim] * 2,

520 bias=bias,

521 quant_config=quant_config,

522 prefix=f"{prefix}.gate_up_proj",

523 )

476 - self.down_proj = RowParallelLinear(

524 + cls_down = (ReplicatedLinear

525 + if use_data_parallel else RowParallelLinear)

526 + self.down_proj = cls_down(

527 context_dim,

528 self.hidden_size,

529 bias=bias,

@@ -548,14 +598,33 @@ class Glm4vVisionEmbeddings(nn.Module):

598 dtype=torch.float32))

599

600 # Calculate target dimensions for each patch

551 - target_h = torch.cat([

552 - image_shapes[i, 1].repeat(lengths[i])

553 - for i in range(len(lengths))

554 - ]).to(device=device, dtype=torch.float32)

555 - target_w = torch.cat([

556 - image_shapes[i, 2].repeat(lengths[i])

557 - for i in range(len(lengths))

558 - ]).to(device=device, dtype=torch.float32)

601 + # Add bounds checking for data parallel mode

602 + if len(lengths) > image_shapes.shape[0]:

603 + # In data parallel mode, some GPUs might not have all

604 + # image shapes

605 + # Use available image shapes, cycling if necessary

606 + target_h_list = []

607 + target_w_list = []

608 + for i in range(len(lengths)):

609 + # Cycle through available shapes

610 + shape_idx = i % image_shapes.shape[0]

611 + target_h_list.append(image_shapes[shape_idx,

612 + 1].repeat(lengths[i]))

613 + target_w_list.append(image_shapes[shape_idx,

614 + 2].repeat(lengths[i]))

615 + target_h = torch.cat(target_h_list).to(device=device,

616 + dtype=torch.float32)

617 + target_w = torch.cat(target_w_list).to(device=device,

618 + dtype=torch.float32)

619 + else:

620 + target_h = torch.cat([

621 + image_shapes[i, 1].repeat(lengths[i])

622 + for i in range(len(lengths))

623 + ]).to(device=device, dtype=torch.float32)

624 + target_w = torch.cat([

625 + image_shapes[i, 2].repeat(lengths[i])

626 + for i in range(len(lengths))

627 + ]).to(device=device, dtype=torch.float32)

628

629 # Normalize coordinates to [-1, 1] range for grid_sample

630 h_coords = h_coords.to(device=device, dtype=torch.float32)

@@ -629,6 +698,7 @@ class Glm4vVisionTransformer(nn.Module):

698 norm_eps: float = 1e-6,

699 quant_config: Optional[QuantizationConfig] = None,

700 prefix: str = "",

701 + use_data_parallel: bool = False,

702 ) -> None:

703 super().__init__()

704

@@ -638,6 +708,7 @@ class Glm4vVisionTransformer(nn.Module):

708 depth = vision_config.depth

709 self.hidden_size = vision_config.hidden_size

710 self.num_heads = vision_config.num_heads

711 + self.use_data_parallel = use_data_parallel

712

713 self.patch_size = vision_config.patch_size

714 self.spatial_merge_size = vision_config.spatial_merge_size

@@ -661,6 +732,7 @@ class Glm4vVisionTransformer(nn.Module):

732 norm_layer=norm_layer,

733 quant_config=quant_config,

734 prefix=f"{prefix}.blocks.{layer_idx}",

735 + use_data_parallel=self.use_data_parallel,

736 ) for layer_idx in range(depth)

737 ])

738 self.merger = Glm4vPatchMerger(

@@ -669,6 +741,7 @@ class Glm4vVisionTransformer(nn.Module):

741 quant_config=quant_config,

742 bias=False,

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

744 + use_data_parallel=self.use_data_parallel,

745 )

746 self.embeddings = Glm4vVisionEmbeddings(vision_config)

747

@@ -731,8 +804,11 @@ class Glm4vVisionTransformer(nn.Module):

804 def forward(

805 self,

806 x: torch.Tensor,

734 - grid_thw: torch.Tensor,

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

808 ) -> torch.Tensor:

809 + # Convert grid_thw to tensor (always expecting list format now)

810 + grid_thw = torch.tensor(grid_thw, device=x.device, dtype=torch.long)

811 +

812 # patchify

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

814 x = self.patch_embed(x)

@@ -1250,6 +1326,8 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,

1326 "model.visual.": "visual.",

1327 })

1328

1329 + supports_encoder_tp_data = True

1330 +

1331 @classmethod

1332 def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1333 if modality.startswith("image"):

@@ -1267,12 +1345,14 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,

1345

1346 self.config = config

1347 self.multimodal_config = multimodal_config

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

1349

1350 self.visual = Glm4vVisionTransformer(

1351 config.vision_config,

1352 norm_eps=getattr(config, "rms_norm_eps", 1e-5),

1353 quant_config=quant_config,

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

1355 + use_data_parallel=self.use_data_parallel,

1356 )

1357

1358 if config.model_type == "glm4v":

@@ -1382,8 +1462,14 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,

1462 image_embeds = image_input["image_embeds"].type(self.visual.dtype)

1463 else:

1464 pixel_values = image_input["pixel_values"].type(self.visual.dtype)

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

1386 -

1465 + if self.use_data_parallel:

1466 + return run_dp_sharded_mrope_vision_model(self.visual,

1467 + pixel_values,

1468 + grid_thw.tolist(),

1469 + rope_type="rope_3d")

1470 + else:

1471 + image_embeds = self.visual(pixel_values,

1472 + grid_thw=grid_thw.tolist())

1473 merge_size = self.visual.spatial_merge_size

1474 sizes = grid_thw.prod(-1) // merge_size // merge_size

1475 return image_embeds.split(sizes.tolist())

@@ -1393,23 +1479,22 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,

1479 grid_thw = video_input["video_grid_thw"]

1480 assert grid_thw.ndim == 2

1481

1396 - device = self.visual.device

1397 - flat_grid_thw = torch.cat([

1398 - torch.tensor([[1, h, w]] * t, device=device)

1399 - for t, h, w in grid_thw

1400 - ])

1482 if video_input["type"] == "video_embeds":

1483 video_embeds = video_input["video_embeds"].type(self.visual.dtype)

1484 else:

1485 pixel_values_videos = video_input["pixel_values_videos"].type(

1486 self.visual.dtype)

1406 - video_embeds = self.visual(pixel_values_videos,

1407 - grid_thw=flat_grid_thw)

1408 -

1487 + if self.use_data_parallel:

1488 + return run_dp_sharded_mrope_vision_model(self.visual,

1489 + pixel_values_videos,

1490 + grid_thw.tolist(),

1491 + rope_type="rope_3d")

1492 + else:

1493 + video_embeds = self.visual(pixel_values_videos,

1494 + grid_thw=grid_thw.tolist())

1495 # Split concatenated embeddings for each video item.

1496 merge_size = self.visual.spatial_merge_size

1497 sizes = grid_thw.prod(-1) // merge_size // merge_size

1412 -

1498 return video_embeds.split(sizes.tolist())

1499

1500 def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: