@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5400014d55

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 9, 10:00 UTC

[Chore] Remove `use_data_parallel` kwargs from ViT implementation (#33310)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Isotr0py committed Jan 29, 2026 at 18:20 UTC 5400014d55b4cb388fdf7dbaab6c90d2161e2764

9 files changed +36 -89

vllm/model_executor/models/idefics2_vision_model.py

+5 -11

@@ -39,7 +39,7 @@ from vllm.model_executor.layers.linear import (

39 from vllm.model_executor.layers.quantization import QuantizationConfig

40 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

42 -from .vision import run_dp_sharded_vision_model

42 +from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model

45 class Idefics2VisionEmbeddings(nn.Module):

@@ -126,9 +126,9 @@ class Idefics2VisionAttention(nn.Module):

126 config: Idefics2VisionConfig,

127 quant_config: QuantizationConfig | None = None,

128 prefix: str = "",

129 - use_data_parallel: bool = False,

129 ) -> None:

130 super().__init__()

131 + use_data_parallel = is_vit_use_data_parallel()

132 self.config = config

133 self.embed_dim = config.hidden_size

134 self.num_heads = config.num_attention_heads

@@ -187,11 +187,12 @@ class Idefics2VisionMLP(nn.Module):

187 config: Idefics2VisionConfig,

188 quant_config: QuantizationConfig | None = None,

189 prefix: str = "",

190 - use_data_parallel: bool = False,

190 ) -> None:

191 super().__init__()

192 self.config = config

193 self.activation_fn = get_act_fn(config.hidden_act)

194 +

195 + use_data_parallel = is_vit_use_data_parallel()

196 self.fc1 = ColumnParallelLinear(

197 config.hidden_size,

198 config.intermediate_size,

@@ -222,7 +223,6 @@ class Idefics2EncoderLayer(nn.Module):

223 config: Idefics2Config,

224 quant_config: QuantizationConfig | None = None,

225 prefix: str = "",

225 - use_data_parallel: bool = False,

226 ) -> None:

227 super().__init__()

228 self.embed_dim = config.hidden_size

@@ -230,14 +230,12 @@ class Idefics2EncoderLayer(nn.Module):

230 config,

231 quant_config=quant_config,

232 prefix=f"{prefix}.self_attn",

233 - use_data_parallel=use_data_parallel,

233 )

234 self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

235 self.mlp = Idefics2VisionMLP(

236 config,

237 quant_config=quant_config,

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

240 - use_data_parallel=use_data_parallel,

239 )

240 self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

241

@@ -279,7 +277,6 @@ class Idefics2Encoder(nn.Module):

277 *,

278 num_hidden_layers_override: int | None = None,

279 prefix: str = "",

282 - use_data_parallel: bool = False,

280 ) -> None:

281 super().__init__()

282

@@ -296,7 +293,6 @@ class Idefics2Encoder(nn.Module):

293 config,

294 quant_config=quant_config,

295 prefix=f"{prefix}.layers.{layer_idx}",

299 - use_data_parallel=use_data_parallel,

296 )

297 for layer_idx in range(num_hidden_layers)

298 ]

@@ -331,20 +327,18 @@ class Idefics2VisionTransformer(nn.Module):

327 num_hidden_layers_override: int | None = None,

328 require_post_norm: bool = True,

329 prefix: str = "",

334 - use_data_parallel: bool = False,

330 ) -> None:

331 super().__init__()

332

333 embed_dim = config.hidden_size

334 self.config = config

340 - self.use_data_parallel = use_data_parallel

335 + self.use_data_parallel = is_vit_use_data_parallel()

336 self.embeddings = Idefics2VisionEmbeddings(config)

337 self.encoder = Idefics2Encoder(

338 config,

339 quant_config=quant_config,

340 num_hidden_layers_override=num_hidden_layers_override,

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

347 - use_data_parallel=use_data_parallel,

342 )

343

344 num_hidden_layers = config.num_hidden_layers

vllm/model_executor/models/intern_vit.py

+10 -25

@@ -34,7 +34,7 @@ from vllm.model_executor.layers.linear import (

34 from vllm.model_executor.layers.quantization import QuantizationConfig

35 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

37 -from .vision import run_dp_sharded_vision_model

37 +from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model

39 NORM2FN = {

40 "rms_norm": RMSNorm,

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

148 *,

149 num_dummy_heads: int = 0,

150 prefix: str = "",

151 - use_data_parallel: bool = False,

151 ) -> None:

152 super().__init__()

153

@@ -163,9 +162,14 @@ class InternParallelAttention(nn.Module):

162 f" {self.num_heads})."

163 )

164

166 - self.tp_size = (

167 - 1 if use_data_parallel else get_tensor_model_parallel_world_size()

165 + use_data_parallel = is_vit_use_data_parallel()

166 + # if the number of heads is not divisible by tp_size,

167 + # we also disable Attention's TP

168 + tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()

169 + use_data_parallel = (

170 + use_data_parallel or (self.num_heads + num_dummy_heads) % tp_size != 0

171 )

172 + self.tp_size = 1 if use_data_parallel else tp_size

173 self.tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank()

174

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

@@ -242,12 +246,12 @@ class InternMLP(nn.Module):

246 config: PretrainedConfig,

247 quant_config: QuantizationConfig | None = None,

248 prefix: str = "",

245 - use_data_parallel: bool = False,

249 ) -> None:

250 super().__init__()

251

252 self.config = config

253 self.activation_fn = get_act_fn(config.hidden_act)

254 + use_data_parallel = is_vit_use_data_parallel()

255 self.fc1 = ColumnParallelLinear(

256 config.hidden_size,

257 config.intermediate_size,

@@ -281,11 +285,9 @@ class InternVisionEncoderLayer(nn.Module):

285 *,

286 num_dummy_heads: int = 0,

287 prefix: str = "",

284 - use_data_parallel: bool = False,

288 attn_cls: type[InternParallelAttention] = InternParallelAttention,

289 ) -> None:

290 super().__init__()

288 -

291 self.embed_dim = config.hidden_size

292 self.intermediate_size = config.intermediate_size

293 self.norm_type = config.norm_type

@@ -296,14 +298,12 @@ class InternVisionEncoderLayer(nn.Module):

298 quant_config,

299 num_dummy_heads=num_dummy_heads,

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

299 - use_data_parallel=use_data_parallel,

301 )

302

303 self.mlp = InternMLP(

304 config,

305 quant_config=quant_config,

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

306 - use_data_parallel=use_data_parallel,

307 )

308 self.norm1 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)

309 self.norm2 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)

@@ -318,23 +318,12 @@ class InternVisionEncoderLayer(nn.Module):

318 *,

319 num_dummy_heads: int,

320 prefix: str = "",

321 - use_data_parallel: bool = False,

321 ):

323 - # fallback to sdpa attention if tp unavailable

324 - tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()

325 - num_heads = config.num_attention_heads

326 -

327 - # if the number of heads is not divisible by tp_size,

328 - # we also disable Attention's TP

329 - use_data_parallel = (

330 - use_data_parallel or (num_heads + num_dummy_heads) % tp_size != 0

331 - )

322 return self.attn_cls(

323 config,

324 quant_config=quant_config,

325 num_dummy_heads=num_dummy_heads,

326 prefix=prefix,

337 - use_data_parallel=use_data_parallel,

327 )

328

329 def forward(

@@ -357,7 +346,6 @@ class InternVisionEncoder(nn.Module):

346 num_hidden_layers_override: int | None = None,

347 num_dummy_heads: int = 0,

348 prefix: str = "",

360 - use_data_parallel: bool = False,

349 layer_cls: type[InternVisionEncoderLayer] = InternVisionEncoderLayer,

350 ):

351 super().__init__()

@@ -377,7 +365,6 @@ class InternVisionEncoder(nn.Module):

365 quant_config,

366 num_dummy_heads=num_dummy_heads,

367 prefix=f"{prefix}.layers.{layer_idx}",

380 - use_data_parallel=use_data_parallel,

368 )

369 for layer_idx in range(num_hidden_layers)

370 ]

@@ -404,12 +391,11 @@ class InternVisionModel(nn.Module):

391 num_hidden_layers_override: int | None = None,

392 num_dummy_heads: int = 0,

393 prefix: str = "",

407 - use_data_parallel: bool = False,

394 ) -> None:

395 super().__init__()

396

397 self.config = config

412 - self.use_data_parallel = use_data_parallel

398 + self.use_data_parallel = is_vit_use_data_parallel()

399

400 self.embeddings = InternVisionEmbeddings(config)

401 self.encoder = InternVisionEncoder(

@@ -418,7 +404,6 @@ class InternVisionModel(nn.Module):

404 num_hidden_layers_override=num_hidden_layers_override,

405 num_dummy_heads=num_dummy_heads,

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

421 - use_data_parallel=use_data_parallel,

407 )

408

409 def get_input_embeddings(self):

vllm/model_executor/models/internvl.py

-1

@@ -1153,7 +1153,6 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)

1153 quant_config=quant_config,

1154 num_hidden_layers_override=num_hidden_layers,

1155 prefix=prefix,

1156 - use_data_parallel=self.use_data_parallel,

1156 )

1157 else:

1158 return InternVisionPatchModel(config.vision_config)

vllm/model_executor/models/kimi_vl.py

+5 -4

@@ -81,7 +81,7 @@ from vllm.transformers_utils.configs import KimiVLConfig, MoonViTConfig

81 from vllm.utils.tensor_schema import TensorSchema, TensorShape

83 from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix

84 -from .vision import run_dp_sharded_mrope_vision_model

84 +from .vision import is_vit_use_data_parallel, run_dp_sharded_mrope_vision_model

87 # For dummy input only

@@ -93,10 +93,12 @@ class MaxImageTokenMeta:

94 class KimiVLMultiModalProjector(nn.Module):

95 def __init__(

96 - self, config: KimiVLConfig, use_data_parallel: bool = False, prefix: str = ""

96 + self,

97 + config: KimiVLConfig,

98 + prefix: str = "",

99 ):

100 super().__init__()

99 - self.use_data_parallel = use_data_parallel

101 + self.use_data_parallel = is_vit_use_data_parallel()

102

103 self.hidden_size = (

104 config.vision_config.hidden_size

@@ -321,7 +323,6 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

323 )

324 self.multi_modal_projector = KimiVLMultiModalProjector(

325 config=config,

324 - use_data_parallel=self.use_data_parallel,

326 prefix=maybe_prefix(prefix, "multi_modal_projector"),

327 )

328

vllm/model_executor/models/lfm2_vl.py

+5 -3

@@ -57,6 +57,7 @@ from .utils import (

57 init_vllm_registered_model,

58 maybe_prefix,

59 )

60 +from .vision import is_vit_use_data_parallel

63 class Lfm2VLImagePixelInputs(TensorSchema):

@@ -426,10 +427,12 @@ class Lfm2VLMultiModalProcessor(BaseMultiModalProcessor[Lfm2VLProcessingInfo]):

427

428 class Lfm2VLMultiModalProjector(nn.Module):

429 def __init__(

429 - self, config: Lfm2VlConfig, use_data_parallel: bool = False, prefix: str = ""

430 + self,

431 + config: Lfm2VlConfig,

432 + prefix: str = "",

433 ):

434 super().__init__()

432 - self.use_data_parallel = use_data_parallel

435 + self.use_data_parallel = is_vit_use_data_parallel()

436

437 in_channels = config.vision_config.hidden_size * (config.downsample_factor**2)

438 self.factor = config.downsample_factor

@@ -607,7 +610,6 @@ class Lfm2VLForConditionalGeneration(

610

611 self.multi_modal_projector = Lfm2VLMultiModalProjector(

612 config=config,

610 - use_data_parallel=self.use_data_parallel,

613 prefix=maybe_prefix(prefix, "multi_modal_projector"),

614 )

615

vllm/model_executor/models/minicpmv.py

-4

@@ -1335,7 +1335,6 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA):

1335 config.vision_config,

1336 quant_config=quant_config,

1337 prefix=prefix,

1338 - use_data_parallel=self.use_data_parallel,

1338 )

1339 if self.config.drop_vision_last_layer:

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

@@ -1428,7 +1427,6 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):

1427 config.vision_config,

1428 quant_config=quant_config,

1429 prefix=prefix,

1431 - use_data_parallel=self.use_data_parallel,

1430 )

1431 if self.config.drop_vision_last_layer:

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

@@ -1526,7 +1524,6 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):

1524 config.vision_config,

1525 quant_config=quant_config,

1526 prefix=prefix,

1529 - use_data_parallel=self.use_data_parallel,

1527 )

1528 if self.config.drop_vision_last_layer:

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

@@ -1624,7 +1621,6 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA):

1621 config.vision_config,

1622 quant_config=quant_config,

1623 prefix=prefix,

1627 - use_data_parallel=self.use_data_parallel,

1624 )

1625 if self.config.drop_vision_last_layer:

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

vllm/model_executor/models/mllama4.py

+4 -16

@@ -79,7 +79,7 @@ from .interfaces import (

79 )

80 from .llama4 import Llama4ForCausalLM

81 from .utils import AutoWeightsLoader, StageMissingLayer, maybe_prefix

82 -from .vision import run_dp_sharded_vision_model

82 +from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model

85 class Llama4ImagePatchInputs(TensorSchema):

@@ -124,9 +124,9 @@ class Llama4VisionMLP(nn.Module):

124 output_activation: bool,

125 quant_config: QuantizationConfig | None = None,

126 prefix: str = "",

127 - use_data_parallel: bool = False,

127 ):

128 super().__init__()

129 + use_data_parallel = is_vit_use_data_parallel()

130 self.fc1 = ColumnParallelLinear(

131 input_size=input_size,

132 output_size=intermediate_size,

@@ -208,7 +208,6 @@ class Llama4VisionPixelShuffleMLP(nn.Module):

208 config,

209 quant_config: QuantizationConfig | None = None,

210 prefix: str = "",

211 - use_data_parallel: bool = False,

211 ):

212 super().__init__()

213 self.pixel_shuffle_ratio = config.pixel_shuffle_ratio

@@ -224,7 +223,6 @@ class Llama4VisionPixelShuffleMLP(nn.Module):

223 output_activation=True,

224 quant_config=quant_config,

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

227 - use_data_parallel=use_data_parallel,

226 )

227

228 def forward(self, encoded_patches: torch.Tensor) -> torch.Tensor:

@@ -238,10 +236,10 @@ class Llama4VisionAttention(nn.Module):

236 config: Llama4VisionConfig,

237 quant_config: QuantizationConfig | None,

238 prefix: str = "",

241 - use_data_parallel: bool = False,

239 ):

240 super().__init__()

241 self.config = config

242 + use_data_parallel = is_vit_use_data_parallel()

243 self.tp_size = (

244 1 if use_data_parallel else get_tensor_model_parallel_world_size()

245 )

@@ -336,7 +334,6 @@ class Llama4VisionEncoderLayer(nn.Module):

334 config: Llama4VisionConfig,

335 quant_config: QuantizationConfig | None,

336 prefix: str = "",

339 - use_data_parallel: bool = False,

337 ):

338 super().__init__()

339 self.hidden_size = config.hidden_size

@@ -347,7 +344,6 @@ class Llama4VisionEncoderLayer(nn.Module):

344 config,

345 quant_config=quant_config,

346 prefix=f"{prefix}.self_attn",

350 - use_data_parallel=use_data_parallel,

347 )

348 self.mlp = Llama4VisionMLP(

349 input_size=config.hidden_size,

@@ -357,7 +353,6 @@ class Llama4VisionEncoderLayer(nn.Module):

353 output_activation=False,

354 quant_config=quant_config,

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

360 - use_data_parallel=use_data_parallel,

356 )

357

358 self.input_layernorm = nn.LayerNorm(config.hidden_size)

@@ -389,7 +384,6 @@ class Llama4VisionEncoder(nn.Module):

384 config: Llama4VisionConfig,

385 quant_config: QuantizationConfig | None,

386 prefix: str = "",

392 - use_data_parallel: bool = False,

387 ):

388 super().__init__()

389 self.config = config

@@ -399,7 +393,6 @@ class Llama4VisionEncoder(nn.Module):

393 config,

394 quant_config=quant_config,

395 prefix=f"{prefix}.layers.{layer_idx}",

402 - use_data_parallel=use_data_parallel,

396 )

397 for layer_idx in range(config.num_hidden_layers)

398 ]

@@ -432,13 +425,13 @@ class Llama4UnfoldConvolution(nn.Module):

425 config: Llama4VisionConfig,

426 quant_config: QuantizationConfig | None = None,

427 prefix: str = "",

435 - use_data_parallel: bool = False,

428 ):

429 super().__init__()

430 kernel_size = config.patch_size

431 if isinstance(kernel_size, int):

432 kernel_size = (kernel_size, kernel_size)

433 self.unfold = torch.nn.Unfold(kernel_size=kernel_size, stride=config.patch_size)

434 + use_data_parallel = is_vit_use_data_parallel()

435 self.linear = ColumnParallelLinear(

436 input_size=config.num_channels * kernel_size[0] * kernel_size[1],

437 output_size=config.hidden_size,

@@ -465,7 +458,6 @@ class Llama4VisionModel(nn.Module):

458 config: Llama4VisionConfig,

459 quant_config: QuantizationConfig | None = None,

460 prefix: str = "",

468 - use_data_parallel: bool = False,

461 ):

462 super().__init__()

463 self.config = config

@@ -481,7 +473,6 @@ class Llama4VisionModel(nn.Module):

473 config,

474 quant_config=quant_config,

475 prefix=f"{prefix}.patch_embedding",

484 - use_data_parallel=use_data_parallel,

476 )

477

478 self.class_embedding = nn.Parameter(self.scale * torch.randn(self.hidden_size))

@@ -498,14 +489,12 @@ class Llama4VisionModel(nn.Module):

489 config,

490 quant_config=quant_config,

491 prefix=f"{prefix}.model",

501 - use_data_parallel=use_data_parallel,

492 )

493

494 self.vision_adapter = Llama4VisionPixelShuffleMLP(

495 config,

496 quant_config,

497 prefix=f"{prefix}.vision_adapter",

508 - use_data_parallel=use_data_parallel,

498 )

499

500 def forward(

@@ -780,7 +769,6 @@ class Llama4ForConditionalGeneration(

769 config=config.vision_config,

770 quant_config=None,

771 prefix=maybe_prefix(prefix, "vision_model"),

783 - use_data_parallel=self.use_data_parallel,

772 )

773

774 self.multi_modal_projector = Llama4MultiModalProjector(

vllm/model_executor/models/step3_vl.py

+4 -14

@@ -54,7 +54,7 @@ from .utils import (

54 init_vllm_registered_model,

55 maybe_prefix,

56 )

57 -from .vision import run_dp_sharded_vision_model

57 +from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model

60 class Step3VLImagePixelInputs(TensorSchema):

@@ -724,7 +724,6 @@ class Step3VisionAttention(nn.Module):

724 config,

725 quant_config: QuantizationConfig | None = None,

726 prefix: str = "",

727 - use_data_parallel: bool = False,

727 ):

728 super().__init__()

729 self.config = config

@@ -734,6 +733,7 @@ class Step3VisionAttention(nn.Module):

733

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

735

736 + use_data_parallel = is_vit_use_data_parallel()

737 tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()

738 assert self.total_num_heads % tp_size == 0

739 self.num_heads = self.total_num_heads // tp_size

@@ -786,11 +786,11 @@ class Step3VisionMLP(nn.Module):

786 config,

787 quant_config: QuantizationConfig | None = None,

788 prefix: str = "",

789 - use_data_parallel: bool = False,

789 ):

790 super().__init__()

791 self.config = config

792 self.activation_fn = get_act_fn(config.hidden_act)

793 + use_data_parallel = is_vit_use_data_parallel()

794 self.fc1 = ColumnParallelLinear(

795 config.hidden_size,

796 config.intermediate_size,

@@ -821,23 +821,19 @@ class Step3VisionEncoderLayer(nn.Module):

821 config: Step3VisionEncoderConfig,

822 quant_config: QuantizationConfig | None = None,

823 prefix: str = "",

824 - use_data_parallel: bool = False,

824 ):

825 super().__init__()

827 - self.use_data_parallel = use_data_parallel

826 self.embed_dim = config.hidden_size

827 self.self_attn = Step3VisionAttention(

828 config,

829 quant_config,

830 prefix=f"{prefix}.self_attn",

833 - use_data_parallel=self.use_data_parallel,

831 )

832 self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

833 self.mlp = Step3VisionMLP(

834 config,

835 quant_config,

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

840 - use_data_parallel=self.use_data_parallel,

837 )

838 self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

839

@@ -856,18 +852,15 @@ class Step3VisionEncoder(nn.Module):

852 config: Step3VisionEncoderConfig,

853 quant_config: QuantizationConfig | None = None,

854 prefix: str = "",

859 - use_data_parallel: bool = False,

855 ):

856 super().__init__()

857 self.config = config

863 - self.use_data_parallel = use_data_parallel

858 self.layers = nn.ModuleList(

859 [

860 Step3VisionEncoderLayer(

861 config,

862 quant_config,

863 prefix=f"{prefix}.layers.{i}",

870 - use_data_parallel=self.use_data_parallel,

864 )

865 for i in range(config.num_hidden_layers)

866 ]

@@ -889,18 +882,16 @@ class Step3VisionTransformer(nn.Module):

882 config: Step3VisionEncoderConfig,

883 quant_config: QuantizationConfig | None = None,

884 prefix: str = "",

892 - use_data_parallel: bool = False,

885 ):

886 super().__init__()

887 self.config = config

896 - self.use_data_parallel = use_data_parallel

888 + self.use_data_parallel = is_vit_use_data_parallel()

889 self.image_size = config.image_size

890 self.embeddings = Step3VisionEmbeddings(config)

891 self.transformer = Step3VisionEncoder(

892 config,

893 quant_config,

894 prefix=f"{prefix}.transformer",

903 - use_data_parallel=self.use_data_parallel,

895 )

896

897 def forward(

@@ -952,7 +943,6 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)

943 config.vision_config,

944 None,

945 prefix=maybe_prefix(prefix, "vision_model"),

955 - use_data_parallel=self.use_data_parallel,

946 )

947 self.vit_downsampler = Conv2dLayer(

948 config.vision_config.hidden_size,

vllm/model_executor/models/step_vl.py

+3 -11

@@ -24,7 +24,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig

25 from .step3_vl import Step3VLForConditionalGeneration

26 from .utils import WeightsMapper, init_vllm_registered_model, maybe_prefix

27 -from .vision import run_dp_sharded_vision_model

27 +from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model

29 _DEFAULT_NORM_LAYER = partial(nn.LayerNorm, eps=1e-5)

@@ -151,9 +151,9 @@ class PerceptionEncoderMLP(nn.Module):

151 act_layer: Callable[[], nn.Module],

152 quant_config: QuantizationConfig | None = None,

153 prefix: str = "",

154 - use_data_parallel: bool = False,

154 ):

155 super().__init__()

156 + use_data_parallel = is_vit_use_data_parallel()

157 self.fc1 = ColumnParallelLinear(

158 input_dim,

159 hidden_dim,

@@ -189,7 +189,6 @@ class PerceptionEncoderVisionAttention(nn.Module):

189 use_cls_token: bool = False,

190 quant_config: QuantizationConfig | None = None,

191 prefix: str = "",

192 - use_data_parallel: bool = False,

192 ):

193 super().__init__()

194 self.embed_dim = embed_dim

@@ -197,6 +196,7 @@ class PerceptionEncoderVisionAttention(nn.Module):

196 self.head_dim = embed_dim // num_heads

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

198

199 + use_data_parallel = is_vit_use_data_parallel()

200 tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()

201 assert self.total_num_heads % tp_size == 0, (

202 "embed_dim must be divisible by num_heads"

@@ -258,7 +258,6 @@ class PerceptionEncoderVisionBlock(nn.Module):

258 use_cls_token: bool = False,

259 quant_config: QuantizationConfig | None = None,

260 prefix: str = "",

261 - use_data_parallel: bool = False,

261 ):

262 super().__init__()

263 self.attn = PerceptionEncoderVisionAttention(

@@ -269,7 +268,6 @@ class PerceptionEncoderVisionBlock(nn.Module):

268 use_cls_token=use_cls_token,

269 quant_config=quant_config,

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

272 - use_data_parallel=use_data_parallel,

271 )

272 self.ls_1 = (

273 PerceptionEncoderLayerScale(d_model, ls_init_value)

@@ -290,7 +288,6 @@ class PerceptionEncoderVisionBlock(nn.Module):

288 act_layer,

289 quant_config=quant_config,

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

293 - use_data_parallel=use_data_parallel,

291 )

292

293 def forward(self, x: torch.Tensor, grid_hw: tuple[int, int]):

@@ -314,7 +311,6 @@ class PerceptionEncoderVisionTransformer(nn.Module):

311 use_cls_token: bool = False,

312 quant_config: QuantizationConfig | None = None,

313 prefix: str = "",

317 - use_data_parallel: bool = False,

314 ):

315 super().__init__()

316 self.width = width

@@ -333,7 +329,6 @@ class PerceptionEncoderVisionTransformer(nn.Module):

329 use_cls_token=use_cls_token,

330 quant_config=quant_config,

331 prefix=f"{prefix}.resblocks.{i}",

336 - use_data_parallel=use_data_parallel,

332 )

333 for i in range(layers)

334 ]

@@ -353,7 +348,6 @@ class PerceptionEncoder(nn.Module):

348 norm_layer: Callable = _DEFAULT_NORM_LAYER,

349 quant_config: QuantizationConfig | None = None,

350 prefix: str = "",

356 - use_data_parallel: bool = False,

351 ):

352 super().__init__()

353 self.patch_size = config.patch_size

@@ -394,7 +388,6 @@ class PerceptionEncoder(nn.Module):

388 use_cls_token=self.use_cls_token,

389 quant_config=quant_config,

390 prefix=f"{prefix}.transformer",

397 - use_data_parallel=use_data_parallel,

391 )

392

393 self.vit_downsampler1 = Conv2dLayer(

@@ -511,7 +504,6 @@ class StepVLForConditionalGeneration(Step3VLForConditionalGeneration):

504 get_act_fn(config.vision_config.hidden_act),

505 quant_config=quant_config,

506 prefix=maybe_prefix(prefix, "vision_model"),

514 - use_data_parallel=self.use_data_parallel,

507 )

508 self.vit_large_projector = ColumnParallelLinear(

509 config.vision_config.width * 4,