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
[VLM] Migrate remain DP-supported ViT models to use `disable_tp` (#24363)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Sep 12, 2025 at 02:30 UTC
bb2b5126dac8cbad4c27c379144bbaaa572ccd6c
3 files changed
+54
-106
vllm/model_executor/models/idefics2_vision_model.py
+20
-65
@@ -31,7 +31,6 @@ 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,
34
RowParallelLinear)
35
from vllm.model_executor.layers.quantization import QuantizationConfig
36
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -139,37 +138,22 @@ class Idefics2VisionAttention(nn.Module):
138
assert self.num_heads % tp_size == 0
139
self.num_heads_per_partition = self.num_heads // tp_size
140
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
- )
141
+ self.qkv_proj = QKVParallelLinear(
142
+ self.embed_dim,
143
+ self.head_dim,
144
+ self.num_heads,
145
+ quant_config=quant_config,
146
+ prefix=f"{prefix}.qkv_proj",
147
+ disable_tp=use_data_parallel,
148
+ )
149
+ self.out_proj = RowParallelLinear(
150
+ self.embed_dim,
151
+ self.embed_dim,
152
+ bias=True,
153
+ quant_config=quant_config,
154
+ prefix=f"{prefix}.out_proj",
155
+ disable_tp=use_data_parallel,
156
+ )
157
# Use unified MultiHeadAttention with Flash Attention support
158
self.attn = MultiHeadAttention(self.num_heads_per_partition,
159
self.head_dim, self.scale)
@@ -201,23 +185,21 @@ class Idefics2VisionMLP(nn.Module):
185
super().__init__()
186
self.config = config
187
self.activation_fn = get_act_fn(config.hidden_act)
204
- cls_fc1 = (ReplicatedLinear
205
- if use_data_parallel else ColumnParallelLinear)
206
- self.fc1 = cls_fc1(
188
+ self.fc1 = ColumnParallelLinear(
189
config.hidden_size,
190
config.intermediate_size,
191
bias=True,
192
quant_config=quant_config,
193
prefix=f"{prefix}.fc1",
194
+ disable_tp=use_data_parallel,
195
)
213
- cls_fc2 = (ReplicatedLinear
214
- if use_data_parallel else RowParallelLinear)
215
- self.fc2 = cls_fc2(
196
+ self.fc2 = RowParallelLinear(
197
config.intermediate_size,
198
config.hidden_size,
199
bias=True,
200
quant_config=quant_config,
201
prefix=f"{prefix}.fc2",
202
+ disable_tp=use_data_parallel,
203
)
204
205
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -389,30 +371,6 @@ class Idefics2VisionTransformer(nn.Module):
371
last_hidden_state = self.post_layernorm(encoder_outputs)
372
return last_hidden_state
373
392
- def _consolidate_qkv_weights(
393
- self, weights: Iterable[tuple[str, torch.Tensor]]
394
- ) -> Iterable[tuple[str, torch.Tensor]]:
395
- qkv_idx_mappings = {
396
- ".self_attn.q_proj": 0,
397
- ".self_attn.k_proj": 1,
398
- ".self_attn.v_proj": 2,
399
- }
400
- qkv_weights = {}
401
- for name, loaded_weight in weights:
402
- for weight_name, idx in qkv_idx_mappings.items():
403
- if weight_name not in name:
404
- continue
405
- new_name = name.replace(weight_name, ".self_attn.qkv_proj")
406
- if new_name not in qkv_weights:
407
- qkv_weights[new_name] = [None] * 3
408
- qkv_weights[new_name][idx] = loaded_weight
409
- break
410
- else:
411
- yield name, loaded_weight
412
- for key, weight in qkv_weights.items():
413
- qkv_weight = torch.cat(weight, dim=0)
414
- yield key, qkv_weight
415
-
374
def load_weights(self, weights: Iterable[tuple[str,
375
torch.Tensor]]) -> set[str]:
376
stacked_params_mapping = [
@@ -425,9 +383,6 @@ class Idefics2VisionTransformer(nn.Module):
383
loaded_params: set[str] = set()
384
layer_count = len(self.encoder.layers)
385
428
- if self.use_data_parallel:
429
- weights = self._consolidate_qkv_weights(weights)
430
-
386
for name, loaded_weight in weights:
387
# skip pooling header
388
if name.startswith("head."):
vllm/model_executor/models/mllama4.py
+13
-19
@@ -106,22 +106,21 @@ class Llama4VisionMLP(nn.Module):
106
use_data_parallel: bool = False,
107
):
108
super().__init__()
109
- cls_fc1 = (ReplicatedLinear
110
- if use_data_parallel else ColumnParallelLinear)
111
- self.fc1 = cls_fc1(
109
+ self.fc1 = ColumnParallelLinear(
110
input_size=input_size,
111
output_size=intermediate_size,
112
bias=bias,
113
quant_config=quant_config,
114
prefix=f"{prefix}.fc1",
115
+ disable_tp=use_data_parallel,
116
)
118
- cls_fc2 = ReplicatedLinear if use_data_parallel else RowParallelLinear
119
- self.fc2 = cls_fc2(
117
+ self.fc2 = RowParallelLinear(
118
input_size=intermediate_size,
119
output_size=output_size,
120
bias=bias,
121
quant_config=quant_config,
122
prefix=f"{prefix}.fc2",
123
+ disable_tp=use_data_parallel,
124
)
125
self.activation_fn = nn.GELU()
126
self.output_activation = output_activation
@@ -419,20 +418,15 @@ class Llama4UnfoldConvolution(nn.Module):
418
kernel_size = (kernel_size, kernel_size)
419
self.unfold = torch.nn.Unfold(kernel_size=kernel_size,
420
stride=config.patch_size)
422
- params = {
423
- "input_size":
424
- config.num_channels * kernel_size[0] * kernel_size[1],
425
- "output_size": config.hidden_size,
426
- "bias": False,
427
- "quant_config": quant_config,
428
- "prefix": f"{prefix}.linear",
429
- }
430
- if use_data_parallel:
431
- cls = ReplicatedLinear
432
- else:
433
- cls = ColumnParallelLinear
434
- params["gather_output"] = True
435
- self.linear = cls(**params)
421
+ self.linear = ColumnParallelLinear(
422
+ input_size=config.num_channels * kernel_size[0] * kernel_size[1],
423
+ output_size=config.hidden_size,
424
+ bias=False,
425
+ gather_output=True,
426
+ quant_config=quant_config,
427
+ prefix=f"{prefix}.linear",
428
+ disable_tp=use_data_parallel,
429
+ )
430
431
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
432
hidden_states = self.unfold(hidden_states)
vllm/model_executor/models/qwen2_5_vl.py
+21
-22
@@ -49,7 +49,6 @@ from vllm.model_executor.layers.layernorm import RMSNorm
49
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
50
MergedColumnParallelLinear,
51
QKVParallelLinear,
52
- ReplicatedLinear,
52
RowParallelLinear)
53
# yapf: enable
54
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -510,32 +509,32 @@ class Qwen2_5_VisionPatchMerger(nn.Module):
509
norm_layer = partial(nn.LayerNorm, eps=1e-6)
510
self.ln_q = norm_layer(context_dim)
511
513
- cls_fc1 = (ReplicatedLinear
514
- if use_data_parallel else ColumnParallelLinear)
515
- cls_fc2 = (ReplicatedLinear
516
- if use_data_parallel else RowParallelLinear)
517
- self.mlp = nn.ModuleList([
518
- cls_fc1(self.hidden_size,
519
- self.hidden_size,
520
- bias=True,
521
- quant_config=quant_config,
522
- prefix=f"{prefix}.mlp.0"),
512
+ self.mlp = nn.Sequential(
513
+ ColumnParallelLinear(
514
+ self.hidden_size,
515
+ self.hidden_size,
516
+ bias=True,
517
+ quant_config=quant_config,
518
+ prefix=f"{prefix}.mlp.0",
519
+ return_bias=False,
520
+ disable_tp=use_data_parallel,
521
+ ),
522
nn.GELU(),
524
- cls_fc2(self.hidden_size,
525
- d_model,
526
- bias=True,
527
- quant_config=quant_config,
528
- prefix=f"{prefix}.mlp.2"),
529
- ])
523
+ RowParallelLinear(
524
+ self.hidden_size,
525
+ d_model,
526
+ bias=True,
527
+ quant_config=quant_config,
528
+ prefix=f"{prefix}.mlp.2",
529
+ return_bias=False,
530
+ disable_tp=use_data_parallel,
531
+ ),
532
+ )
533
534
def forward(self, x: torch.Tensor) -> torch.Tensor:
535
x = self.ln_q(x)
536
x = x.view(-1, self.hidden_size)
534
-
535
- mlp_fc1, mlp_act, mlp_fc2 = self.mlp
536
- x_parallel, _ = mlp_fc1(x)
537
- x_parallel = mlp_act(x_parallel)
538
- out, _ = mlp_fc2(x_parallel)
537
+ out = self.mlp(x)
538
return out
539
540