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] Cleanup InternViT's data parallel implementation (#25306)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Sep 20, 2025 at 20:46 UTC
3c713a97116f34ffdc75ba10c9dfbb2850437984
1 file changed
+37
-121
vllm/model_executor/models/intern_vit.py
+37
-121
@@ -25,7 +25,6 @@ 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,
28
RowParallelLinear)
29
from vllm.model_executor.layers.quantization import QuantizationConfig
30
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -164,23 +163,15 @@ class InternParallelAttention(nn.Module):
163
self.tp_size)
164
165
self.scale = self.head_dim**-0.5
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
- )
166
+ self.qkv = QKVParallelLinear(
167
+ self.embed_dim,
168
+ self.head_dim,
169
+ num_dummy_heads + self.num_heads,
170
+ bias=config.qkv_bias,
171
+ quant_config=quant_config,
172
+ prefix=f"{prefix}.qkv",
173
+ disable_tp=use_data_parallel,
174
+ )
175
176
self.qk_normalization = config.qk_normalization
177
@@ -192,20 +183,13 @@ class InternParallelAttention(nn.Module):
183
eps=config.layer_norm_eps,
184
var_hidden_size=self.embed_dim)
185
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
- )
186
+ self.proj = RowParallelLinear(
187
+ self.dummy_dim,
188
+ self.embed_dim,
189
+ quant_config=quant_config,
190
+ prefix=f"{prefix}.proj",
191
+ disable_tp=use_data_parallel,
192
+ )
193
194
self.attn = MultiHeadAttention(self.num_heads_per_partition,
195
self.head_dim, self.scale)
@@ -236,72 +220,6 @@ class InternParallelAttention(nn.Module):
220
return out
221
222
239
-class InternSdpaAttention(nn.Module):
240
- """Multi-headed attention from 'Attention Is All You Need' paper"""
241
-
242
- def __init__(
243
- self,
244
- config: PretrainedConfig,
245
- *,
246
- num_dummy_heads: int = 0,
247
- ) -> None:
248
- super().__init__()
249
-
250
- self.config = config
251
- self.embed_dim = config.hidden_size
252
- self.num_heads = config.num_attention_heads
253
- self.head_dim = self.embed_dim // self.num_heads
254
- if self.head_dim * self.num_heads != self.embed_dim:
255
- raise ValueError(
256
- f'embed_dim must be divisible by num_heads '
257
- f'(got `embed_dim`: {self.embed_dim} and `num_heads`:'
258
- f' {self.num_heads}).')
259
-
260
- # Additional dummy heads are used to enable TP for common GPU counts.
261
- self.dummy_dim = (num_dummy_heads + self.num_heads) * self.head_dim
262
-
263
- self.scale = self.head_dim**-0.5
264
- self.qkv = nn.Linear(self.embed_dim,
265
- 3 * self.dummy_dim,
266
- bias=config.qkv_bias)
267
-
268
- self.qk_normalization = config.qk_normalization
269
-
270
- if self.qk_normalization:
271
- self.q_norm = RMSNorm(self.dummy_dim,
272
- eps=config.layer_norm_eps,
273
- var_hidden_size=self.embed_dim)
274
- self.k_norm = RMSNorm(self.dummy_dim,
275
- eps=config.layer_norm_eps,
276
- var_hidden_size=self.embed_dim)
277
-
278
- self.proj = nn.Linear(self.dummy_dim, self.embed_dim)
279
-
280
- # Use unified MultiHeadAttention with automatic backend selection
281
- self.attn = MultiHeadAttention(self.num_heads, self.head_dim,
282
- self.scale)
283
-
284
- def forward(self, x: torch.Tensor) -> torch.Tensor:
285
- B, N, C = x.shape
286
- qkv = self.qkv(x)
287
- q, k, v = qkv.chunk(3, dim=-1)
288
-
289
- q = q.view(B, N, self.num_heads, self.head_dim)
290
- k = k.view(B, N, self.num_heads, self.head_dim)
291
- v = v.view(B, N, self.num_heads, self.head_dim)
292
-
293
- if self.qk_normalization:
294
- B_, N_, H_, D_ = q.shape
295
- q = self.q_norm(q.flatten(-2, -1)).view(B_, N_, H_, D_)
296
- k = self.k_norm(k.flatten(-2, -1)).view(B_, N_, H_, D_)
297
-
298
- # Use unified MultiHeadAttention with automatic backend selection
299
- x = self.attn(q, k, v)
300
-
301
- x = self.proj(x)
302
- return x
303
-
304
-
223
class InternMLP(nn.Module):
224
225
def __init__(
@@ -315,20 +233,18 @@ class InternMLP(nn.Module):
233
234
self.config = config
235
self.activation_fn = get_act_fn(config.hidden_act)
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")
236
+ self.fc1 = ColumnParallelLinear(config.hidden_size,
237
+ config.intermediate_size,
238
+ bias=True,
239
+ quant_config=quant_config,
240
+ prefix=f"{prefix}.fc1",
241
+ disable_tp=use_data_parallel)
242
+ self.fc2 = RowParallelLinear(config.intermediate_size,
243
+ config.hidden_size,
244
+ bias=True,
245
+ quant_config=quant_config,
246
+ prefix=f"{prefix}.fc2",
247
+ disable_tp=use_data_parallel)
248
249
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
250
hidden_states, _ = self.fc1(hidden_states)
@@ -385,19 +301,19 @@ class InternVisionEncoderLayer(nn.Module):
301
use_data_parallel: bool = False,
302
):
303
# fallback to sdpa attention if tp unavailable
388
- # tp_size = get_tensor_model_parallel_world_size()
304
tp_size = (1 if use_data_parallel else
305
get_tensor_model_parallel_world_size())
306
num_heads = config.num_attention_heads
307
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,
397
- prefix=prefix,
398
- use_data_parallel=use_data_parallel)
399
-
400
- return InternSdpaAttention(config, num_dummy_heads=num_dummy_heads)
308
+ # if the number of heads is not divisible by tp_size,
309
+ # we also disable Attention's TP
310
+ use_data_parallel = (use_data_parallel
311
+ or (num_heads + num_dummy_heads) % tp_size != 0)
312
+ return InternParallelAttention(config,
313
+ quant_config=quant_config,
314
+ num_dummy_heads=num_dummy_heads,
315
+ prefix=prefix,
316
+ use_data_parallel=use_data_parallel)
317
318
def forward(
319
self,