@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 23c939fd30

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 for ViT on MiniCPM-V-4 (#23327)

Signed-off-by: ycyaw66 <497410282@qq.com> Co-authored-by: ycyaw66 <497410282@qq.com>

WeiQing Chen committed Aug 23, 2025 at 10:14 UTC 23c939fd30b845d06568304a0e9d168a48cc8cb4

4 files changed +105 -30

docs/configuration/optimization.md

@@ -172,6 +172,7 @@ The availablilty of batch-level DP is based on model implementation.

172 Currently, the following models support `mm_encoder_tp_mode="data"`:

173

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

175 +- MiniCPM-V-4 (<gh-pr:23327>)

176 - Qwen2.5-VL (<gh-pr:22742>)

177 - Step3 (<gh-pr:22697>)

178

vllm/model_executor/models/idefics2_vision_model.py

+96 -27

@@ -27,13 +27,15 @@ from transformers.models.idefics2.configuration_idefics2 import (

27 Idefics2Config, Idefics2VisionConfig)

29 from vllm.attention.layer import MultiHeadAttention

30 -from vllm.distributed import divide, get_tensor_model_parallel_world_size

30 +from vllm.distributed import get_tensor_model_parallel_world_size

31 from vllm.model_executor.layers.activation import get_act_fn

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

33 QKVParallelLinear,

34 + ReplicatedLinear,

35 RowParallelLinear)

36 from vllm.model_executor.layers.quantization import QuantizationConfig

37 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

38 +from vllm.multimodal.utils import run_dp_sharded_vision_model

41 class Idefics2VisionEmbeddings(nn.Module):

@@ -118,6 +120,7 @@ class Idefics2VisionAttention(nn.Module):

120 config: Idefics2VisionConfig,

121 quant_config: Optional[QuantizationConfig] = None,

122 prefix: str = "",

123 + use_data_parallel: bool = False,

124 ) -> None:

125 super().__init__()

126 self.config = config

@@ -130,22 +133,43 @@ class Idefics2VisionAttention(nn.Module):

133 f" {self.num_heads}).")

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

135 self.dropout = config.attention_dropout

133 - self.qkv_proj = QKVParallelLinear(

134 - self.embed_dim,

135 - self.head_dim,

136 - self.num_heads,

137 - quant_config=quant_config,

138 - prefix=f"{prefix}.qkv_proj",

139 - )

140 - self.out_proj = RowParallelLinear(

141 - self.embed_dim,

142 - self.embed_dim,

143 - bias=True,

144 - quant_config=quant_config,

145 - prefix=f"{prefix}.out_proj",

146 - )

147 - self.tp_size = get_tensor_model_parallel_world_size()

148 - self.num_heads_per_partition = divide(self.num_heads, self.tp_size)

136 +

137 + tp_size = (1 if use_data_parallel else

138 + get_tensor_model_parallel_world_size())

139 + assert self.num_heads % tp_size == 0

140 + self.num_heads_per_partition = self.num_heads // tp_size

141 +

142 + if use_data_parallel:

143 + self.q_size = self.num_heads * self.head_dim

144 + self.qkv_proj = ReplicatedLinear(

145 + self.embed_dim,

146 + 3 * self.q_size,

147 + bias=True,

148 + quant_config=quant_config,

149 + prefix=f"{prefix}.qkv_proj",

150 + )

151 + self.out_proj = ReplicatedLinear(

152 + self.embed_dim,

153 + self.embed_dim,

154 + bias=True,

155 + quant_config=quant_config,

156 + prefix=f"{prefix}.out_proj",

157 + )

158 + else:

159 + self.qkv_proj = QKVParallelLinear(

160 + self.embed_dim,

161 + self.head_dim,

162 + self.num_heads,

163 + quant_config=quant_config,

164 + prefix=f"{prefix}.qkv_proj",

165 + )

166 + self.out_proj = RowParallelLinear(

167 + self.embed_dim,

168 + self.embed_dim,

169 + bias=True,

170 + quant_config=quant_config,

171 + prefix=f"{prefix}.out_proj",

172 + )

173 self.attn = MultiHeadAttention(self.num_heads_per_partition,

174 self.head_dim, self.scale)

175

@@ -169,18 +193,23 @@ class Idefics2VisionMLP(nn.Module):

193 config: Idefics2VisionConfig,

194 quant_config: Optional[QuantizationConfig] = None,

195 prefix: str = "",

196 + use_data_parallel: bool = False,

197 ) -> None:

198 super().__init__()

199 self.config = config

200 self.activation_fn = get_act_fn(config.hidden_act)

176 - self.fc1 = ColumnParallelLinear(

201 + cls_fc1 = (ReplicatedLinear

202 + if use_data_parallel else ColumnParallelLinear)

203 + self.fc1 = cls_fc1(

204 config.hidden_size,

205 config.intermediate_size,

206 bias=True,

207 quant_config=quant_config,

208 prefix=f"{prefix}.fc1",

209 )

183 - self.fc2 = RowParallelLinear(

210 + cls_fc2 = (ReplicatedLinear

211 + if use_data_parallel else RowParallelLinear)

212 + self.fc2 = cls_fc2(

213 config.intermediate_size,

214 config.hidden_size,

215 bias=True,

@@ -202,17 +231,21 @@ class Idefics2EncoderLayer(nn.Module):

231 config: Idefics2Config,

232 quant_config: Optional[QuantizationConfig] = None,

233 prefix: str = "",

234 + use_data_parallel: bool = False,

235 ) -> None:

236 super().__init__()

237 self.embed_dim = config.hidden_size

208 - self.self_attn = Idefics2VisionAttention(config,

209 - quant_config=quant_config,

210 - prefix=f"{prefix}.self_attn")

238 + self.self_attn = Idefics2VisionAttention(

239 + config,

240 + quant_config=quant_config,

241 + prefix=f"{prefix}.self_attn",

242 + use_data_parallel=use_data_parallel)

243 self.layer_norm1 = nn.LayerNorm(self.embed_dim,

244 eps=config.layer_norm_eps)

245 self.mlp = Idefics2VisionMLP(config,

246 quant_config=quant_config,

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

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

248 + use_data_parallel=use_data_parallel)

249 self.layer_norm2 = nn.LayerNorm(self.embed_dim,

250 eps=config.layer_norm_eps)

251

@@ -254,6 +287,7 @@ class Idefics2Encoder(nn.Module):

287 *,

288 num_hidden_layers_override: Optional[int] = None,

289 prefix: str = "",

290 + use_data_parallel: bool = False,

291 ) -> None:

292 super().__init__()

293

@@ -267,7 +301,8 @@ class Idefics2Encoder(nn.Module):

301 self.layers = nn.ModuleList([

302 Idefics2EncoderLayer(config,

303 quant_config=quant_config,

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

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

305 + use_data_parallel=use_data_parallel)

306 for layer_idx in range(num_hidden_layers)

307 ])

308

@@ -301,17 +336,20 @@ class Idefics2VisionTransformer(nn.Module):

336 num_hidden_layers_override: Optional[int] = None,

337 require_post_norm: bool = True,

338 prefix: str = "",

339 + use_data_parallel: bool = False,

340 ) -> None:

341 super().__init__()

342

343 embed_dim = config.hidden_size

344 self.config = config

345 + self.use_data_parallel = use_data_parallel

346 self.embeddings = Idefics2VisionEmbeddings(config)

347 self.encoder = Idefics2Encoder(

348 config,

349 quant_config=quant_config,

350 num_hidden_layers_override=num_hidden_layers_override,

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

351 + prefix=f"{prefix}.encoder",

352 + use_data_parallel=use_data_parallel)

353

354 num_hidden_layers = config.num_hidden_layers

355 if len(self.encoder.layers) > config.num_hidden_layers:

@@ -340,10 +378,38 @@ class Idefics2VisionTransformer(nn.Module):

378 patch_attention_mask=patch_attention_mask,

379 tgt_sizes=tgt_sizes,

380 )

343 - encoder_outputs = self.encoder(hidden_states)

381 + if self.use_data_parallel:

382 + encoder_outputs = run_dp_sharded_vision_model(

383 + hidden_states, self.encoder)

384 + else:

385 + encoder_outputs = self.encoder(hidden_states)

386 last_hidden_state = self.post_layernorm(encoder_outputs)

387 return last_hidden_state

388

389 + def _consolidate_qkv_weights(

390 + self, weights: Iterable[tuple[str, torch.Tensor]]

391 + ) -> Iterable[tuple[str, torch.Tensor]]:

392 + qkv_idx_mappings = {

393 + ".self_attn.q_proj": 0,

394 + ".self_attn.k_proj": 1,

395 + ".self_attn.v_proj": 2,

396 + }

397 + qkv_weights = {}

398 + for name, loaded_weight in weights:

399 + for weight_name, idx in qkv_idx_mappings.items():

400 + if weight_name not in name:

401 + continue

402 + new_name = name.replace(weight_name, ".self_attn.qkv_proj")

403 + if new_name not in qkv_weights:

404 + qkv_weights[new_name] = [None] * 3

405 + qkv_weights[new_name][idx] = loaded_weight

406 + break

407 + else:

408 + yield name, loaded_weight

409 + for key, weight in qkv_weights.items():

410 + qkv_weight = torch.cat(weight, dim=0)

411 + yield key, qkv_weight

412 +

413 def load_weights(self, weights: Iterable[tuple[str,

414 torch.Tensor]]) -> set[str]:

415 stacked_params_mapping = [

@@ -356,6 +422,9 @@ class Idefics2VisionTransformer(nn.Module):

422 loaded_params: set[str] = set()

423 layer_count = len(self.encoder.layers)

424

425 + if self.use_data_parallel:

426 + weights = self._consolidate_qkv_weights(weights)

427 +

428 for name, loaded_weight in weights:

429 # skip pooling header

430 if name.startswith("head."):

@@ -373,7 +442,7 @@ class Idefics2VisionTransformer(nn.Module):

442 continue

443

444 for param_name, weight_name, shard_id in stacked_params_mapping:

376 - if weight_name not in name:

445 + if weight_name not in name or self.use_data_parallel:

446 continue

447 name = name.replace(weight_name, param_name)

448 param = params_dict[name]

vllm/model_executor/models/minicpmv.py

+6 -3

@@ -778,6 +778,7 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

778 # and config class

779 self.config = config

780 self.multimodal_config = multimodal_config

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

782

783 self.version = get_version_by_config(self.config)

784 self.llm = self.init_llm(vllm_config=vllm_config,

@@ -1325,9 +1326,11 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):

1326 prefix: str = "",

1327 ) -> nn.Module:

1328 quant_config = self._maybe_ignore_quant_config(quant_config)

1328 - model = Idefics2VisionTransformer(config.vision_config,

1329 - quant_config=quant_config,

1330 - prefix=prefix)

1329 + model = Idefics2VisionTransformer(

1330 + config.vision_config,

1331 + quant_config=quant_config,

1332 + prefix=prefix,

1333 + use_data_parallel=self.use_data_parallel)

1334 if self.config.drop_vision_last_layer:

1335 model.encoder.layers = model.encoder.layers[:-1]

1336 return model

vllm/multimodal/utils.py

@@ -461,6 +461,8 @@ def run_dp_sharded_vision_model(image_input: torch.Tensor,

461 num_chunks_per_rank, ...]

462

463 vision_embeddings = vision_model(image_input_per_rank)

464 + # Ensure tensor is contiguous before all_gather

465 + vision_embeddings = vision_embeddings.contiguous()

466 vision_embeddings = tensor_model_parallel_all_gather(vision_embeddings,

467 dim=0)

468 vision_embeddings = vision_embeddings[:num_chunks, ...]