@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 52bc9d5b3e

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] enable data parallel for InternVL vision encoder (#23909)

Signed-off-by: Yiwen Chen <yiwen66@berkeley.edu> Signed-off-by: YiwenC <54658925+666even666@users.noreply.github.com> Co-authored-by: Roger Wang <hey@rogerw.io>

YiwenC committed Sep 17, 2025 at 21:11 UTC 52bc9d5b3edbf8804758d46cde28024d6c362e42

3 files changed +80 -33

docs/configuration/optimization.md

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

175 Known supported models:

176

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

178 +- InternVL (<gh-pr:23909>)

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

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

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

vllm/model_executor/models/intern_vit.py

+75 -32

@@ -25,9 +25,11 @@ from vllm.model_executor.layers.activation import get_act_fn

25 from vllm.model_executor.layers.layernorm import RMSNorm

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

27 QKVParallelLinear,

28 + ReplicatedLinear,

29 RowParallelLinear)

30 from vllm.model_executor.layers.quantization import QuantizationConfig

31 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

32 +from vllm.multimodal.utils import run_dp_sharded_vision_model

34 NORM2FN = {

35 'rms_norm': RMSNorm,

@@ -137,6 +139,7 @@ class InternParallelAttention(nn.Module):

139 *,

140 num_dummy_heads: int = 0,

141 prefix: str = "",

142 + use_data_parallel: bool = False,

143 ) -> None:

144 super().__init__()

145

@@ -150,8 +153,10 @@ class InternParallelAttention(nn.Module):

153 f'(got `embed_dim`: {self.embed_dim} and `num_heads`:'

154 f' {self.num_heads}).')

155

153 - self.tp_size = get_tensor_model_parallel_world_size()

154 - self.tp_rank = get_tensor_model_parallel_rank()

156 + self.tp_size = (1 if use_data_parallel else

157 + get_tensor_model_parallel_world_size())

158 + self.tp_rank = (0 if use_data_parallel else

159 + get_tensor_model_parallel_rank())

160

161 # Additional dummy heads are used to enable TP for common GPU counts.

162 self.dummy_dim = (num_dummy_heads + self.num_heads) * self.head_dim

@@ -159,14 +164,23 @@ class InternParallelAttention(nn.Module):

164 self.tp_size)

165

166 self.scale = self.head_dim**-0.5

162 - self.qkv = QKVParallelLinear(

163 - self.embed_dim,

164 - self.head_dim,

165 - num_dummy_heads + self.num_heads,

166 - bias=config.qkv_bias,

167 - quant_config=quant_config,

168 - prefix=f"{prefix}.qkv",

169 - )

167 + if use_data_parallel:

168 + self.qkv = ReplicatedLinear(

169 + self.embed_dim,

170 + 3 * self.head_dim * self.num_heads,

171 + bias=config.qkv_bias,

172 + quant_config=quant_config,

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

174 + )

175 + else:

176 + self.qkv = QKVParallelLinear(

177 + self.embed_dim,

178 + self.head_dim,

179 + num_dummy_heads + self.num_heads,

180 + bias=config.qkv_bias,

181 + quant_config=quant_config,

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

183 + )

184

185 self.qk_normalization = config.qk_normalization

186

@@ -178,12 +192,20 @@ class InternParallelAttention(nn.Module):

192 eps=config.layer_norm_eps,

193 var_hidden_size=self.embed_dim)

194

181 - self.proj = RowParallelLinear(

182 - self.dummy_dim,

183 - self.embed_dim,

184 - quant_config=quant_config,

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

186 - )

195 + if use_data_parallel:

196 + self.proj = ReplicatedLinear(

197 + self.dummy_dim,

198 + self.embed_dim,

199 + quant_config=quant_config,

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

201 + )

202 + else:

203 + self.proj = RowParallelLinear(

204 + self.dummy_dim,

205 + self.embed_dim,

206 + quant_config=quant_config,

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

208 + )

209

210 self.attn = MultiHeadAttention(self.num_heads_per_partition,

211 self.head_dim, self.scale)

@@ -287,21 +309,26 @@ class InternMLP(nn.Module):

309 config: PretrainedConfig,

310 quant_config: Optional[QuantizationConfig] = None,

311 prefix: str = "",

312 + use_data_parallel: bool = False,

313 ) -> None:

314 super().__init__()

315

316 self.config = config

317 self.activation_fn = get_act_fn(config.hidden_act)

295 - self.fc1 = ColumnParallelLinear(config.hidden_size,

296 - config.intermediate_size,

297 - bias=True,

298 - quant_config=quant_config,

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

300 - self.fc2 = RowParallelLinear(config.intermediate_size,

301 - config.hidden_size,

302 - bias=True,

303 - quant_config=quant_config,

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

318 + cls_fc1 = (ReplicatedLinear

319 + if use_data_parallel else ColumnParallelLinear)

320 + self.fc1 = cls_fc1(config.hidden_size,

321 + config.intermediate_size,

322 + bias=True,

323 + quant_config=quant_config,

324 + prefix=f"{prefix}.fc1")

325 + cls_fc2 = (ReplicatedLinear

326 + if use_data_parallel else RowParallelLinear)

327 + self.fc2 = cls_fc2(config.intermediate_size,

328 + config.hidden_size,

329 + bias=True,

330 + quant_config=quant_config,

331 + prefix=f"{prefix}.fc2")

332

333 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

334 hidden_states, _ = self.fc1(hidden_states)

@@ -320,6 +347,7 @@ class InternVisionEncoderLayer(nn.Module):

347 *,

348 num_dummy_heads: int = 0,

349 prefix: str = "",

350 + use_data_parallel: bool = False,

351 ) -> None:

352 super().__init__()

353

@@ -330,11 +358,13 @@ class InternVisionEncoderLayer(nn.Module):

358 self.attn = self._init_attn(config,

359 quant_config,

360 num_dummy_heads=num_dummy_heads,

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

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

362 + use_data_parallel=use_data_parallel)

363

364 self.mlp = InternMLP(config,

365 quant_config=quant_config,

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

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

367 + use_data_parallel=use_data_parallel)

368 self.norm1 = NORM2FN[self.norm_type](self.embed_dim,

369 eps=config.layer_norm_eps)

370 self.norm2 = NORM2FN[self.norm_type](self.embed_dim,

@@ -352,16 +382,20 @@ class InternVisionEncoderLayer(nn.Module):

382 *,

383 num_dummy_heads: int,

384 prefix: str = "",

385 + use_data_parallel: bool = False,

386 ):

387 # fallback to sdpa attention if tp unavailable

357 - tp_size = get_tensor_model_parallel_world_size()

388 + # tp_size = get_tensor_model_parallel_world_size()

389 + tp_size = (1 if use_data_parallel else

390 + get_tensor_model_parallel_world_size())

391 num_heads = config.num_attention_heads

392

393 if (num_heads + num_dummy_heads) % tp_size == 0:

394 return InternParallelAttention(config,

395 quant_config=quant_config,

396 num_dummy_heads=num_dummy_heads,

364 - prefix=prefix)

397 + prefix=prefix,

398 + use_data_parallel=use_data_parallel)

399

400 return InternSdpaAttention(config, num_dummy_heads=num_dummy_heads)

401

@@ -388,6 +422,7 @@ class InternVisionEncoder(nn.Module):

422 num_hidden_layers_override: Optional[int] = None,

423 num_dummy_heads: int = 0,

424 prefix: str = "",

425 + use_data_parallel: bool = False,

426 ):

427 super().__init__()

428

@@ -402,7 +437,8 @@ class InternVisionEncoder(nn.Module):

437 InternVisionEncoderLayer(config,

438 quant_config,

439 num_dummy_heads=num_dummy_heads,

405 - prefix=f"{prefix}.layers.{layer_idx}")

440 + prefix=f"{prefix}.layers.{layer_idx}",

441 + use_data_parallel=use_data_parallel)

442 for layer_idx in range(num_hidden_layers)

443 ])

444

@@ -429,10 +465,12 @@ class InternVisionModel(nn.Module):

465 num_hidden_layers_override: Optional[int] = None,

466 num_dummy_heads: int = 0,

467 prefix: str = "",

468 + use_data_parallel: bool = False,

469 ) -> None:

470 super().__init__()

471

472 self.config = config

473 + self.use_data_parallel = use_data_parallel

474

475 self.embeddings = InternVisionEmbeddings(config)

476 self.encoder = InternVisionEncoder(

@@ -441,6 +479,7 @@ class InternVisionModel(nn.Module):

479 num_hidden_layers_override=num_hidden_layers_override,

480 num_dummy_heads=num_dummy_heads,

481 prefix=f"{prefix}.encoder",

482 + use_data_parallel=use_data_parallel,

483 )

484

485 def get_input_embeddings(self):

@@ -464,7 +503,11 @@ class InternVisionModel(nn.Module):

503 raise ValueError(

504 f'wrong pixel_values size: {pixel_values.shape}')

505

467 - encoder_outputs = self.encoder(inputs_embeds=hidden_states)

506 + if self.use_data_parallel:

507 + encoder_outputs = run_dp_sharded_vision_model(

508 + hidden_states, self.encoder)

509 + else:

510 + encoder_outputs = self.encoder(inputs_embeds=hidden_states)

511

512 return encoder_outputs

513

vllm/model_executor/models/internvl.py

+4 -1

@@ -1035,6 +1035,8 @@ class InternVLMultiModalProcessor(

1035 class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP,

1036 SupportsLoRA):

1037

1038 + supports_encoder_tp_data = True

1039 +

1040 @classmethod

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

1042 if modality.startswith("image"):

@@ -1053,6 +1055,7 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP,

1055

1056 self.config = config

1057 self.multimodal_config = multimodal_config

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

1059 self._patch_quant_config(config, quant_config)

1060

1061 image_size = config.force_image_size or config.vision_config.image_size

@@ -1120,7 +1123,7 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP,

1123 quant_config=quant_config,

1124 num_hidden_layers_override=num_hidden_layers,

1125 prefix=prefix,

1123 - )

1126 + use_data_parallel=self.use_data_parallel)

1127 else:

1128 return InternVisionPatchModel(config.vision_config)

1129