@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 32d2b4064f

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] Add Ovis2.5 PP support (#23405)

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

Isotr0py committed Aug 23, 2025 at 01:46 UTC 32d2b4064feea38802489b71e47703d1f901a17e

5 files changed +185 -105

tests/distributed/test_pipeline_parallel.py

@@ -233,6 +233,7 @@ MULTIMODAL_MODELS = {

233 "openbmb/MiniCPM-Llama3-V-2_5": PPTestSettings.fast(),

234 "allenai/Molmo-7B-D-0924": PPTestSettings.fast(),

235 "AIDC-AI/Ovis2-1B": PPTestSettings.fast(),

236 + "AIDC-AI/Ovis2.5-2B": PPTestSettings.fast(),

237 "microsoft/Phi-3.5-vision-instruct": PPTestSettings.fast(),

238 "mistralai/Pixtral-12B-2409": PPTestSettings.fast(load_format="dummy"),

239 "Qwen/Qwen-VL-Chat": PPTestSettings.fast(),

tests/models/multimodal/generation/test_common.py

+1 -5

@@ -11,7 +11,6 @@ from pathlib import PosixPath

11 import pytest

12 from transformers import (AutoModel, AutoModelForImageTextToText,

13 AutoModelForTextToWaveform, AutoModelForVision2Seq)

14 -from transformers.utils import is_flash_attn_2_available

15 from vllm.platforms import current_platform

16 from vllm.utils import identity

@@ -637,10 +636,7 @@ VLM_TEST_SETTINGS = {

636 dtype="half",

637 num_logprobs=10,

638 patch_hf_runner=model_utils.ovis2_5_patch_hf_runner,

640 - marks=[pytest.mark.skipif(

641 - not is_flash_attn_2_available(),

642 - reason="HF model needs `flash_attn` installed"

643 - )],

639 + hf_model_kwargs={"revision": "refs/pr/5"},

640 ),

641 "phi3v": VLMTestInfo(

642 models=["microsoft/Phi-3.5-vision-instruct"],

tests/models/registry.py

+1 -3

@@ -468,9 +468,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {

468 extras={"1.6-llama": "AIDC-AI/Ovis1.6-Llama3.2-3B",

469 "1.6-gemma": "AIDC-AI/Ovis1.6-Gemma2-9B"}), # noqa: E501

470 "Ovis2_5": _HfExamplesInfo("AIDC-AI/Ovis2.5-2B",

471 - trust_remote_code=True,

472 - max_transformers_version="4.53",

473 - transformers_version_reason="HF model is not compatible"), # noqa: E501

471 + trust_remote_code=True),

472 "PaliGemmaForConditionalGeneration": _HfExamplesInfo("google/paligemma-3b-mix-224", # noqa: E501

473 extras={"v2": "google/paligemma2-3b-ft-docci-448"}), # noqa: E501

474 "Phi3VForCausalLM": _HfExamplesInfo("microsoft/Phi-3-vision-128k-instruct",

vllm/model_executor/models/ovis2_5.py

+20 -16

@@ -30,7 +30,7 @@ from vllm.multimodal.profiling import BaseDummyInputsBuilder

30 from vllm.sequence import IntermediateTensors

31 from vllm.transformers_utils.processors.ovis2_5 import Ovis2_5Processor

33 -from .interfaces import MultiModalEmbeddings, SupportsMultiModal

33 +from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

35 IMAGE_TOKEN = "<image>"

36 VIDEO_TOKEN = "<video>"

@@ -70,6 +70,7 @@ class VisualTokenizer(torch.nn.Module):

70 visual_vocab_size: int,

71 quant_config: Optional[QuantizationConfig] = None,

72 prefix: str = "",

73 + use_data_parallel: bool = False,

74 ):

75 super().__init__()

76 self.config = config

@@ -77,6 +78,7 @@ class VisualTokenizer(torch.nn.Module):

78 config=config,

79 quant_config=quant_config,

80 prefix=f"{prefix}.vit",

81 + use_data_parallel=use_data_parallel,

82 )

83 # reserved tokens for INDICATOR_IDS

84 head_dim = visual_vocab_size - len(INDICATOR_IDS)

@@ -93,31 +95,33 @@ class VisualTokenizer(torch.nn.Module):

95 config: PretrainedConfig,

96 quant_config: Optional[QuantizationConfig] = None,

97 prefix: str = "",

98 + use_data_parallel: bool = False,

99 ):

100 model_type = config.model_type

101 if model_type == "siglip2_navit":

99 - return Siglip2NavitModel(config=config, )

102 + return Siglip2NavitModel(config=config,

103 + quant_config=quant_config,

104 + prefix=prefix,

105 + use_data_parallel=use_data_parallel)

106 raise ValueError(

107 f"Unsupported visual tokenizer model_type: {model_type}")

108

109 @property

104 - def dtype(self):

110 + def dtype(self) -> torch.dtype:

111 return next(self.head.parameters()).dtype

112

113 @property

108 - def device(self):

114 + def device(self) -> torch.device:

115 return next(self.head.parameters()).device

116

111 - def tokenize(self, logits):

117 + def tokenize(self, logits: torch.Tensor) -> torch.Tensor:

118 tokens = torch.softmax(logits, dim=-1,

119 dtype=torch.float32).to(logits.dtype)

120 return tokens

121

116 - def encode(self, pixel_values, grid_thws):

117 - features = self.vit(pixel_values,

118 - grid_thws,

119 - output_hidden_states=True,

120 - return_dict=True)

122 + def encode(self, pixel_values: torch.Tensor,

123 + grid_thws: torch.Tensor) -> torch.Tensor:

124 + features = self.vit(pixel_values, grid_thws)

125 # refer to qwen2.5-vl patchmerger

126 seq_len, _ = features.shape

127 features = features.reshape(seq_len // (self.config.hidden_stride**2),

@@ -125,7 +129,8 @@ class VisualTokenizer(torch.nn.Module):

129

130 return features

131

128 - def forward(self, pixel_values, grid_thws) -> torch.Tensor:

132 + def forward(self, pixel_values: torch.Tensor,

133 + grid_thws: torch.Tensor) -> torch.Tensor:

134 features = self.encode(pixel_values, grid_thws)

135 logits = self.head(features)

136 tokens = self.tokenize(logits)

@@ -395,7 +400,7 @@ class Ovis2_5MultiModalProcessor(BaseMultiModalProcessor[Ovis2_5ProcessingInfo]

400 @MULTIMODAL_REGISTRY.register_processor(Ovis2_5MultiModalProcessor,

401 info=Ovis2_5ProcessingInfo,

402 dummy_inputs=Ovis2_5DummyInputsBuilder)

398 -class Ovis2_5(nn.Module, SupportsMultiModal):

403 +class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

404

405 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

406 super().__init__()

@@ -421,9 +426,8 @@ class Ovis2_5(nn.Module, SupportsMultiModal):

426 text_model_type = self.config.get_text_config().model_type

427 self.image_pad_token_id = IMAGE_PAD_TOKEN_ID_MAP[text_model_type]

428

424 - # TODO(Isotr0py): PP support

425 - # self.make_empty_intermediate_tensors = (

426 - # self.language_model.make_empty_intermediate_tensors)

429 + self.make_empty_intermediate_tensors = (

430 + self.get_language_model().make_empty_intermediate_tensors)

431

432 def _parse_and_validate_visual_input(

433 self, is_video,

@@ -567,4 +571,4 @@ class Ovis2_5(nn.Module, SupportsMultiModal):

571 return loader.load_weights(weights)

572

573 def get_language_model(self) -> torch.nn.Module:

570 - return self.llm

\ No newline at end of file

574 + return self.llm

vllm/model_executor/models/siglip2navit.py

+162 -81

@@ -3,16 +3,24 @@

3 """Implementation of SiglipVisionModel intended to be only used

4 within a vision language model."""

6 -from typing import Optional, Union

6 +from collections.abc import Iterable

7 +from typing import Optional

9 import torch

10 from einops import rearrange, repeat

11 from torch import nn

12 from torch.nn import functional as F

12 -from transformers.activations import ACT2FN

13 +from transformers import Siglip2VisionConfig

14 from transformers.configuration_utils import PretrainedConfig

14 -from transformers.modeling_outputs import BaseModelOutputWithNoAttention

16 +from vllm.config import QuantizationConfig

17 +from vllm.distributed import divide, get_tensor_model_parallel_world_size

18 +from vllm.model_executor.layers.activation import get_act_fn

19 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,

20 + LinearBase, QKVParallelLinear,

21 + ReplicatedLinear,

22 + RowParallelLinear)

23 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

24 from vllm.platforms import _Backend

26 from .vision import get_vit_attn_backend

@@ -48,10 +56,11 @@ class Siglip2VisionEmbeddings(nn.Module):

57 # siglip2 naflex

58 if self.num_patches > 0:

51 - self.patch_embedding = nn.Linear(

52 - in_features=config.num_channels * self.patch_size *

59 + self.patch_embedding = ReplicatedLinear(

60 + input_size=config.num_channels * self.patch_size *

61 self.patch_size,

54 - out_features=self.embed_dim,

62 + output_size=self.embed_dim,

63 + return_bias=False,

64 )

65 if self.preserve_original_pe:

66 self.position_embedding_size = int(self.num_patches**0.5)

@@ -89,7 +98,7 @@ class Siglip2VisionEmbeddings(nn.Module):

99 # Apply patch embeddings to already patchified pixel values

100 target_dtype = self.patch_embedding.weight.dtype

92 - if isinstance(self.patch_embedding, nn.Linear):

101 + if isinstance(self.patch_embedding, LinearBase):

102 patch_embeds = self.patch_embedding(

103 pixel_values.to(dtype=target_dtype))

104 elif isinstance(self.patch_embedding, nn.Conv2d):

@@ -184,7 +193,13 @@ def apply_rotary_pos_emb(

193 class Siglip2Attention(nn.Module):

194 """Multi-headed attention from 'Attention Is All You Need' paper"""

195

187 - def __init__(self, config):

196 + def __init__(

197 + self,

198 + config: Siglip2VisionConfig,

199 + quant_config: Optional[QuantizationConfig] = None,

200 + prefix: str = "",

201 + use_data_parallel: bool = False,

202 + ):

203 super().__init__()

204 self.config = config

205 self.embed_dim = config.hidden_size

@@ -199,11 +214,25 @@ class Siglip2Attention(nn.Module):

214 self.dropout = config.attention_dropout

215 self.is_causal = False

216

202 - self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)

203 - self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)

204 - self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)

205 - self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)

217 + # TODO(Isotr0py): Enable data parallel after we support

218 + # disabling TP on parallel linear layer

219 + self.qkv_proj = QKVParallelLinear(

220 + hidden_size=self.embed_dim,

221 + head_size=self.head_dim,

222 + total_num_heads=self.num_heads,

223 + quant_config=quant_config,

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

225 + )

226 + self.out_proj = RowParallelLinear(

227 + input_size=self.embed_dim,

228 + output_size=self.embed_dim,

229 + quant_config=quant_config,

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

231 + )

232

233 + self.tp_size = (1 if use_data_parallel else

234 + get_tensor_model_parallel_world_size())

235 + self.num_heads_per_partition = divide(self.num_heads, self.tp_size)

236 self.use_rope = config.use_rope

237

238 # Detect attention implementation.

@@ -228,13 +257,15 @@ class Siglip2Attention(nn.Module):

257

258 seq_length, embed_dim = hidden_states.shape

259

231 - queries = self.q_proj(hidden_states)

232 - keys = self.k_proj(hidden_states)

233 - values = self.v_proj(hidden_states)

260 + qkv_states, _ = self.qkv_proj(hidden_states)

261 + queries, keys, values = qkv_states.chunk(3, dim=-1)

262

235 - queries = queries.view(seq_length, self.num_heads, self.head_dim)

236 - keys = keys.view(seq_length, self.num_heads, self.head_dim)

237 - values = values.view(seq_length, self.num_heads, self.head_dim)

263 + queries = queries.view(seq_length, self.num_heads_per_partition,

264 + self.head_dim)

265 + keys = keys.view(seq_length, self.num_heads_per_partition,

266 + self.head_dim)

267 + values = values.view(seq_length, self.num_heads_per_partition,

268 + self.head_dim)

269

270 if self.use_rope:

271 cos, sin = position_embeddings

@@ -276,41 +307,72 @@ class Siglip2Attention(nn.Module):

307 v_i,

308 dropout_p=0.0)

309 # (1, num_heads, seq_len, head_dim) -> (seq_len, embed_dim)

279 - output_i = output_i.transpose(1, 2).reshape(-1, self.embed_dim)

310 + output_i = output_i.transpose(1, 2).reshape(

311 + end_idx - start_idx, -1)

312 outputs.append(output_i)

313

314 attn_output = torch.cat(outputs, dim=0)

283 - attn_output = self.out_proj(attn_output)

315 + attn_output, _ = self.out_proj(attn_output)

316 return attn_output

317

318

319 class Siglip2MLP(nn.Module):

320

289 - def __init__(self, config):

321 + def __init__(

322 + self,

323 + config: Siglip2VisionConfig,

324 + quant_config: Optional[QuantizationConfig] = None,

325 + prefix: str = "",

326 + use_data_parallel: bool = False,

327 + ):

328 super().__init__()

329 self.config = config

292 - self.activation_fn = ACT2FN[config.hidden_act]

293 - self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)

294 - self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)

330 + self.activation_fn = get_act_fn(config.hidden_act)

331 + # TODO(Isotr0py): Enable data parallel after we support

332 + # disabling TP on parallel linear layer

333 + self.fc1 = ColumnParallelLinear(

334 + config.hidden_size,

335 + config.intermediate_size,

336 + quant_config=quant_config,

337 + prefix=f"{prefix}.fc1",

338 + )

339 + self.fc2 = RowParallelLinear(

340 + config.intermediate_size,

341 + config.hidden_size,

342 + quant_config=quant_config,

343 + prefix=f"{prefix}.fc2",

344 + )

345

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

297 - hidden_states = self.fc1(hidden_states)

347 + hidden_states, _ = self.fc1(hidden_states)

348 hidden_states = self.activation_fn(hidden_states)

299 - hidden_states = self.fc2(hidden_states)

349 + hidden_states, _ = self.fc2(hidden_states)

350 return hidden_states

351

352

353 class Siglip2EncoderLayer(nn.Module):

354

305 - def __init__(self, config: PretrainedConfig):

355 + def __init__(

356 + self,

357 + config: Siglip2VisionConfig,

358 + quant_config: Optional[QuantizationConfig] = None,

359 + prefix: str = "",

360 + use_data_parallel: bool = False,

361 + ):

362 super().__init__()

363 self.embed_dim = config.hidden_size

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

365 eps=config.layer_norm_eps)

310 - self.self_attn = Siglip2Attention(config)

366 + self.self_attn = Siglip2Attention(config,

367 + quant_config=quant_config,

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

369 + use_data_parallel=use_data_parallel)

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

371 eps=config.layer_norm_eps)

313 - self.mlp = Siglip2MLP(config)

372 + self.mlp = Siglip2MLP(config,

373 + quant_config=quant_config,

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

375 + use_data_parallel=use_data_parallel)

376

377 def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor,

378 position_embeddings: torch.Tensor) -> tuple[torch.FloatTensor]:

@@ -347,14 +409,22 @@ class Siglip2Encoder(nn.Module):

409 config: PretrainedConfig

410 """

411

350 - def __init__(self, config: PretrainedConfig):

412 + def __init__(

413 + self,

414 + config: Siglip2VisionConfig,

415 + quant_config: Optional[QuantizationConfig] = None,

416 + prefix: str = "",

417 + use_data_parallel: bool = False,

418 + ):

419 super().__init__()

420 self.config = config

421 self.layers = nn.ModuleList([

354 - Siglip2EncoderLayer(config)

355 - for _ in range(config.num_hidden_layers)

422 + Siglip2EncoderLayer(config,

423 + quant_config=quant_config,

424 + prefix=f"{prefix}.layers.{idx}",

425 + use_data_parallel=use_data_parallel)

426 + for idx in range(config.num_hidden_layers)

427 ])

357 - self.gradient_checkpointing = False

428

429 self.rotary_pos_emb = VisionRotaryEmbedding(

430 config.hidden_size // config.num_attention_heads // 2)

@@ -445,13 +515,11 @@ class Siglip2Encoder(nn.Module):

515

516 return window_index, cu_window_seqlens

517

448 - # Ignore copy

518 def forward(

519 self,

451 - inputs_embeds,

520 + inputs_embeds: torch.Tensor,

521 grid_thws: torch.Tensor,

453 - output_hidden_states: bool = False,

454 - ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, ...]]]:

522 + ) -> torch.Tensor:

523 r"""

524 Args:

525 inputs_embeds (`torch.FloatTensor` of shape

@@ -506,7 +574,6 @@ class Siglip2Encoder(nn.Module):

574 cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)

575

576 reverse_indices = torch.argsort(window_index)

509 - encoder_states = () if output_hidden_states else None

577

578 hidden_states = inputs_embeds

579 for index, block in enumerate(self.layers):

@@ -517,45 +584,40 @@ class Siglip2Encoder(nn.Module):

584 cu_seqlens_tmp = cu_window_seqlens

585 hidden_states = block(hidden_states, cu_seqlens_tmp,

586 position_embeddings)

520 - if output_hidden_states:

521 - hidden_states_ = hidden_states.reshape(

522 - seq_len // self.spatial_merge_unit,

523 - self.spatial_merge_unit, -1)

524 - encoder_states += (hidden_states_[reverse_indices, :].reshape(

525 - seq_len, -1), )

526 - # tokens = self.post_trunk_norm(tokens)

587 +

588 hidden_states = hidden_states.reshape(

589 seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)

590 hidden_states = hidden_states[reverse_indices, :].reshape(seq_len, -1)

591

531 - return hidden_states, encoder_states

592 + return hidden_states

593

594

595 class Siglip2VisionTransformer(nn.Module):

596

536 - def __init__(self, config: PretrainedConfig):

597 + def __init__(

598 + self,

599 + config: Siglip2VisionConfig,

600 + quant_config: Optional[QuantizationConfig] = None,

601 + prefix: str = "",

602 + use_data_parallel: bool = False,

603 + ):

604 super().__init__()

605 self.config = config

606 embed_dim = config.hidden_size

607

608 self.embeddings = Siglip2VisionEmbeddings(config)

542 - self.encoder = Siglip2Encoder(config)

609 + self.encoder = Siglip2Encoder(config,

610 + quant_config=quant_config,

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

612 + use_data_parallel=use_data_parallel)

613 self.post_layernorm = nn.LayerNorm(embed_dim,

614 eps=config.layer_norm_eps)

545 - self._use_flash_attention_2 = \

546 - (config._attn_implementation == "flash_attention_2")

615

616 def forward(

617 self,

618 pixel_values: torch.FloatTensor,

619 grid_thws: torch.LongTensor,

552 - output_hidden_states: Optional[bool] = True,

553 - return_dict: Optional[bool] = True,

554 - ) -> Union[

555 - tuple[torch.Tensor],

556 - tuple[torch.Tensor, tuple[torch.Tensor, ...]],

557 - BaseModelOutputWithNoAttention,

558 - ]:

620 + ) -> torch.Tensor:

621 r"""

622 spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):

623 Tensor containing the spatial dimensions (height, width)

@@ -563,45 +625,64 @@ class Siglip2VisionTransformer(nn.Module):

625 """

626 hidden_states = self.embeddings(pixel_values, grid_thws)

627

566 - last_hidden_state, hidden_states = self.encoder(

567 - hidden_states, grid_thws, output_hidden_states)

628 + last_hidden_state = self.encoder(hidden_states, grid_thws)

629 last_hidden_state = self.post_layernorm(last_hidden_state)

630

570 - if not return_dict:

571 - output = (last_hidden_state, )

572 - output += (hidden_states, ) if output_hidden_states else ()

573 - return output

574 -

631 return last_hidden_state

632

633

634 class Siglip2NavitModel(torch.nn.Module):

635

580 - def __init__(self, config: PretrainedConfig):

636 + def __init__(

637 + self,

638 + config: Siglip2VisionConfig,

639 + quant_config: Optional[QuantizationConfig] = None,

640 + prefix: str = "",

641 + use_data_parallel: bool = False,

642 + ):

643 super().__init__()

644

583 - self.vision_model = Siglip2VisionTransformer(config)

645 + self.vision_model = Siglip2VisionTransformer(

646 + config,

647 + quant_config=quant_config,

648 + prefix=f"{prefix}.vision_model",

649 + use_data_parallel=use_data_parallel)

650

651 def forward(

652 self,

653 pixel_values: torch.FloatTensor,

654 grid_thws: torch.LongTensor,

589 - output_hidden_states: Optional[bool] = None,

590 - return_dict: Optional[bool] = None,

591 - ) -> Union[

592 - tuple[torch.Tensor],

593 - tuple[torch.Tensor, tuple[torch.Tensor, ...]],

594 - BaseModelOutputWithNoAttention,

595 - ]:

596 -

597 - if output_hidden_states is None:

598 - output_hidden_states = self.config.output_hidden_states

599 - if return_dict is None:

600 - return_dict = self.config.use_return_dict

601 -

655 + ) -> torch.Tensor:

656 return self.vision_model(

657 pixel_values=pixel_values,

658 grid_thws=grid_thws,

605 - output_hidden_states=output_hidden_states,

606 - return_dict=return_dict,

659 )

660 +

661 + def load_weights(self, weights: Iterable[tuple[str,

662 + torch.Tensor]]) -> set[str]:

663 + stacked_params_mapping = [

664 + # (param_name, shard_name, shard_id)

665 + ("qkv_proj", "q_proj", "q"),

666 + ("qkv_proj", "k_proj", "k"),

667 + ("qkv_proj", "v_proj", "v"),

668 + ]

669 + params_dict = dict(self.named_parameters())

670 + loaded_params: set[str] = set()

671 +

672 + for name, loaded_weight in weights:

673 + for (param_name, weight_name, shard_id) in stacked_params_mapping:

674 + if weight_name not in name:

675 + continue

676 + name = name.replace(weight_name, param_name)

677 +

678 + param = params_dict[name]

679 + weight_loader = param.weight_loader

680 + weight_loader(param, loaded_weight, shard_id)

681 + break

682 + else:

683 + param = params_dict[name]

684 + weight_loader = getattr(param, "weight_loader",

685 + default_weight_loader)

686 + weight_loader(param, loaded_weight)

687 + loaded_params.add(name)

688 + return loaded_params