@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 644d57d531

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 Ernie4.5 VL Model Support (#22514)

Signed-off-by: wangyafeng <wangyafeng@baidu.com>

CSWYF3634076 committed Aug 27, 2025 at 12:02 UTC 644d57d53191b94d9e50a4765891c498790d924b

11 files changed +2463

docs/models/supported_models.md

@@ -616,6 +616,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

616 | `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I+ | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ | ✅︎ |

617 | `DeepseekVLV2ForCausalLM`^ | DeepSeek-VL2 | T + I+ | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ | ✅︎ |

619 +| `Ernie4_5_VLMoeForConditionalGeneration` | Ernie4.5-VL | T + I+/ V+ | `baidu/ERNIE-4.5-VL-28B-A3B-PT`, `baidu/ERNIE-4.5-VL-424B-A47B-PT` | | ✅︎ | ✅︎ |

621 | `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ | ✅︎ |

622 | `Gemma3ForConditionalGeneration` | Gemma 3 | T + I+ | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ | ⚠️ |

examples/offline_inference/vision_language.py

+32

@@ -173,6 +173,37 @@ def run_deepseek_vl2(questions: list[str], modality: str) -> ModelRequestData:

173 )

174

175

176 +# Ernie4.5-VL

177 +def run_ernie45_vl(questions: list[str], modality: str) -> ModelRequestData:

178 + model_name = "baidu/ERNIE-4.5-VL-28B-A3B-PT"

179 +

180 + engine_args = EngineArgs(

181 + model=model_name,

182 + max_model_len=4096,

183 + max_num_seqs=5,

184 + limit_mm_per_prompt={modality: 1},

185 + trust_remote_code=True,

186 + )

187 +

188 + if modality == "image":

189 + placeholder = "Picture 1:<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>"

190 + elif modality == "video":

191 + placeholder = "Video 1:<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>"

192 +

193 + prompts = [

194 + (

195 + f"<|begin_of_sentence|>User: {question}{placeholder}\n"

196 + "Assistant: <think></think>"

197 + )

198 + for question in questions

199 + ]

200 +

201 + return ModelRequestData(

202 + engine_args=engine_args,

203 + prompts=prompts,

204 + )

205 +

206 +

207 # Florence2

208 def run_florence2(questions: list[str], modality: str) -> ModelRequestData:

209 assert modality == "image"

@@ -1602,6 +1633,7 @@ model_example_map = {

1633 "chameleon": run_chameleon,

1634 "command_a_vision": run_command_a_vision,

1635 "deepseek_vl_v2": run_deepseek_vl2,

1636 + "ernie45_vl": run_ernie45_vl,

1637 "florence2": run_florence2,

1638 "fuyu": run_fuyu,

1639 "gemma3": run_gemma3,

requirements/test.in

@@ -54,3 +54,4 @@ runai-model-streamer-s3==0.11.0

54 fastsafetensors>=0.1.10

55 pydantic>=2.10 # 2.9 leads to error on python 3.10

56 terratorch==1.1rc2 # required for PrithviMAE test

57 +decord==0.6.0

requirements/test.txt

@@ -156,6 +156,8 @@ datasets==3.0.2

156 # mteb

157 decorator==5.1.1

158 # via librosa

159 +decord==0.6.0

160 + # via -r requirements/test.in

161 dill==0.3.8

162 # via

163 # datasets

@@ -493,6 +495,7 @@ numpy==1.26.4

495 # contourpy

496 # cupy-cuda12x

497 # datasets

498 + # decord

499 # einx

500 # encodec

501 # evaluate

tests/models/multimodal/processing/test_common.py

@@ -272,6 +272,7 @@ def _test_processing_correctness_one(

272 "CohereLabs/command-a-vision-07-2025",

273 "deepseek-ai/deepseek-vl2-tiny",

274 "naver-clova-ix/donut-base-finetuned-docvqa",

275 + "baidu/ERNIE-4.5-VL-28B-A3B-PT",

276 "microsoft/Florence-2-base",

277 "adept/fuyu-8b",

278 "google/gemma-3-4b-it",

tests/models/registry.py

@@ -396,6 +396,8 @@ _MULTIMODAL_EXAMPLE_MODELS = {

396 transformers_version_reason="HF model is not compatible.", # noqa: E501

397 hf_overrides={"architectures": ["DeepseekVLV2ForCausalLM"]}), # noqa: E501

398 "Emu3ForConditionalGeneration": _HfExamplesInfo("BAAI/Emu3-Chat-hf"),

399 + "Ernie4_5_VLMoeForConditionalGeneration": _HfExamplesInfo("baidu/ERNIE-4.5-VL-28B-A3B-PT", # noqa: E501

400 + trust_remote_code=True),

401 "FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"),

402 "Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it"),

403 "Gemma3nForConditionalGeneration": _HfExamplesInfo("google/gemma-3n-E2B-it", # noqa: E501

vllm/model_executor/layers/rotary_embedding/ernie45_vl_rope.py new

+72

@@ -0,0 +1,72 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from typing import Optional

5 +

6 +import torch

7 +

8 +from .common import apply_rotary_emb_dispatch

9 +from .mrope import MRotaryEmbedding

10 +

11 +

12 +class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):

13 + """3D rotary positional embedding. 3D is t:time h:height w:width"""

14 +

15 + def forward(

16 + self,

17 + positions: torch.Tensor,

18 + query: torch.Tensor,

19 + key: Optional[torch.Tensor] = None,

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

21 + assert positions.ndim == 1 or positions.ndim == 2

22 + assert key is not None

23 +

24 + num_tokens = positions.shape[-1]

25 + cos_sin = self.cos_sin_cache[positions]

26 + cos, sin = cos_sin.chunk(2, dim=-1)

27 + if positions.ndim == 2:

28 + assert self.mrope_section

29 +

30 + section_h = self.mrope_section[0] # 22

31 + section_w = self.mrope_section[1] # 22

32 + section_t = self.mrope_section[2] # 20

33 + assert section_h == section_w

34 + # Split according to [h w h w h w h w... t t t...]

35 + section_cos_t = cos[..., -section_t:]

36 + section_cos_h = cos[..., :section_h + section_w:2]

37 + section_cos_w = cos[..., 1:section_h + section_w:2]

38 +

39 + cos_t, cos_h, cos_w = section_cos_t[0], section_cos_h[

40 + 1], section_cos_w[2]

41 + cos_hw = torch.stack([cos_h, cos_w],

42 + dim=-1).reshape(cos_h.shape[:-1] +

43 + (cos_h.shape[-1] * 2, ))

44 + cos = torch.cat([cos_hw, cos_t], dim=-1)

45 +

46 + section_sin_t = sin[..., -section_t:]

47 + section_sin_h = sin[..., :section_h + section_w:2]

48 + section_sin_w = sin[..., 1:section_h + section_w:2]

49 +

50 + sin_t, sin_h, sin_w = section_sin_t[0], section_sin_h[

51 + 1], section_sin_w[2]

52 + sin_hw = torch.stack([sin_h, sin_w],

53 + dim=-1).reshape(sin_h.shape[:-1] +

54 + (sin_h.shape[-1] * 2, ))

55 + sin = torch.cat([sin_hw, sin_t], dim=-1)

56 +

57 + query_shape = query.shape

58 + query = query.view(num_tokens, -1, self.head_size)

59 + query_rot = query[..., :self.rotary_dim]

60 + query_pass = query[..., self.rotary_dim:]

61 + query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin,

62 + self.is_neox_style)

63 + query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)

64 +

65 + key_shape = key.shape

66 + key = key.view(num_tokens, -1, self.head_size)

67 + key_rot = key[..., :self.rotary_dim]

68 + key_pass = key[..., self.rotary_dim:]

69 + key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin,

70 + self.is_neox_style)

71 + key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)

72 + return query, key

vllm/model_executor/layers/rotary_embedding/mrope.py

+123

@@ -393,6 +393,15 @@ class MRotaryEmbedding(RotaryEmbedding):

393 context_len=context_len,

394 seq_len=seq_len,

395 )

396 + elif hf_config.model_type in ["ernie4_5_moe_vl", "ernie4_5_vl"]:

397 + return cls._ernie_get_input_positions_tensor(

398 + input_tokens=input_tokens,

399 + hf_config=hf_config,

400 + image_grid_thw=image_grid_thw,

401 + video_grid_thw=video_grid_thw,

402 + context_len=context_len,

403 + seq_len=seq_len,

404 + )

405 else:

406 return cls._vl_get_input_positions_tensor(

407 input_tokens=input_tokens,

@@ -513,6 +522,120 @@ class MRotaryEmbedding(RotaryEmbedding):

522 len(input_tokens)).item()

523 return llm_positions, mrope_position_delta

524

525 + @classmethod

526 + def _ernie_get_input_positions_tensor(

527 + cls,

528 + input_tokens: list[int],

529 + hf_config: PretrainedConfig,

530 + image_grid_thw: Union[list[list[int]], torch.Tensor],

531 + video_grid_thw: Union[list[list[int]], torch.Tensor],

532 + context_len: int = 0,

533 + seq_len: Optional[int] = None,

534 + ) -> tuple[torch.Tensor, int]:

535 + """Get mrope input positions and delta value for Ernie VL."""

536 +

537 + image_token_id = hf_config.im_patch_id

538 + video_start_token_id = hf_config.video_start_token_id

539 + video_end_token_id = hf_config.video_end_token_id

540 + spatial_conv_size = hf_config.spatial_conv_size

541 + temporal_conv_size = hf_config.temporal_conv_size

542 + llm_pos_ids_list: list = []

543 +

544 + if not (image_grid_thw is None and video_grid_thw is None):

545 + if isinstance(image_grid_thw, torch.Tensor):

546 + image_grid_thw = image_grid_thw.tolist()

547 +

548 + input_token_type: list[str] = []

549 + video_check_flg = False

550 + for token in input_tokens:

551 + if token == video_start_token_id:

552 + video_check_flg = True

553 + elif token == video_end_token_id:

554 + video_check_flg = False

555 +

556 + if (token == image_token_id) and (video_check_flg is False):

557 + input_token_type.append("image")

558 + elif (token == image_token_id) and (video_check_flg is True):

559 + input_token_type.append("video")

560 + else:

561 + input_token_type.append("text")

562 +

563 + input_type_group: list[tuple[str, int, int]] = []

564 + for key, group_iter in itertools.groupby(

565 + enumerate(input_token_type), lambda x: x[1]):

566 + group_list = list(group_iter)

567 + start_index = group_list[0][0]

568 + end_index = group_list[-1][0] + 1

569 + input_type_group.append((key, start_index, end_index))

570 +

571 + video_frame_num = 1

572 + mm_data_idx = 0

573 + for modality_type, start_idx, end_idx in input_type_group:

574 + st_idx = llm_pos_ids_list[-1].max() + 1 if len(

575 + llm_pos_ids_list) > 0 else 0

576 + if modality_type == "image":

577 + t, h, w = (

578 + image_grid_thw[mm_data_idx][0],

579 + image_grid_thw[mm_data_idx][1],

580 + image_grid_thw[mm_data_idx][2],

581 + )

582 + llm_grid_t, llm_grid_h, llm_grid_w = \

583 + t, h // spatial_conv_size, w // spatial_conv_size

584 +

585 + t_index = torch.arange(llm_grid_t).view(-1, 1).expand(

586 + -1, llm_grid_h * llm_grid_w).flatten()

587 + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(

588 + llm_grid_t, -1, llm_grid_w).flatten()

589 + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(

590 + llm_grid_t, llm_grid_h, -1).flatten()

591 + llm_pos_ids_list.append(

592 + torch.stack([t_index, h_index, w_index]) + st_idx)

593 + mm_data_idx += 1

594 +

595 + elif modality_type == "video":

596 + t, h, w = (

597 + video_grid_thw[mm_data_idx][0],

598 + video_grid_thw[mm_data_idx][1],

599 + video_grid_thw[mm_data_idx][2],

600 + )

601 + llm_grid_t, llm_grid_h, llm_grid_w = (t //

602 + temporal_conv_size,

603 + h //

604 + spatial_conv_size,

605 + w //

606 + spatial_conv_size)

607 +

608 + for t_idx in range(llm_grid_t):

609 + t_index = torch.tensor(t_idx).view(-1, 1).expand(

610 + -1, llm_grid_h * llm_grid_w).flatten()

611 + h_index = torch.arange(llm_grid_h).view(

612 + 1, -1, 1).expand(1, -1, llm_grid_w).flatten()

613 + w_index = torch.arange(llm_grid_w).view(

614 + 1, 1, -1).expand(1, llm_grid_h, -1).flatten()

615 + llm_pos_ids_list.append(

616 + torch.stack([t_index, h_index, w_index]) + st_idx)

617 +

618 + mm_data_idx += 1

619 + video_frame_num += 1

620 +

621 + else:

622 + text_len = end_idx - start_idx

623 + llm_pos_ids_list.append(

624 + torch.arange(text_len).view(1, -1).expand(3, -1) +

625 + st_idx)

626 + video_frame_num = 1

627 +

628 + else:

629 + text_len = len(input_tokens)

630 + llm_pos_ids_list.append(

631 + torch.arange(text_len).view(1, -1).expand(3, -1))

632 +

633 + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)

634 + llm_positions = llm_positions[:, context_len:seq_len]

635 + mrope_position_delta = (llm_positions.max() + 1 -

636 + len(input_tokens)).item()

637 + return llm_positions, mrope_position_delta

638 +

639 @classmethod

640 def _vl_get_input_positions_tensor(

641 cls,

vllm/model_executor/models/ernie45_vl.py new

+1504

@@ -0,0 +1,1504 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

7 +#

8 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX

9 +# and OPT implementations in this library. It has been modified from its

10 +# original forms to accommodate minor architectural differences compared

11 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.

12 +#

13 +# Licensed under the Apache License, Version 2.0 (the "License");

14 +# you may not use this file except in compliance with the License.

15 +# You may obtain a copy of the License at

16 +#

17 +# http://www.apache.org/licenses/LICENSE-2.0

18 +#

19 +# Unless required by applicable law or agreed to in writing, software

20 +# distributed under the License is distributed on an "AS IS" BASIS,

21 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

22 +# See the License for the specific language governing permissions and

23 +# limitations under the License.

24 +"""Inference-only Erine VL model compatible with HuggingFace weights."""

25 +import math

26 +from collections.abc import Iterable, Mapping, Sequence

27 +from functools import partial

28 +from typing import Any, Callable, Literal, Optional, TypedDict, Union

29 +

30 +import numpy as np

31 +import torch

32 +import torch.nn as nn

33 +import torch.nn.functional as F

34 +from einops import rearrange, repeat

35 +from transformers import BatchFeature

36 +

37 +from vllm.config import VllmConfig

38 +from vllm.distributed import parallel_state

39 +from vllm.distributed import utils as dist_utils

40 +from vllm.logger import init_logger

41 +from vllm.model_executor import SamplingMetadata

42 +from vllm.model_executor.layers.activation import QuickGELU

43 +from vllm.model_executor.layers.layernorm import RMSNorm

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

45 + QKVParallelLinear,

46 + RowParallelLinear)

47 +from vllm.model_executor.layers.quantization import QuantizationConfig

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

49 +from vllm.multimodal import MULTIMODAL_REGISTRY

50 +from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,

51 + MultiModalKwargsItems)

52 +from vllm.multimodal.parse import ImageSize, MultiModalDataItems

53 +from vllm.multimodal.processing import (BaseMultiModalProcessor,

54 + BaseProcessingInfo, PromptReplacement,

55 + PromptUpdate)

56 +from vllm.multimodal.profiling import BaseDummyInputsBuilder

57 +from vllm.platforms import _Backend, current_platform

58 +from vllm.sequence import IntermediateTensors

59 +

60 +from .ernie45_vl_moe import Ernie4_5_VLMoeForCausalLM

61 +from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

62 + SupportsMultiModal, SupportsPP)

63 +from .utils import (AutoWeightsLoader, WeightsMapper, maybe_prefix,

64 + merge_multimodal_embeddings)

65 +from .vision import get_vit_attn_backend

66 +

67 +logger = init_logger(__name__)

68 +

69 +_MAX_FRAMES_PER_VIDEO = 16

70 +

71 +# === Vision Transformer === #

72 +

73 +

74 +def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:

75 + if not interleaved:

76 + x1, x2 = x.chunk(2, dim=-1)

77 + return torch.cat((-x2, x1), dim=-1)

78 + else:

79 + x1, x2 = x[..., ::2], x[..., 1::2]

80 + return rearrange(torch.stack((-x2, x1), dim=-1),

81 + "... d two -> ... (d two)",

82 + two=2)

83 +

84 +

85 +def apply_rotary_emb_torch(x: torch.Tensor,

86 + cos: torch.Tensor,

87 + sin: torch.Tensor,

88 + interleaved: bool = False) -> torch.Tensor:

89 + """

90 + x: (batch_size, seqlen, nheads, headdim)

91 + cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)

92 + """

93 + ro_dim = cos.shape[-1] * 2

94 + assert ro_dim <= x.shape[-1]

95 + cos = repeat(

96 + cos,

97 + "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")

98 + sin = repeat(

99 + sin,

100 + "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")

101 + return torch.cat(

102 + [

103 + x[..., :ro_dim] * cos +

104 + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]

105 + ],

106 + dim=-1,

107 + )

108 +

109 +

110 +def apply_rotary_pos_emb_vision(t: torch.Tensor,

111 + freqs: torch.Tensor) -> torch.Tensor:

112 + t_ = t.float()

113 + cos = freqs.cos()

114 + sin = freqs.sin()

115 + apply_rotary_emb = apply_rotary_emb_torch

116 + if current_platform.is_cuda():

117 + from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb

118 + output = apply_rotary_emb(t_, cos, sin).type_as(t)

119 + return output

120 +

121 +

122 +def all_gather_interleave(local_tensor, hidden_size: int, tp_size: int):

123 + """All-gather the input tensor interleavely across model parallel group."""

124 + import torch.distributed as dist

125 + gathered_tensors = [torch.zeros_like(local_tensor) for _ in range(tp_size)]

126 + dist.all_gather(gathered_tensors,

127 + local_tensor,

128 + group=parallel_state.get_tp_group().device_group)

129 +

130 + gathered_tensors_split = [

131 + torch.split(tensor, hidden_size // tp_size, -1)

132 + for tensor in gathered_tensors

133 + ]

134 + ordered_tensors = [

135 + tensor for pair in zip(*gathered_tensors_split) for tensor in pair

136 + ]

137 + result_tensor = torch.cat(ordered_tensors, dim=-1)

138 + return result_tensor

139 +

140 +

141 +class Ernie4_5_VisionAttention(nn.Module):

142 + """VisionAttention using VLLM framework APIs"""

143 +

144 + def __init__(

145 + self,

146 + embed_dim: int,

147 + num_heads: int,

148 + projection_size: int,

149 + quant_config: Optional[QuantizationConfig] = None,

150 + prefix: str = "",

151 + ) -> None:

152 + super().__init__()

153 + # Per attention head and per partition values.

154 + self.tp_size = parallel_state.get_tensor_model_parallel_world_size()

155 + self.tp_rank = parallel_state.get_tensor_model_parallel_rank()

156 + self.hidden_size_per_attention_head = dist_utils.divide(

157 + projection_size, num_heads)

158 + self.num_attention_heads_per_partition = dist_utils.divide(

159 + num_heads, self.tp_size)

160 +

161 + self.qkv = QKVParallelLinear(

162 + hidden_size=embed_dim,

163 + head_size=self.hidden_size_per_attention_head,

164 + total_num_heads=num_heads,

165 + total_num_kv_heads=num_heads,

166 + bias=True,

167 + quant_config=quant_config,

168 + prefix=f"{prefix}.qkv")

169 + self.proj = RowParallelLinear(input_size=projection_size,

170 + output_size=embed_dim,

171 + quant_config=quant_config,

172 + prefix=f"{prefix}.proj")

173 +

174 + # Detect attention implementation.

175 + self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)

176 + if self.attn_backend not in {

177 + _Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS,

178 + _Backend.ROCM_AITER_FA

179 + }:

180 + raise RuntimeError(

181 + f"Ernie45-VL does not support {self.attn_backend} backend now."

182 + )

183 + self.is_flash_attn_backend = self.attn_backend in {

184 + _Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA

185 + }

186 +

187 + def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:

188 + # [s, b, 3 * head * head_dim]

189 + seq_len, bs, _ = qkv.shape

190 + if self.tp_size > 1:

191 + qkv = all_gather_interleave(qkv, self.qkv.hidden_size,

192 + self.tp_size)

193 +

194 + # [s, b, 3 * head * head_dim] -> 3 * [s, b, head * head_dim]

195 + q, k, v = qkv.chunk(3, dim=2)

196 +

197 + # 3 * [s, b, head * head_dim]

198 + if self.tp_size > 1:

199 + splitter = partial(dist_utils.split_tensor_along_last_dim,

200 + num_partitions=self.tp_size)

201 + q = splitter(q)[self.tp_rank]

202 + k = splitter(k)[self.tp_rank]

203 + v = splitter(v)[self.tp_rank]

204 +

205 + # 3 * [s, b, head * head_dim] -> 3 * [s, b, head, head_dim]

206 + new_shape = (seq_len, bs, self.num_attention_heads_per_partition,

207 + self.hidden_size_per_attention_head)

208 + q, k, v = (x.view(*new_shape) for x in (q, k, v))

209 + return q, k, v

210 +

211 + def forward(

212 + self,

213 + x: torch.Tensor,

214 + cu_seqlens: torch.Tensor,

215 + rotary_pos_emb: torch.Tensor,

216 + max_seqlen: Optional[int] = None, # Only used for Flash Attention

217 + seqlens: Optional[list[int]] = None, # Only used for xFormers

218 + ) -> torch.Tensor:

219 + # [s, b, c] --> [s, b, head * 3 * head_dim]

220 + x, _ = self.qkv(x)

221 +

222 + # [s, b, 3 * head * head_dim] -> 3 * [s, b, head, head_dim]

223 + q, k, v = self.split_qkv(x)

224 + batch_size = q.shape[1]

225 +

226 + q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous()

227 + for x in (q, k, v))

228 + if rotary_pos_emb is not None:

229 + q = apply_rotary_pos_emb_vision(q, rotary_pos_emb)

230 + k = apply_rotary_pos_emb_vision(k, rotary_pos_emb)

231 +

232 + if self.is_flash_attn_backend:

233 + # from vllm_flash_attn.flash_attn_interface import (

234 + # flash_attn_varlen_func)

235 + if self.attn_backend == _Backend.ROCM_AITER_FA:

236 + from aiter import flash_attn_varlen_func

237 + else:

238 + from flash_attn import flash_attn_varlen_func

239 +

240 + q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])

241 +

242 + output = flash_attn_varlen_func(q,

243 + k,

244 + v,

245 + cu_seqlens_q=cu_seqlens,

246 + cu_seqlens_k=cu_seqlens,

247 + max_seqlen_q=max_seqlen,

248 + max_seqlen_k=max_seqlen,

249 + dropout_p=0.0,

250 + causal=False)

251 +

252 + context_layer = rearrange(output,

253 + "(b s) ... -> b s ...",

254 + b=batch_size)

255 + elif self.attn_backend == _Backend.TORCH_SDPA:

256 + # Execute attention entry by entry for speed & less VRAM.

257 + outputs = []

258 + for i in range(1, len(cu_seqlens)):

259 + start_idx = cu_seqlens[i - 1]

260 + end_idx = cu_seqlens[i]

261 + q_i = q[:, start_idx:end_idx]

262 + k_i = k[:, start_idx:end_idx]

263 + v_i = v[:, start_idx:end_idx]

264 + q_i, k_i, v_i = (rearrange(x, "b s h d -> b h s d")

265 + for x in [q_i, k_i, v_i])

266 + output_i = F.scaled_dot_product_attention(q_i,

267 + k_i,

268 + v_i,

269 + dropout_p=0.0)

270 + output_i = rearrange(output_i, "b h s d -> b s h d ")

271 + outputs.append(output_i)

272 + context_layer = torch.cat(outputs, dim=1)

273 + elif self.attn_backend == _Backend.XFORMERS:

274 + from xformers import ops as xops

275 + from xformers.ops.fmha.attn_bias import BlockDiagonalMask

276 +

277 + attn_bias = BlockDiagonalMask.from_seqlens(q_seqlen=seqlens,

278 + kv_seqlen=None,

279 + device=q.device)

280 +

281 + context_layer = xops.memory_efficient_attention_forward(

282 + q, k, v, attn_bias=attn_bias, p=0, scale=None)

283 + context_layer = rearrange(context_layer,

284 + "b s h d -> s b (h d)").contiguous()

285 +

286 + output, _ = self.proj(context_layer)

287 + return output

288 +

289 +

290 +class Ernie4_5_VisionMLP(nn.Module):

291 +

292 + def __init__(

293 + self,

294 + in_features: int,

295 + hidden_features: int,

296 + act_layer: type[nn.Module] = QuickGELU,

297 + quant_config: Optional[QuantizationConfig] = None,

298 + prefix: str = "",

299 + ):

300 + super().__init__()

301 + self.fc1 = ColumnParallelLinear(in_features,

302 + hidden_features,

303 + quant_config=quant_config,

304 + prefix=f"{prefix}.fc1")

305 + self.act = act_layer()

306 + self.fc2 = RowParallelLinear(hidden_features,

307 + in_features,

308 + quant_config=quant_config,

309 + prefix=f"{prefix}.fc2")

310 +

311 + def forward(self, x: torch.Tensor) -> torch.Tensor:

312 + x_parallel, _ = self.fc1(x)

313 + x_parallel = self.act(x_parallel)

314 + x, _ = self.fc2(x_parallel)

315 + return x

316 +

317 +

318 +class Ernie4_5_VisionBlock(nn.Module):

319 +

320 + def __init__(

321 + self,

322 + dim: int,

323 + num_heads: int,

324 + mlp_ratio: float,

325 + act_layer: type[nn.Module] = QuickGELU,

326 + norm_layer: Optional[Callable[[int], nn.Module]] = None,

327 + quant_config: Optional[QuantizationConfig] = None,

328 + prefix: str = "",

329 + ) -> None:

330 + super().__init__()

331 +

332 + if norm_layer is None:

333 + norm_layer = partial(nn.LayerNorm, eps=1e-6)

334 + self.norm1 = norm_layer(dim)

335 + self.norm2 = norm_layer(dim)

336 + mlp_hidden_dim = int(dim * mlp_ratio)

337 +

338 + self.attn = Ernie4_5_VisionAttention(embed_dim=dim,

339 + num_heads=num_heads,

340 + projection_size=dim,

341 + quant_config=quant_config,

342 + prefix=f"{prefix}.attn")

343 +

344 + self.mlp = Ernie4_5_VisionMLP(dim,

345 + mlp_hidden_dim,

346 + act_layer=act_layer,

347 + quant_config=quant_config,

348 + prefix=f"{prefix}.mlp")

349 +

350 + def forward(

351 + self,

352 + hidden_states: torch.Tensor,

353 + cu_seqlens: torch.Tensor,

354 + rotary_pos_emb: torch.Tensor,

355 + max_seqlen: Optional[int] = None, # Only used for Flash Attention

356 + seqlens: Optional[list[int]] = None, # Only used for xFormers

357 + ) -> torch.Tensor:

358 +

359 + hidden_states = hidden_states + self.attn(

360 + self.norm1(hidden_states),

361 + cu_seqlens=cu_seqlens,

362 + rotary_pos_emb=rotary_pos_emb,

363 + max_seqlen=max_seqlen,

364 + seqlens=seqlens,

365 + )

366 + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))

367 + return hidden_states

368 +

369 +

370 +class Ernie4_5_VisionPatchEmbed(nn.Module):

371 +

372 + def __init__(

373 + self,

374 + patch_size: int = 14,

375 + in_channels: int = 3,

376 + embed_dim: int = 1280,

377 + prefix="",

378 + ) -> None:

379 +

380 + super().__init__()

381 + self.patch_size = patch_size

382 + self.in_channels = in_channels

383 + self.embed_dim = embed_dim

384 +

385 + self.proj = nn.Linear(in_channels * patch_size * patch_size,

386 + embed_dim,

387 + bias=False)

388 +

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

390 +

391 + target_dtype = self.proj.weight.dtype

392 + hidden_states = hidden_states.to(target_dtype)

393 + hidden_states = self.proj(hidden_states)

394 +

395 + return hidden_states

396 +

397 +

398 +class Ernie4_5_VisionRotaryEmbedding(nn.Module):

399 +

400 + def __init__(self, dim: int, theta: float = 10000.0) -> None:

401 + super().__init__()

402 + self.inv_freq = 1.0 / theta**(

403 + torch.arange(start=0, end=dim, step=2, dtype=torch.float32) / dim)

404 +

405 + def forward(self, seqlen: int) -> torch.Tensor:

406 + seq = torch.arange(seqlen,

407 + device=self.inv_freq.device,

408 + dtype=self.inv_freq.dtype)

409 + freqs = torch.outer(input=seq, vec2=self.inv_freq)

410 + return freqs

411 +

412 +

413 +class Ernie4_5_VisionTransformer(nn.Module):

414 +

415 + def __init__(

416 + self,

417 + vision_config,

418 + norm_eps: float = 1e-6,

419 + quant_config: Optional[QuantizationConfig] = None,

420 + prefix: str = "",

421 + ) -> None:

422 +

423 + super().__init__()

424 + patch_size = vision_config.patch_size

425 + spatial_merge_size = vision_config.spatial_merge_size

426 + in_channels = vision_config.in_channels

427 + hidden_size = vision_config.hidden_size

428 + embed_dim = vision_config.embed_dim

429 + depth = vision_config.depth

430 + num_heads = vision_config.num_heads

431 + mlp_ratio = vision_config.mlp_ratio

432 +

433 + self.spatial_merge_size = spatial_merge_size

434 + self.num_heads = num_heads

435 + self.embed_dim = embed_dim

436 +

437 + self.patch_embed = Ernie4_5_VisionPatchEmbed(

438 + patch_size=patch_size,

439 + in_channels=in_channels,

440 + embed_dim=embed_dim,

441 + prefix=f"{prefix}.patch_embed",

442 + )

443 +

444 + norm_layer = partial(nn.LayerNorm, eps=norm_eps)

445 + head_dim = embed_dim // num_heads

446 + self.rotary_pos_emb = Ernie4_5_VisionRotaryEmbedding(head_dim // 2)

447 +

448 + self.blocks = nn.ModuleList([

449 + Ernie4_5_VisionBlock(dim=embed_dim,

450 + num_heads=num_heads,

451 + mlp_ratio=mlp_ratio,

452 + norm_layer=norm_layer,

453 + quant_config=quant_config,

454 + prefix=f"{prefix}.blocks.{layer_idx}")

455 + for layer_idx in range(depth)

456 + ])

457 +

458 + assert (hidden_size == embed_dim

459 + ), "vit's config.hidden must be equal to config.embed_dim"

460 + self.ln = nn.LayerNorm(hidden_size, eps=1e-6)

461 +

462 + self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)

463 +

464 + @property

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

466 + return self.patch_embed.proj.weight.dtype

467 +

468 + @property

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

470 + return self.patch_embed.proj.weight.device

471 +

472 + def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:

473 + pos_ids = []

474 + for t, h, w in grid_thw:

475 + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)

476 + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)

477 + hpos_ids = hpos_ids.reshape(

478 + h // self.spatial_merge_size,

479 + self.spatial_merge_size,

480 + w // self.spatial_merge_size,

481 + self.spatial_merge_size,

482 + ).permute(0, 2, 1, 3).flatten()

483 + wpos_ids = wpos_ids.reshape(

484 + h // self.spatial_merge_size,

485 + self.spatial_merge_size,

486 + w // self.spatial_merge_size,

487 + self.spatial_merge_size,

488 + ).permute(0, 2, 1, 3).flatten()

489 + pos_ids.append(

490 + torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))

491 + pos_ids = torch.cat(pos_ids, dim=0)

492 + max_grid_size = grid_thw[:, 1:].max()

493 + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)

494 + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)

495 + return rotary_pos_emb

496 +

497 + def compute_attn_mask_seqlen(

498 + self, cu_seqlens: torch.Tensor

499 + ) -> tuple[Optional[int], Optional[list[int]]]:

500 + max_seqlen, seqlens = None, None

501 + if self.attn_backend == _Backend.FLASH_ATTN:

502 + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

503 + elif self.attn_backend == _Backend.XFORMERS:

504 + seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()

505 + return max_seqlen, seqlens

506 +

507 + def forward(self,

508 + hidden_states: torch.Tensor,

509 + grid_thw: torch.Tensor,

510 + num_pad=0) -> torch.Tensor:

511 +

512 + hidden_states = self.patch_embed(hidden_states)

513 +

514 + rotary_pos_emb = self.rot_pos_emb(grid_thw)

515 + rotary_pos_emb = rotary_pos_emb.to(hidden_states.device)

516 +

517 + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2],

518 + grid_thw[:, 0]).cumsum(

519 + dim=0, dtype=torch.int32)

520 +

521 + if num_pad > 0:

522 + cu_seqlens = F.pad(cu_seqlens, (1, 1), value=0)

523 + cu_seqlens[-1] = cu_seqlens[-2] + num_pad

524 + else:

525 + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)

526 +

527 + # add batch size

528 + if hidden_states.ndim == 2:

529 + hidden_states = hidden_states.unsqueeze(dim=1)

530 +

531 + # pre-compute seqlens for attn mask to reduce cuMemcpy operations

532 + max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)

533 +

534 + for i, blk in enumerate(self.blocks):

535 + hidden_states = blk(

536 + hidden_states,

537 + cu_seqlens=cu_seqlens,

538 + rotary_pos_emb=rotary_pos_emb,

539 + max_seqlen=max_seqlen,

540 + seqlens=seqlens,

541 + )

542 +

543 + final_output = self.ln(hidden_states)

544 +

545 + if final_output.ndim == 3:

546 + final_output = final_output.squeeze(dim=1)

547 +

548 + return final_output

549 +

550 + def load_weights(self, weights) -> set[str]:

551 + params_dict = dict(self.named_parameters(remove_duplicate=False))

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

553 +

554 + for name, loaded_weight in weights:

555 + param = params_dict[name]

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

557 + default_weight_loader)

558 + weight_loader(param, loaded_weight)

559 + loaded_params.add(name)

560 + return loaded_params

561 +

562 +

563 +# === Vision Inputs === #

564 +

565 +

566 +class Ernie4_5_VLImagePixelInputs(TypedDict):

567 + type: Literal["pixel_values"]

568 + pixel_values: torch.Tensor

569 + """Shape:

570 + `(num_patches, num_channels * patch_size * patch_size)`

571 + """

572 +

573 + grid_thw: torch.Tensor

574 + """Shape: `(num_images, 3)`

575 + This should be in `(grid_t, grid_h, grid_w)` format.

576 + """

577 +

578 +

579 +Ernie4_5_VLImageInputs = Ernie4_5_VLImagePixelInputs

580 +

581 +

582 +class Ernie4_5_VLVideoPixelInputs(TypedDict):

583 + type: Literal["pixel_values_videos"]

584 + pixel_values_videos: torch.Tensor

585 + """Shape:

586 + `(num_patches,

587 + num_channels * temporal_patch_size * patch_size * patch_size)`

588 + """

589 +

590 + video_grid_thw: torch.Tensor

591 + """Shape: `(num_videos, 3)`

592 +

593 + This should be in `(grid_t, grid_h, grid_w)` format.

594 + """

595 +

596 +

597 +Ernie4_5_VLVideoInputs = Ernie4_5_VLImagePixelInputs

598 +

599 +# === Vision Processor === #

600 +

601 +

602 +def round_by_factor(number: Union[int, float], factor: int) -> int:

603 + return round(number / factor) * factor

604 +

605 +

606 +def ceil_by_factor(number: Union[int, float], factor: int) -> int:

607 + return math.ceil(number / factor) * factor

608 +

609 +

610 +def floor_by_factor(number: Union[int, float], factor: int) -> int:

611 + return math.floor(number / factor) * factor

612 +

613 +

614 +def smart_resize(

615 + height: int,

616 + width: int,

617 + factor: int = 28,

618 + min_pixels: int = 4 * 28 * 28,

619 + max_pixels: int = 16384 * 28 * 28,

620 +):

621 + MAX_RATIO = 200

622 + if max(height, width) / min(height, width) > MAX_RATIO:

623 + if height > width:

624 + new_width = max(factor, round_by_factor(width, factor))

625 + new_height = floor_by_factor(new_width * MAX_RATIO, factor)

626 + else:

627 + new_height = max(factor, round_by_factor(height, factor))

628 + new_width = floor_by_factor(new_height * MAX_RATIO, factor)

629 +

630 + height = new_height

631 + width = new_width

632 +

633 + h_bar = max(factor, round_by_factor(height, factor))

634 + w_bar = max(factor, round_by_factor(width, factor))

635 + if h_bar * w_bar > max_pixels:

636 + beta = math.sqrt((height * width) / max_pixels)

637 + h_bar = floor_by_factor(height / beta, factor)

638 + w_bar = floor_by_factor(width / beta, factor)

639 + elif h_bar * w_bar < min_pixels:

640 + beta = math.sqrt(min_pixels / (height * width))

641 + h_bar = ceil_by_factor(height * beta, factor)

642 + w_bar = ceil_by_factor(width * beta, factor)

643 +

644 + if min_pixels > h_bar * w_bar or h_bar * w_bar > max_pixels:

645 + raise ValueError(f"encounter invalid h_bar: {h_bar}, w_bar: {w_bar}")

646 +

647 + return h_bar, w_bar

648 +

649 +

650 +class VariableResolutionResamplerModel(nn.Module):

651 +

652 + def __init__(self,

653 + in_dim,

654 + out_dim,

655 + spatial_conv_size,

656 + temporal_conv_size,

657 + config,

658 + prefix: str = "") -> None:

659 + super().__init__()

660 + self.in_dim = in_dim

661 + self.out_dim = out_dim

662 + self.config = config

663 + self.spatial_conv_size = spatial_conv_size

664 + self.temporal_conv_size = temporal_conv_size

665 + self.use_temporal_conv = config.use_temporal_conv

666 +

667 + # compress 2d conv(picture) to 1d

668 + self.spatial_dim = (self.in_dim * self.spatial_conv_size *

669 + self.spatial_conv_size)

670 + # compress 3d conv(video) to 1d

671 + self.temporal_dim = (self.in_dim * self.spatial_conv_size *

672 + self.spatial_conv_size * self.temporal_conv_size)

673 +

674 + self.spatial_linear1 = ColumnParallelLinear(

675 + self.spatial_dim,

676 + self.spatial_dim,

677 + bias=True,

678 + gather_output=True,

679 + quant_config=getattr(config, 'quant_config', None),

680 + prefix=f"{prefix}.spatial_linear1",

681 + )

682 +

683 + self.spatial_gelu = nn.GELU()

684 +

685 + self.spatial_linear2 = ColumnParallelLinear(

686 + self.spatial_dim,

687 + self.spatial_dim,

688 + bias=True,

689 + gather_output=True,

690 + quant_config=getattr(config, 'quant_config', None),

691 + prefix=f"{prefix}.spatial_linear2",

692 + )

693 +

694 + self.spatial_norm = nn.LayerNorm(self.spatial_dim, eps=1e-6)

695 +

696 + if self.use_temporal_conv:

697 + self.temporal_linear1 = ColumnParallelLinear(

698 + self.temporal_dim,

699 + self.spatial_dim,

700 + bias=True,

701 + gather_output=True,

702 + quant_config=getattr(config, 'quant_config', None),

703 + prefix=f"{prefix}.temporal_linear1",

704 + )

705 +

706 + self.temporal_gelu = nn.GELU()

707 +

708 + self.temporal_linear2 = ColumnParallelLinear(

709 + self.spatial_dim,

710 + self.spatial_dim,

711 + bias=True,

712 + gather_output=True,

713 + quant_config=getattr(config, 'quant_config', None),

714 + prefix=f"{prefix}.temporal_linear2",

715 + )

716 +

717 + self.temporal_norm = nn.LayerNorm(self.spatial_dim, eps=1e-6)

718 +

719 + self.mlp = ColumnParallelLinear(

720 + self.spatial_dim,

721 + self.out_dim,

722 + bias=True,

723 + gather_output=True,

724 + quant_config=getattr(config, 'quant_config', None),

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

726 + )

727 +

728 + self.after_norm = RMSNorm(hidden_size=out_dim,

729 + eps=getattr(config, 'rms_norm_eps', 1e-6))

730 +

731 + def spatial_conv_reshape(self, x, spatial_conv_size):

732 + S, C = x.shape

733 + x = x.reshape([-1, C * (spatial_conv_size**2)])

734 + return x

735 +

736 + def forward(self, x, grid_thw):

737 +

738 + def fwd_spatial(x):

739 + x = self.spatial_conv_reshape(x, self.spatial_conv_size)

740 +

741 + x, _ = self.spatial_linear1(x)

742 + x = self.spatial_gelu(x)

743 + x, _ = self.spatial_linear2(x)

744 + x = self.spatial_norm(x)

745 +

746 + return x

747 +

748 + def fwd_placeholder(x, grid_thw, to_tensor=False):

749 +

750 + grid_thw_cpu = grid_thw.cpu().numpy()

751 + grid_t, grid_hw = grid_thw_cpu[:, 0], grid_thw_cpu[:, 1:]

752 + grid_hw_after_conv = grid_hw.prod(-1) // (self.spatial_conv_size**

753 + 2)

754 +

755 + tokens_per_img_or_vid = grid_thw_cpu.prod(-1) // (

756 + self.spatial_conv_size**2)

757 + batch_offset = np.empty(tokens_per_img_or_vid.size,

758 + dtype=tokens_per_img_or_vid.dtype)

759 + batch_offset[0] = 0

760 + batch_offset[1:] = tokens_per_img_or_vid.cumsum()[:-1]

761 +

762 + slice_offsets = []

763 + for temporoal_size, spatial_size, b_offset in zip(

764 + grid_t, grid_hw_after_conv, batch_offset):

765 + for temp_offset in range(0, temporoal_size, 2):

766 + slice_offsets.append(

767 + np.arange(

768 + b_offset + (temp_offset) * spatial_size,

769 + b_offset + (temp_offset + 1) * spatial_size,

770 + ))

771 + slice_offsets = torch.tensor(np.concatenate(slice_offsets,

772 + axis=-1)).to(x.device)

773 +

774 + slice_offsets2 = []

775 + for temporoal_size, spatial_size, b_offset in zip(

776 + grid_t, grid_hw_after_conv, batch_offset):

777 + for temp_offset in range(1 if temporoal_size > 1 else 0,

778 + temporoal_size, 2):

779 + slice_offsets2.append(

780 + np.arange(

781 + b_offset + (temp_offset) * spatial_size,

782 + b_offset + (temp_offset + 1) * spatial_size,

783 + ))

784 + slice_offsets2 = torch.tensor(

785 + np.concatenate(slice_offsets2, axis=-1)).to(x.device)

786 +

787 + x_timestep_1 = torch.index_select(x, dim=0, index=slice_offsets)

788 + x_timestep_2 = torch.index_select(x, dim=0, index=slice_offsets2)

789 + x = torch.concat([x_timestep_1, x_timestep_2], dim=-1)

790 + return x

791 +

792 + def fwd_temporal(x):

793 + x, _ = self.temporal_linear1(x)

794 + x = self.temporal_gelu(x)

795 + x, _ = self.temporal_linear2(x)

796 + x = self.temporal_norm(x)

797 + return x

798 +

799 + def fwd_mlp(x):

800 + x, _ = self.mlp(x)

801 + x = self.after_norm(x)

802 + return x

803 +

804 + x = fwd_spatial(x)

805 + if self.use_temporal_conv:

806 + x = fwd_placeholder(x, grid_thw)

807 + x = fwd_temporal(x)

808 + x = fwd_mlp(x)

809 + return x

810 +

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

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

813 +

814 + params_dict = dict(self.named_parameters(remove_duplicate=False))

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

816 +

817 + for name, loaded_weight in weights:

818 + if name not in params_dict:

819 + continue

820 + param = params_dict[name]

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

822 + default_weight_loader)

823 + weight_loader(param, loaded_weight)

824 + loaded_params.add(name)

825 + return loaded_params

826 +

827 +

828 +class Ernie4_5_VLProcessingInfo(BaseProcessingInfo):

829 +

830 + def get_hf_config(self):

831 + return self.ctx.model_config.hf_config

832 +

833 + def get_hf_processor(self, **kwargs: object):

834 + return self.ctx.get_hf_processor(use_fast=True, **kwargs)

835 +

836 + def get_image_processor(self, **kwargs: object):

837 + return self.get_hf_processor(**kwargs).image_processor

838 +

839 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

840 + return {"image": None, "video": None}

841 +

842 + def _get_vision_info(

843 + self,

844 + *,

845 + image_width: int,

846 + image_height: int,

847 + num_frames: int = 1,

848 + do_resize: bool = True,

849 + image_processor: Optional[Any],

850 + ) -> tuple[ImageSize, int]:

851 + if image_processor is None:

852 + image_processor = self.get_image_processor()

853 + hf_config = self.get_hf_config()

854 + vision_config = hf_config.vision_config

855 +

856 + patch_size = vision_config.patch_size

857 + spatial_conv_size = hf_config.spatial_conv_size

858 + temporal_conv_size = hf_config.temporal_conv_size

859 +

860 + if do_resize:

861 + resized_height, resized_width = smart_resize(

862 + height=image_height,

863 + width=image_width,

864 + factor=patch_size * spatial_conv_size,

865 + min_pixels=image_processor.min_pixels,

866 + max_pixels=image_processor.max_pixels,

867 + )

868 + preprocessed_size = ImageSize(width=resized_width,

869 + height=resized_height)

870 + else:

871 + preprocessed_size = ImageSize(width=image_width,

872 + height=image_height)

873 +

874 + grid_t = max(num_frames // temporal_conv_size, 1)

875 + grid_h = preprocessed_size.height // patch_size

876 + grid_w = preprocessed_size.width // patch_size

877 +

878 + num_patches = grid_t * grid_h * grid_w

879 + num_vision_tokens = num_patches // (spatial_conv_size**2)

880 +

881 + return preprocessed_size, num_vision_tokens

882 +

883 + def get_num_image_tokens(

884 + self,

885 + *,

886 + image_width: int,

887 + image_height: int,

888 + image_processor: Optional[Any],

889 + ) -> int:

890 + _, num_image_tokens = self._get_vision_info(

891 + image_width=image_width,

892 + image_height=image_height,

893 + image_processor=image_processor,

894 + )

895 + return num_image_tokens

896 +

897 + def get_num_video_tokens(

898 + self,

899 + *,

900 + image_width: int,

901 + image_height: int,

902 + num_frames: int,

903 + image_processor: Optional[Any],

904 + ) -> int:

905 + _, num_video_tokens = self._get_vision_info(

906 + image_width=image_width,

907 + image_height=image_height,

908 + num_frames=num_frames,

909 + image_processor=image_processor,

910 + )

911 + return num_video_tokens

912 +

913 + def get_image_size_with_most_features(self) -> ImageSize:

914 + max_image_size, _ = self._get_vision_info(

915 + image_width=9999999,

916 + image_height=9999999,

917 + image_processor=None,

918 + )

919 + return max_image_size

920 +

921 + def get_max_image_tokens(self) -> int:

922 + target_width, target_height = self.get_image_size_with_most_features()

923 +

924 + num_image_tokens = self.get_num_image_tokens(

925 + image_width=target_width,

926 + image_height=target_height,

927 + image_processor=None,

928 + )

929 + return num_image_tokens

930 +

931 + def _get_max_video_frames(self, max_tokens: int) -> int:

932 + target_width, target_height = self.get_image_size_with_most_features()

933 +

934 + num_frames = 0

935 +

936 + while True:

937 + next_num_frames = num_frames + 1

938 + next_max_tokens = self.get_num_video_tokens(

939 + image_width=target_width,

940 + image_height=target_height,

941 + num_frames=next_num_frames,

942 + image_processor=None,

943 + )

944 +

945 + if next_max_tokens > max_tokens:

946 + break

947 +

948 + num_frames = next_num_frames

949 +

950 + # If the number of frames is odd, discard one frame.

951 + if num_frames % 2 != 0:

952 + num_frames -= 1

953 +

954 + return num_frames

955 +

956 + def get_num_frames_with_most_features(

957 + self,

958 + seq_len: int,

959 + mm_counts: Mapping[str, int],

960 + ) -> int:

961 + max_images = mm_counts.get("image", 0)

962 + max_videos = mm_counts.get("video", 0)

963 +

964 + max_image_tokens = self.get_max_image_tokens() * max_images

965 + max_total_frames = self._get_max_video_frames(seq_len -

966 + max_image_tokens)

967 + max_frames_per_video = min(max_total_frames // max(max_videos, 1),

968 + _MAX_FRAMES_PER_VIDEO)

969 +

970 + return max(max_frames_per_video, 2)

971 +

972 + def get_max_video_tokens(

973 + self,

974 + seq_len: int,

975 + mm_counts: Mapping[str, int],

976 + ) -> int:

977 + target_width, target_height = self.get_image_size_with_most_features()

978 +

979 + return self.get_num_video_tokens(

980 + image_width=target_width,

981 + image_height=target_height,

982 + num_frames=self.get_num_frames_with_most_features(

983 + seq_len, mm_counts),

984 + image_processor=None,

985 + )

986 +

987 +

988 +class Ernie4_5VLMultiModalProcessor(

989 + BaseMultiModalProcessor[Ernie4_5_VLProcessingInfo]):

990 +

991 + def _pixel_values_norm(

992 + self,

993 + pixel_values: torch.Tensor,

994 + mm_kwargs: object,

995 + ) -> torch.Tensor:

996 + hf_config = self.info.get_hf_config()

997 + vision_config = hf_config.vision_config

998 + image_processor = self.info.get_image_processor(**mm_kwargs)

999 + image_mean_tensor = torch.tensor(image_processor.image_mean,

1000 + dtype=torch.float32).reshape(

1001 + [1, 3, 1, 1])

1002 + image_std_tensor = torch.tensor(image_processor.image_std,

1003 + dtype=torch.float32).reshape(

1004 + [1, 3, 1, 1])

1005 + rescale_factor = torch.tensor(image_processor.rescale_factor,

1006 + dtype=torch.float32)

1007 + patch_size_squared = vision_config.patch_size**2

1008 +

1009 + image_mean_tensor = (image_mean_tensor.squeeze(

1010 + [-2, -1]).repeat_interleave(patch_size_squared, -1))

1011 + image_std_tensor = (image_std_tensor.squeeze(

1012 + [-2, -1]).repeat_interleave(patch_size_squared, -1))

1013 +

1014 + if not image_mean_tensor.is_contiguous():

1015 + image_mean_tensor = image_mean_tensor.contiguous()

1016 + if not image_std_tensor.is_contiguous():

1017 + image_std_tensor = image_std_tensor.contiguous()

1018 +

1019 + pixel_values = (rescale_factor * pixel_values.to(torch.float32) -

1020 + image_mean_tensor) / image_std_tensor

1021 + pixel_values = pixel_values.to(hf_config.torch_dtype)

1022 + return pixel_values

1023 +

1024 + def _call_hf_processor(

1025 + self,

1026 + prompt: str,

1027 + mm_data: Mapping[str, object],

1028 + mm_kwargs: Mapping[str, object],

1029 + tok_kwargs: Mapping[str, object],

1030 + ) -> BatchFeature:

1031 + # when the prompt is not empty but the multimodal data is empty,

1032 + # directly invoke the tokenizer.

1033 + if "images" not in mm_data and "videos" not in mm_data and prompt != "":

1034 + tokenizer = self.info.get_tokenizer()

1035 + prompt_ids = tokenizer.encode(prompt)

1036 + tokenizer_output = BatchFeature(dict(input_ids=[prompt_ids]),

1037 + tensor_type="pt")

1038 + return tokenizer_output

1039 +

1040 + if "images" not in mm_data:

1041 + mm_data["images"] = []

1042 + if "videos" not in mm_data:

1043 + mm_data["videos"] = []

1044 + processor_output = self.info.ctx.call_hf_processor(

1045 + self.info.get_hf_processor(**mm_kwargs),

1046 + dict(text=[prompt],

1047 + images=mm_data["images"],

1048 + videos=mm_data["videos"]),

1049 + dict(**mm_kwargs, **tok_kwargs),

1050 + )

1051 +

1052 + # Divide the processor_output into two modalities: image and video.

1053 + if processor_output is not None:

1054 + pixel_values = processor_output['images']

1055 + if pixel_values is not None:

1056 + processor_output['images'] = self._pixel_values_norm(

1057 + pixel_values, mm_kwargs)

1058 + for key in list(processor_output.keys()):

1059 + if processor_output[key] is None:

1060 + del processor_output[key]

1061 + continue

1062 + if key == "grid_thw":

1063 + grid_thw = processor_output['grid_thw']

1064 + pixel_values_all = processor_output['images']

1065 + # Identify elements where the first

1066 + # dimension is greater than 1 and

1067 + # treat them as the video modality

1068 + mask = grid_thw[:, 0] > 1

1069 + processor_output["video_grid_thw"] = grid_thw[mask]

1070 + processor_output["image_grid_thw"] = grid_thw[~mask]

1071 + image_patch_num = processor_output["image_grid_thw"].prod(

1072 + dim=1).sum()

1073 + processor_output[

1074 + 'pixel_values'] = pixel_values_all[:image_patch_num]

1075 + processor_output['pixel_values_videos'] = pixel_values_all[

1076 + image_patch_num:]

1077 + del processor_output['images']

1078 +

1079 + return processor_output

1080 +

1081 + def _get_prompt_updates(

1082 + self,

1083 + mm_items: MultiModalDataItems,

1084 + hf_processor_mm_kwargs: Mapping[str, Any],

1085 + out_mm_kwargs: MultiModalKwargsItems,

1086 + ) -> Sequence[PromptUpdate]:

1087 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

1088 +

1089 + before_placeholder = {

1090 + "image": "<|image@placeholder|>",

1091 + "video": "<|video@placeholder|>"

1092 + }

1093 +

1094 + after_placeholder = {

1095 + # image and video have same placeholder

1096 + "image": "<|IMAGE_PLACEHOLDER|>",

1097 + "video": "<|IMAGE_PLACEHOLDER|>"

1098 + }

1099 +

1100 + merge_length = hf_processor.spatial_conv_size**2

1101 +

1102 + def get_replacement_ernie45vl(item_idx: int, modality: str):

1103 + out_item = out_mm_kwargs[modality][item_idx]

1104 + grid_thw = out_item[f"{modality}_grid_thw"].data

1105 + assert isinstance(grid_thw, torch.Tensor)

1106 + if modality == "video":

1107 + num_tokens = int(grid_thw.prod(

1108 + )) // hf_processor.temporal_conv_size // merge_length

1109 + else:

1110 + num_tokens = int(grid_thw.prod()) // merge_length

1111 + return after_placeholder[modality] * num_tokens

1112 +

1113 + return [

1114 + PromptReplacement(

1115 + modality=modality,

1116 + target=before_placeholder[modality],

1117 + replacement=partial(get_replacement_ernie45vl,

1118 + modality=modality),

1119 + ) for modality in ("image", "video")

1120 + ]

1121 +

1122 + def _get_mm_fields_config(

1123 + self,

1124 + hf_inputs: BatchFeature,

1125 + hf_processor_mm_kwargs: Mapping[str, object],

1126 + ) -> Mapping[str, MultiModalFieldConfig]:

1127 +

1128 + image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))

1129 + image_grid_sizes = image_grid_thw.prod(-1)

1130 +

1131 + video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3)))

1132 + video_grid_sizes = video_grid_thw.prod(-1)

1133 +

1134 + return dict(

1135 + pixel_values=MultiModalFieldConfig.flat_from_sizes(

1136 + "image", image_grid_sizes),

1137 + image_grid_thw=MultiModalFieldConfig.batched("image"),

1138 + pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(

1139 + "video", video_grid_sizes),

1140 + video_grid_thw=MultiModalFieldConfig.batched("video"),

1141 + )

1142 +

1143 +

1144 +class Ernie4_5_VLDummyInputsBuilder(

1145 + BaseDummyInputsBuilder[Ernie4_5_VLProcessingInfo]):

1146 +

1147 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

1148 + num_images = mm_counts.get("image", 0)

1149 + num_videos = mm_counts.get("video", 0)

1150 + prompt = ""

1151 + for i in range(num_images):

1152 + prompt += (f"Picture {i+1}:"

1153 + "<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>")

1154 +

1155 + for i in range(num_videos):

1156 + prompt += (f"Video {i+1}:"

1157 + "<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>")

1158 + return prompt

1159 +

1160 + def get_dummy_mm_data(

1161 + self,

1162 + seq_len: int,

1163 + mm_counts: Mapping[str, int],

1164 + ) -> MultiModalDataDict:

1165 + num_images = mm_counts.get("image", 0)

1166 + num_videos = mm_counts.get("video", 0)

1167 +

1168 + target_width, target_height = \

1169 + self.info.get_image_size_with_most_features()

1170 + target_num_frames = \

1171 + self.info.get_num_frames_with_most_features(seq_len, mm_counts)

1172 +

1173 + return {

1174 + "image":

1175 + self._get_dummy_images(width=target_width,

1176 + height=target_height,

1177 + num_images=num_images),

1178 + "video":

1179 + self._get_dummy_videos(width=target_width,

1180 + height=target_height,

1181 + num_frames=target_num_frames,

1182 + num_videos=num_videos)

1183 + }

1184 +

1185 +

1186 +@MULTIMODAL_REGISTRY.register_processor(

1187 + Ernie4_5VLMultiModalProcessor,

1188 + info=Ernie4_5_VLProcessingInfo,

1189 + dummy_inputs=Ernie4_5_VLDummyInputsBuilder)

1190 +class Ernie4_5_VLMoeForConditionalGeneration(nn.Module, SupportsMultiModal,

1191 + SupportsLoRA, SupportsPP):

1192 +

1193 + packed_modules_mapping = {

1194 + "qkv_proj": [

1195 + "q_proj",

1196 + "k_proj",

1197 + "v_proj",

1198 + ],

1199 + "gate_up_proj": [

1200 + "gate_proj",

1201 + "up_proj",

1202 + ],

1203 + }

1204 +

1205 + # To ensure correct weight loading and mapping.

1206 + hf_to_vllm_mapper = WeightsMapper(

1207 + orig_to_new_prefix={

1208 + "lm_head.": "language_model.lm_head.",

1209 + "model.": "language_model.model.",

1210 + # model.resampler_model.-> language_model.model.resampler_model.

1211 + # language_model.model.resampler_model. -> resampler_model.

1212 + "language_model.model.resampler_model.": "resampler_model.",

1213 + },

1214 + # resampler_weight_mappings

1215 + orig_to_new_substr={

1216 + "spatial_linear.0.": "spatial_linear1.",

1217 + "spatial_linear.2.": "spatial_linear2.",

1218 + "spatial_linear.3.": "spatial_norm.",

1219 + "temporal_linear.0.": "temporal_linear1.",

1220 + "temporal_linear.2.": "temporal_linear2.",

1221 + "temporal_linear.3.": "temporal_norm.",

1222 + })

1223 +

1224 + @classmethod

1225 + def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1226 + if modality.startswith("image"):

1227 + return "<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>"

1228 + if modality.startswith("video"):

1229 + return "<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>"

1230 +

1231 + raise ValueError("Only image or video modality is supported")

1232 +

1233 + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:

1234 + super().__init__()

1235 + config = vllm_config.model_config.hf_config

1236 + quant_config = vllm_config.quant_config

1237 + multimodal_config = vllm_config.model_config.multimodal_config

1238 +

1239 + self.config = config

1240 + self.multimodal_config = multimodal_config

1241 +

1242 + self.vision_model = Ernie4_5_VisionTransformer(

1243 + config.vision_config,

1244 + norm_eps=getattr(config, "rms_norm_eps", 1e-6),

1245 + quant_config=quant_config,

1246 + prefix=maybe_prefix(prefix, "vision_model"),

1247 + )

1248 +

1249 + self.language_model = Ernie4_5_VLMoeForCausalLM(

1250 + vllm_config=vllm_config,

1251 + prefix=maybe_prefix(prefix, "language_model"),

1252 + )

1253 +

1254 + self.resampler_model = VariableResolutionResamplerModel(

1255 + self.config.pixel_hidden_size,

1256 + self.config.hidden_size,

1257 + self.config.spatial_conv_size,

1258 + self.config.temporal_conv_size,

1259 + config=self.config,

1260 + prefix=maybe_prefix(prefix, "resampler_model"))

1261 +

1262 + self.visual_token_mask = None

1263 + self.make_empty_intermediate_tensors = (

1264 + self.language_model.make_empty_intermediate_tensors)

1265 +

1266 + def compute_logits(

1267 + self,

1268 + hidden_states: torch.Tensor,

1269 + sampling_metadata: SamplingMetadata,

1270 + ) -> Optional[torch.Tensor]:

1271 + """compute logits"""

1272 + return self.language_model.compute_logits(hidden_states,

1273 + sampling_metadata)

1274 +

1275 + def _vision_forward(

1276 + self,

1277 + pixel_values: torch.Tensor,

1278 + grid_thw: torch.Tensor,

1279 + ) -> torch.Tensor:

1280 + if grid_thw is not None:

1281 + grid_thw = grid_thw[grid_thw > 0]

1282 + if grid_thw.numel() % 3 != 0:

1283 + raise ValueError(

1284 + f"grid_thw has {grid_thw.numel()} elements after filtering,"

1285 + "which is not divisible by 3.")

1286 + grid_thw = grid_thw.reshape(-1, 3)

1287 + # example: [[1,64,64],[2,80,80]] -> [[1,64,64],[1,80,80],[1,80,80]]

1288 + grid_thw = F.pad(

1289 + torch.repeat_interleave(grid_thw[:, 1:], grid_thw[:, 0], 0),

1290 + [1, 0, 0, 0],

1291 + value=1,

1292 + )

1293 + image_features = self.vision_model(pixel_values, grid_thw)

1294 + return image_features

1295 +

1296 + def _set_visual_token_mask(self, input_ids: torch.Tensor) -> None:

1297 + if getattr(self.config, "im_patch_id", None) is not None:

1298 + self.visual_token_mask = (

1299 + input_ids == self.config.im_patch_id).reshape(-1, 1)

1300 + else:

1301 + self.visual_token_mask = None

1302 +

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

1304 + return self.language_model

1305 +

1306 + def _validate_and_reshape_mm_tensor(self, mm_input: object,

1307 + name: str) -> torch.Tensor:

1308 + if not isinstance(mm_input, (torch.Tensor, list)):

1309 + raise ValueError(f"Incorrect type of {name}. "

1310 + f"Got type: {type(mm_input)}")

1311 + if isinstance(mm_input, torch.Tensor):

1312 + if mm_input.ndim == 2:

1313 + return mm_input

1314 + if mm_input.ndim != 3:

1315 + raise ValueError(f"{name} should be 2D or batched 3D tensor. "

1316 + f"Got ndim: {mm_input.ndim} "

1317 + f"(shape={mm_input.shape})")

1318 + return torch.concat(list(mm_input))

1319 + else:

1320 + return torch.concat(mm_input)

1321 +

1322 + def _parse_and_validate_image_input(

1323 + self, **kwargs: object) -> Optional[Ernie4_5_VLImageInputs]:

1324 + pixel_values = kwargs.pop("pixel_values", None)

1325 + image_grid_thw = kwargs.pop("image_grid_thw", None)

1326 +

1327 + if pixel_values is None:

1328 + return None

1329 +

1330 + if pixel_values is not None:

1331 + pixel_values = self._validate_and_reshape_mm_tensor(

1332 + pixel_values, "image pixel values")

1333 + image_grid_thw = self._validate_and_reshape_mm_tensor(

1334 + image_grid_thw, "image grid_thw")

1335 +

1336 + if not isinstance(pixel_values, (torch.Tensor, list)):

1337 + raise ValueError("Incorrect type of image pixel values. "

1338 + f"Got type: {type(pixel_values)}")

1339 +

1340 + return Ernie4_5_VLImagePixelInputs(type="pixel_values",

1341 + pixel_values=pixel_values,

1342 + image_grid_thw=image_grid_thw)

1343 +

1344 + def _parse_and_validate_video_input(

1345 + self, **kwargs: object) -> Optional[Ernie4_5_VLVideoInputs]:

1346 + pixel_values_videos = kwargs.pop("pixel_values_videos", None)

1347 + video_grid_thw = kwargs.pop("video_grid_thw", None)

1348 +

1349 + if pixel_values_videos is None:

1350 + return None

1351 +

1352 + if pixel_values_videos is not None:

1353 + pixel_values_videos = self._validate_and_reshape_mm_tensor(

1354 + pixel_values_videos, "video pixel values")

1355 + video_grid_thw = self._validate_and_reshape_mm_tensor(

1356 + video_grid_thw, "video grid_thw")

1357 +

1358 + return Ernie4_5_VLVideoPixelInputs(

1359 + type="pixel_values_videos",

1360 + pixel_values_videos=pixel_values_videos,

1361 + video_grid_thw=video_grid_thw,

1362 + )

1363 +

1364 + def _process_image_input(

1365 + self,

1366 + image_input: Ernie4_5_VLImageInputs) -> tuple[torch.Tensor, ...]:

1367 +

1368 + grid_thw = image_input["image_grid_thw"]

1369 + assert grid_thw.ndim == 2

1370 +

1371 + pixel_values = image_input["pixel_values"].type(

1372 + self.vision_model.dtype)

1373 + image_features = self._vision_forward(pixel_values=pixel_values,

1374 + grid_thw=grid_thw)

1375 + image_embeds = self.resampler_model(image_features, grid_thw)

1376 +

1377 + merge_size = self.vision_model.spatial_merge_size

1378 + sizes = grid_thw.prod(-1) // merge_size // merge_size

1379 +

1380 + return image_embeds.split(sizes.tolist())

1381 +

1382 + def _process_video_input(

1383 + self,

1384 + video_input: Ernie4_5_VLVideoInputs) -> tuple[torch.Tensor, ...]:

1385 +

1386 + grid_thw = video_input["video_grid_thw"]

1387 + assert grid_thw.ndim == 2

1388 +

1389 + pixel_values_videos = video_input["pixel_values_videos"].type(

1390 + self.vision_model.dtype)

1391 + video_features = self._vision_forward(pixel_values=pixel_values_videos,

1392 + grid_thw=grid_thw)

1393 + video_embeds = self.resampler_model(video_features, grid_thw)

1394 +

1395 + merge_size = self.vision_model.spatial_merge_size

1396 + sizes = (grid_thw.prod(-1) //

1397 + self.config.temporal_conv_size) // merge_size // merge_size

1398 +

1399 + return video_embeds.split(sizes.tolist())

1400 +

1401 + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

1402 + modalities = {}

1403 +

1404 + # Preserve the order of modalities if there are multiple of them

1405 + # from the order of kwargs.

1406 + for input_key in kwargs:

1407 + if input_key in ("pixel_values",

1408 + "image_embeds") and "images" not in modalities:

1409 + modalities["images"] = self._parse_and_validate_image_input(

1410 + **kwargs)

1411 + if input_key in ("pixel_values_videos",

1412 + "video_embeds") and "videos" not in modalities:

1413 + modalities["videos"] = self._parse_and_validate_video_input(

1414 + **kwargs)

1415 +

1416 + return modalities

1417 +

1418 + def get_multimodal_embeddings(

1419 + self, **kwargs: object) -> Optional[MultiModalEmbeddings]:

1420 +

1421 + modalities = self._parse_and_validate_multimodal_inputs(**kwargs)

1422 + if not modalities:

1423 + return None

1424 +

1425 + # The result multimodal_embeddings is tuple of tensors, with each

1426 + # tensor correspoending to a multimodal data item (image or video).

1427 + multimodal_embeddings: tuple[torch.Tensor, ...] = ()

1428 +

1429 + # NOTE: It is important to iterate over the keys in this dictionary

1430 + # to preserve the order of the modalities.

1431 + for modality in modalities:

1432 + if modality == "images":

1433 + image_input = modalities["images"]

1434 + vision_embeddings = self._process_image_input(image_input)

1435 + multimodal_embeddings += vision_embeddings

1436 + if modality == "videos":

1437 + video_input = modalities["videos"]

1438 + video_embeddings = self._process_video_input(video_input)

1439 + multimodal_embeddings += video_embeddings

1440 +

1441 + return multimodal_embeddings

1442 +

1443 + def get_input_embeddings(

1444 + self,

1445 + input_ids: torch.Tensor,

1446 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1447 + ) -> torch.Tensor:

1448 +

1449 + inputs_embeds = self.language_model.get_input_embeddings(input_ids)

1450 +

1451 + if multimodal_embeddings is None:

1452 + return inputs_embeds

1453 +

1454 + self._set_visual_token_mask(input_ids)

1455 + inputs_embeds = merge_multimodal_embeddings(input_ids, inputs_embeds,

1456 + multimodal_embeddings,

1457 + [self.config.im_patch_id])

1458 + return inputs_embeds

1459 +

1460 + def forward(

1461 + self,

1462 + input_ids: torch.Tensor,

1463 + positions: torch.Tensor,

1464 + intermediate_tensors: Optional[IntermediateTensors] = None,

1465 + inputs_embeds: Optional[torch.Tensor] = None,

1466 + **kwargs,

1467 + ):

1468 +

1469 + forward_kwargs = {

1470 + "input_ids": input_ids,

1471 + "positions": positions,

1472 + "intermediate_tensors": intermediate_tensors,

1473 + "inputs_embeds": inputs_embeds,

1474 + }

1475 +

1476 + if self.visual_token_mask is not None:

1477 +

1478 + if self.visual_token_mask.shape[0] != inputs_embeds.shape[0]:

1479 + padding_len = inputs_embeds.shape[

1480 + 0] - self.visual_token_mask.shape[0]

1481 + # right pad False

1482 + pad = torch.zeros(

1483 + (padding_len, self.visual_token_mask.shape[1]),

1484 + dtype=self.visual_token_mask.dtype,

1485 + device=self.visual_token_mask.device)

1486 + self.visual_token_mask = torch.cat(

1487 + [self.visual_token_mask, pad], dim=0)

1488 +

1489 + forward_kwargs.update(

1490 + {"visual_token_mask": self.visual_token_mask})

1491 + self.visual_token_mask = None

1492 +

1493 + hidden_states = self.language_model.model(

1494 + **forward_kwargs,

1495 + **kwargs,

1496 + )

1497 +

1498 + return hidden_states

1499 +

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

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

1502 +

1503 + loader = AutoWeightsLoader(self)

1504 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/ernie45_vl_moe.py new

+723

@@ -0,0 +1,723 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

7 +#

8 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX

9 +# and OPT implementations in this library. It has been modified from its

10 +# original forms to accommodate minor architectural differences compared

11 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.

12 +#

13 +# Licensed under the Apache License, Version 2.0 (the "License");

14 +# you may not use this file except in compliance with the License.

15 +# You may obtain a copy of the License at

16 +#

17 +# http://www.apache.org/licenses/LICENSE-2.0

18 +#

19 +# Unless required by applicable law or agreed to in writing, software

20 +# distributed under the License is distributed on an "AS IS" BASIS,

21 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

22 +# See the License for the specific language governing permissions and

23 +# limitations under the License.

24 +"""Inference-only Erine VL model compatible with HuggingFace weights."""

25 +from collections.abc import Iterable

26 +from typing import Any, Optional, Union

27 +

28 +import torch

29 +from torch import nn

30 +from transformers import PretrainedConfig

31 +

32 +from vllm.attention import Attention

33 +# from vllm.compilation.decorators import support_torch_compile

34 +from vllm.config import CacheConfig, VllmConfig

35 +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

36 +from vllm.logger import init_logger

37 +from vllm.model_executor.layers.fused_moe import FusedMoE

38 +from vllm.model_executor.layers.layernorm import RMSNorm

39 +from vllm.model_executor.layers.linear import (QKVParallelLinear,

40 + ReplicatedLinear,

41 + RowParallelLinear)

42 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

43 +from vllm.model_executor.layers.quantization import QuantizationConfig

44 +from vllm.model_executor.layers.rotary_embedding.ernie45_vl_rope import (

45 + Ernie4_5_VLRotaryEmbedding)

46 +from vllm.model_executor.layers.vocab_parallel_embedding import (

47 + ParallelLMHead, VocabParallelEmbedding)

48 +from vllm.model_executor.model_loader.weight_utils import (

49 + default_weight_loader, maybe_remap_kv_scale_name)

50 +from vllm.model_executor.sampling_metadata import SamplingMetadata

51 +from vllm.sequence import IntermediateTensors

52 +

53 +from .ernie45_moe import Ernie4_5_MoeMLP

54 +from .interfaces import SupportsPP

55 +from .utils import (PPMissingLayer, extract_layer_index,

56 + is_pp_missing_parameter,

57 + make_empty_intermediate_tensors_factory, make_layers,

58 + maybe_prefix)

59 +

60 +logger = init_logger(__name__)

61 +

62 +

63 +class Ernie4_5_VLMoeMLP(Ernie4_5_MoeMLP):

64 + pass

65 +

66 +

67 +class Ernie4_5_VLMoeAttention(nn.Module):

68 +

69 + def __init__(

70 + self,

71 + hidden_size: int,

72 + num_heads: int,

73 + num_kv_heads: int,

74 + head_dim: Optional[int] = None,

75 + rope_theta: float = 500000,

76 + rope_scaling: Optional[dict[str, Any]] = None,

77 + freq_allocation: int = 20,

78 + max_position_embeddings: int = 131072,

79 + rms_norm_eps: float = 1e-05,

80 + qkv_bias: bool = False,

81 + cache_config: Optional[CacheConfig] = None,

82 + quant_config: Optional[QuantizationConfig] = None,

83 + prefix: str = "",

84 + ) -> None:

85 + super().__init__()

86 + layer_idx = extract_layer_index(prefix) if len(prefix) > 0 else 0

87 + self.layer_idx = layer_idx

88 + self.hidden_size = hidden_size

89 + tp_size = get_tensor_model_parallel_world_size()

90 + self.total_num_heads = num_heads

91 + assert self.total_num_heads % tp_size == 0

92 + self.num_heads = self.total_num_heads // tp_size

93 +

94 + self.total_num_kv_heads = num_kv_heads

95 + if self.total_num_kv_heads >= tp_size:

96 + # Number of KV heads is greater than TP size, so we partition

97 + # the KV heads across multiple tensor parallel GPUs.

98 + assert self.total_num_kv_heads % tp_size == 0

99 + else:

100 + # Number of KV heads is less than TP size, so we replicate

101 + # the KV heads across multiple tensor parallel GPUs.

102 + assert tp_size % self.total_num_kv_heads == 0

103 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

104 + self.head_dim = head_dim or (hidden_size // self.total_num_heads)

105 +

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

107 + self.kv_size = self.num_kv_heads * self.head_dim

108 + self.scaling = self.head_dim**-0.5

109 + self.rope_theta = rope_theta

110 + self.max_position_embeddings = max_position_embeddings

111 +

112 + self.qkv_proj = QKVParallelLinear(hidden_size,

113 + self.head_dim,

114 + self.total_num_heads,

115 + self.total_num_kv_heads,

116 + bias=qkv_bias,

117 + quant_config=quant_config,

118 + prefix=f"{prefix}.qkv_proj")

119 +

120 + self.o_proj = RowParallelLinear(self.total_num_heads * self.head_dim,

121 + hidden_size,

122 + bias=False,

123 + quant_config=quant_config,

124 + prefix=f"{prefix}.o_proj")

125 +

126 + t_rope = freq_allocation

127 + h_rope = (self.head_dim // 2 - freq_allocation) // 2

128 + w_rope = (self.head_dim // 2 - freq_allocation) // 2

129 +

130 + self.rotary_emb = Ernie4_5_VLRotaryEmbedding(

131 + head_size=self.head_dim,

132 + rotary_dim=self.head_dim,

133 + max_position_embeddings=max_position_embeddings,

134 + base=rope_theta,

135 + is_neox_style=False,

136 + dtype=torch.get_default_dtype(),

137 + mrope_section=[h_rope, w_rope, t_rope])

138 +

139 + self.attn = Attention(self.num_heads,

140 + self.head_dim,

141 + self.scaling,

142 + num_kv_heads=self.num_kv_heads,

143 + cache_config=cache_config,

144 + quant_config=quant_config,

145 + prefix=f"{prefix}.attn")

146 +

147 + def forward(

148 + self,

149 + positions: torch.Tensor,

150 + hidden_states: torch.Tensor,

151 + ) -> torch.Tensor:

152 +

153 + qkv, _ = self.qkv_proj(hidden_states)

154 +

155 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

156 + q, k = self.rotary_emb(positions, q, k)

157 +

158 + # Attention

159 + attn_output = self.attn(q, k, v)

160 + # Output projection

161 + output, _ = self.o_proj(attn_output)

162 + return output

163 +

164 +

165 +class Ernie4_5_VLMoeMoE(nn.Module):

166 +

167 + def __init__(

168 + self,

169 + config: PretrainedConfig,

170 + quant_config: Optional[QuantizationConfig] = None,

171 + prefix: str = "",

172 + ):

173 + super().__init__()

174 +

175 + layer_idx = extract_layer_index(prefix)

176 + self.layer_idx = layer_idx

177 + self.tp_size = get_tensor_model_parallel_world_size()

178 + self.has_shared_experts = (getattr(config, "moe_num_shared_experts", 0)

179 + > 0)

180 + self.hidden_size = config.hidden_size

181 +

182 + moe_num_experts = config.moe_num_experts

183 + max_moe_num_experts = max(moe_num_experts)

184 +

185 + if self.tp_size > max_moe_num_experts:

186 + raise ValueError(

187 + f"Tensor parallel size {self.tp_size} is greater than "

188 + f"the number of experts {moe_num_experts}.")

189 +

190 + moe_layer_start_index = config.moe_layer_start_index

191 + text_moe_layer_start_index = moe_layer_start_index[0]

192 + vision_moe_layer_start_index = moe_layer_start_index[1]

193 + moe_layer_end_index = config.moe_layer_end_index

194 + moe_layer_end_index = getattr(

195 + config, "moe_layer_end_index",

196 + [config.num_hidden_layers - 1, config.num_hidden_layers - 1])

197 + text_moe_layer_end_index = moe_layer_end_index[0]

198 + vision_moe_layer_end_index = moe_layer_end_index[1]

199 +

200 + assert config.moe_num_experts[0] == config.moe_num_experts[1]

201 + self.e_score_correction_bias = nn.Parameter(

202 + torch.empty(2, config.moe_num_experts[0]))

203 +

204 + assert text_moe_layer_start_index <= text_moe_layer_end_index

205 +

206 + if layer_idx >= text_moe_layer_start_index and \

207 + layer_idx <= text_moe_layer_end_index:

208 + self.text_experts_gate = ReplicatedLinear(

209 + config.hidden_size,

210 + config.moe_num_experts[0],

211 + bias=False,

212 + quant_config=quant_config,

213 + prefix=f"{prefix}.text_experts_gate")

214 +

215 + self.text_experts = FusedMoE(

216 + num_experts=config.moe_num_experts[0],

217 + top_k=config.moe_k,

218 + hidden_size=config.hidden_size,

219 + intermediate_size=config.moe_intermediate_size[0],

220 + reduce_results=False,

221 + renormalize=True,

222 + quant_config=quant_config,

223 + e_score_correction_bias=self.e_score_correction_bias[0],

224 + prefix=f"{prefix}.text_experts")

225 + else:

226 + self.text_experts = Ernie4_5_VLMoeMLP(

227 + hidden_size=config.hidden_size,

228 + intermediate_size=config.intermediate_size,

229 + hidden_act=config.hidden_act,

230 + use_bias=getattr(config, 'use_bias', False),

231 + quant_config=quant_config,

232 + prefix=f"{prefix}.mlp")

233 +

234 + assert vision_moe_layer_start_index <= vision_moe_layer_end_index

235 + if layer_idx >= vision_moe_layer_start_index and \

236 + layer_idx <= vision_moe_layer_end_index:

237 + self.vision_experts_gate = ReplicatedLinear(

238 + config.hidden_size,

239 + config.moe_num_experts[1],

240 + bias=False,

241 + quant_config=quant_config,

242 + prefix=f"{prefix}.vision_experts_gate")

243 +

244 + self.vision_experts = FusedMoE(

245 + num_experts=config.moe_num_experts[1],

246 + top_k=config.moe_k,

247 + hidden_size=config.hidden_size,

248 + intermediate_size=config.moe_intermediate_size[1],

249 + reduce_results=False,

250 + renormalize=True,

251 + quant_config=quant_config,

252 + e_score_correction_bias=self.e_score_correction_bias[1],

253 + prefix=f"{prefix}.vision_experts")

254 + else:

255 + self.vision_experts = Ernie4_5_VLMoeMLP(

256 + hidden_size=config.hidden_size,

257 + intermediate_size=config.intermediate_size,

258 + hidden_act=config.hidden_act,

259 + use_bias=getattr(config, 'use_bias', False),

260 + quant_config=quant_config,

261 + prefix=f"{prefix}.mlp")

262 +

263 + if self.has_shared_experts:

264 + intermediate_size = (config.moe_intermediate_size[0] *

265 + config.moe_num_shared_experts)

266 + self.shared_experts = Ernie4_5_VLMoeMLP(

267 + hidden_size=config.hidden_size,

268 + intermediate_size=intermediate_size,

269 + hidden_act=config.hidden_act,

270 + quant_config=quant_config,

271 + prefix=f"{prefix}.shared_experts",

272 + reduce_results=self.text_experts.

273 + must_reduce_shared_expert_outputs())

274 +

275 + def forward(

276 + self,

277 + hidden_states: torch.Tensor,

278 + visual_token_mask: torch.Tensor,

279 + **kwargs: object,

280 + ) -> torch.Tensor:

281 +

282 + orig_shape = hidden_states.shape

283 + hidden_dim = hidden_states.shape[-1]

284 + hidden_states = hidden_states.view(-1, hidden_dim)

285 +

286 + if self.has_shared_experts:

287 + shared_output = self.shared_experts(hidden_states)

288 +

289 + if visual_token_mask is not None and visual_token_mask.any():

290 + # assert visual_token_mask.shape[0] != hidden_states.shape[0]

291 + visual_token_mask = visual_token_mask.repeat(

292 + 1, self.hidden_size).bool()

293 + text_token_mask = ~visual_token_mask

294 + final_hidden_states = torch.zeros_like(hidden_states)

295 +

296 + text_hidden_states = hidden_states[text_token_mask].reshape(

297 + -1, self.hidden_size)

298 + vision_hidden_states = hidden_states[visual_token_mask].reshape(

299 + -1, self.hidden_size)

300 +

301 + text_router_logits, _ = self.text_experts_gate(text_hidden_states)

302 + final_hidden_states[text_token_mask] = self.text_experts(

303 + hidden_states=text_hidden_states,

304 + router_logits=text_router_logits).flatten()

305 +

306 + vision_router_logits, _ = self.vision_experts_gate(

307 + vision_hidden_states)

308 + final_hidden_states[visual_token_mask] = self.vision_experts(

309 + hidden_states=vision_hidden_states,

310 + router_logits=vision_router_logits).flatten()

311 + else:

312 + # text modal input processing directly

313 + text_router_logits, _ = self.text_experts_gate(hidden_states)

314 +

315 + final_hidden_states = self.text_experts(

316 + hidden_states=hidden_states, router_logits=text_router_logits)

317 +

318 + if self.has_shared_experts and \

319 + shared_output is not None:

320 + final_hidden_states = final_hidden_states + shared_output

321 +

322 + if self.tp_size > 1:

323 + final_hidden_states = (

324 + self.text_experts.maybe_all_reduce_tensor_model_parallel(

325 + final_hidden_states))

326 +

327 + return final_hidden_states.view(orig_shape)

328 +

329 +

330 +class Ernie4_5_VLMoeDecoderLayer(nn.Module):

331 +

332 + def __init__(

333 + self,

334 + config: PretrainedConfig,

335 + cache_config: Optional[CacheConfig] = None,

336 + quant_config: Optional[QuantizationConfig] = None,

337 + prefix: str = "",

338 + ) -> None:

339 + super().__init__()

340 + self.hidden_size = config.hidden_size

341 + rope_theta = getattr(config, "rope_theta", 500000)

342 + rope_scaling = getattr(config, "rope_scaling", None)

343 + freq_allocation = getattr(config, "freq_allocation", 20)

344 + max_position_embeddings = getattr(config, "max_position_embeddings",

345 + 131072)

346 +

347 + self.self_attn = Ernie4_5_VLMoeAttention(

348 + hidden_size=self.hidden_size,

349 + num_heads=config.num_attention_heads,

350 + num_kv_heads=config.num_key_value_heads,

351 + head_dim=getattr(config, 'head_dim', None),

352 + rope_theta=rope_theta,

353 + rope_scaling=rope_scaling,

354 + freq_allocation=freq_allocation,

355 + max_position_embeddings=max_position_embeddings,

356 + rms_norm_eps=config.rms_norm_eps,

357 + qkv_bias=getattr(config, 'use_bias', False),

358 + cache_config=cache_config,

359 + quant_config=quant_config,

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

361 + )

362 +

363 + layer_idx = extract_layer_index(prefix)

364 + self.layer_idx = layer_idx

365 +

366 + # MoE

367 + moe_layer_start_index = config.moe_layer_start_index

368 + min_moe_layer_start_index = min(moe_layer_start_index)

369 + moe_layer_end_index = getattr(

370 + config, "moe_layer_end_index",

371 + [config.num_hidden_layers - 1, config.num_hidden_layers - 1])

372 + max_moe_layer_end_index = max(moe_layer_end_index)

373 + assert min_moe_layer_start_index <= max_moe_layer_end_index

374 + moe_num_experts = config.moe_num_experts

375 + max_moe_num_experts = max(moe_num_experts)

376 + moe_layer_interval = getattr(config, "moe_layer_interval", 1)

377 + use_moe = getattr(config, "use_moe", max_moe_num_experts > 0)

378 +

379 + if (use_moe and ((layer_idx + 1) % moe_layer_interval == 0)

380 + and layer_idx >= min_moe_layer_start_index

381 + and layer_idx <= max_moe_layer_end_index):

382 + self.mlp = Ernie4_5_VLMoeMoE(config=config,

383 + quant_config=quant_config,

384 + prefix=f"{prefix}.mlp")

385 + else:

386 + self.mlp = Ernie4_5_VLMoeMLP(

387 + hidden_size=config.hidden_size,

388 + intermediate_size=config.intermediate_size,

389 + hidden_act=config.hidden_act,

390 + use_bias=getattr(config, 'use_bias', False),

391 + quant_config=quant_config,

392 + prefix=f"{prefix}.mlp")

393 +

394 + self.input_layernorm = RMSNorm(config.hidden_size,

395 + eps=config.rms_norm_eps)

396 + self.post_attention_layernorm = RMSNorm(config.hidden_size,

397 + eps=config.rms_norm_eps)

398 +

399 + def forward(

400 + self,

401 + positions: torch.Tensor,

402 + hidden_states: torch.Tensor,

403 + residual: Optional[torch.Tensor],

404 + visual_token_mask: Optional[torch.Tensor],

405 + **kwargs: object,

406 + ) -> torch.Tensor:

407 +

408 + # Self Attention

409 + if residual is None:

410 + residual = hidden_states

411 + hidden_states = self.input_layernorm(hidden_states)

412 + else:

413 + hidden_states, residual = self.input_layernorm(

414 + hidden_states, residual)

415 +

416 + hidden_states = self.self_attn(

417 + positions=positions,

418 + hidden_states=hidden_states,

419 + )

420 +

421 + # Fully Connected

422 + hidden_states, residual = self.post_attention_layernorm(

423 + hidden_states, residual)

424 +

425 + if isinstance(self.mlp, Ernie4_5_VLMoeMoE):

426 + hidden_states = self.mlp(hidden_states, visual_token_mask,

427 + **kwargs)

428 + else:

429 + hidden_states = self.mlp(hidden_states)

430 +

431 + return hidden_states, residual

432 +

433 +

434 +# Since Ernie VL distinguishes between text experts and vision experts,

435 +# enabling torch.compile will cause errors.

436 +# @support_torch_compile(

437 +# dynamic_arg_dims={

438 +# "input_ids": 0,

439 +# "positions": -1,

440 +# "intermediate_tensors": 0,

441 +# "inputs_embeds": 0,

442 +# "visual_token_mask": 0,

443 +# })

444 +class Ernie4_5_VLMoeModel(nn.Module):

445 +

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

447 + super().__init__()

448 +

449 + config = vllm_config.model_config.hf_config

450 + cache_config = vllm_config.cache_config

451 + quant_config = vllm_config.quant_config

452 +

453 + self.padding_idx = config.pad_token_id

454 + self.vocab_size = config.vocab_size

455 + self.config = config

456 +

457 + self.im_patch_id = config.im_patch_id

458 +

459 + if get_pp_group().is_first_rank:

460 + self.embed_tokens = VocabParallelEmbedding(

461 + config.vocab_size,

462 + config.hidden_size,

463 + quant_config=quant_config,

464 + prefix=f"{prefix}.embed_tokens")

465 + else:

466 + self.embed_tokens = PPMissingLayer()

467 +

468 + self.start_layer, self.end_layer, self.layers = make_layers(

469 + config.num_hidden_layers,

470 + lambda prefix: Ernie4_5_VLMoeDecoderLayer(

471 + config=config,

472 + cache_config=cache_config,

473 + quant_config=quant_config,

474 + prefix=prefix),

475 + prefix=f"{prefix}.layers",

476 + )

477 +

478 + if get_pp_group().is_last_rank:

479 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

480 + else:

481 + self.norm = PPMissingLayer()

482 +

483 + self.make_empty_intermediate_tensors = (

484 + make_empty_intermediate_tensors_factory(

485 + ["hidden_states", "residual"], config.hidden_size))

486 +

487 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

488 + return self.embed_tokens(input_ids)

489 +

490 + def forward(

491 + self,

492 + input_ids: torch.Tensor,

493 + positions: torch.Tensor,

494 + intermediate_tensors: Optional[IntermediateTensors] = None,

495 + inputs_embeds: Optional[torch.Tensor] = None,

496 + visual_token_mask: Optional[torch.Tensor] = None,

497 + **kwargs: object,

498 + ) -> Union[torch.Tensor, IntermediateTensors]:

499 +

500 + if get_pp_group().is_first_rank:

501 + if inputs_embeds is not None:

502 + hidden_states = inputs_embeds

503 + else:

504 + hidden_states = self.get_input_embeddings(input_ids)

505 + residual = None

506 + else:

507 + assert intermediate_tensors is not None

508 + hidden_states = intermediate_tensors["hidden_states"]

509 + residual = intermediate_tensors["residual"]

510 +

511 + for i in range(self.start_layer, self.end_layer):

512 + layer = self.layers[i]

513 + hidden_states, residual = layer(positions, hidden_states, residual,

514 + visual_token_mask, **kwargs)

515 +

516 + if not get_pp_group().is_last_rank:

517 + return IntermediateTensors({

518 + "hidden_states": hidden_states,

519 + "residual": residual

520 + })

521 +

522 + hidden_states, _ = self.norm(hidden_states, residual)

523 +

524 + return hidden_states

525 +

526 +

527 +# only used as text backbone for ernie4.5-vl

528 +class Ernie4_5_VLMoeForCausalLM(nn.Module, SupportsPP):

529 + packed_modules_mapping = {

530 + "qkv_proj": [

531 + "q_proj",

532 + "k_proj",

533 + "v_proj",

534 + ],

535 + "gate_up_proj": [

536 + "gate_proj",

537 + "up_proj",

538 + ],

539 + }

540 +

541 + fall_back_to_pt_during_load = False

542 +

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

544 + super().__init__()

545 + config = vllm_config.model_config.hf_config

546 + quant_config = vllm_config.quant_config

547 + self.config = config

548 + self.quant_config = quant_config

549 + self.model = Ernie4_5_VLMoeModel(vllm_config=vllm_config,

550 + prefix=maybe_prefix(prefix, "model"))

551 +

552 + if get_pp_group().is_last_rank:

553 + self.lm_head = ParallelLMHead(config.vocab_size,

554 + config.hidden_size,

555 + quant_config=quant_config)

556 + else:

557 + self.lm_head = PPMissingLayer()

558 +

559 + if self.config.tie_word_embeddings:

560 + self.lm_head.weight = self.model.embed_tokens.weight

561 + self.logits_processor = LogitsProcessor(config.vocab_size)

562 + self.make_empty_intermediate_tensors = (

563 + self.model.make_empty_intermediate_tensors)

564 +

565 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

566 + return self.model.get_input_embeddings(input_ids)

567 +

568 + def forward(

569 + self,

570 + input_ids: torch.Tensor,

571 + positions: torch.Tensor,

572 + intermediate_tensors: Optional[IntermediateTensors] = None,

573 + inputs_embeds: Optional[torch.Tensor] = None,

574 + **kwargs: object,

575 + ) -> Union[torch.Tensor, IntermediateTensors]:

576 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

577 + inputs_embeds, **kwargs)

578 + return hidden_states

579 +

580 + def compute_logits(

581 + self,

582 + hidden_states: torch.Tensor,

583 + sampling_metadata: SamplingMetadata,

584 + ) -> Optional[torch.Tensor]:

585 + logits = self.logits_processor(self.lm_head, hidden_states,

586 + sampling_metadata)

587 + return logits

588 +

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

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

591 + stacked_params_mapping = [

592 + # (param_name, shard_name, shard_id)

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

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

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

596 + ("gate_up_proj", "gate_proj", 0),

597 + ("gate_up_proj", "up_proj", 1),

598 + ]

599 +

600 + # Params for weights, fp8 weight scales, fp8 activation scales

601 + # (param_name, weight_name, expert_id, shard_id)

602 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

603 + ckpt_gate_proj_name="gate_proj",

604 + ckpt_down_proj_name="down_proj",

605 + ckpt_up_proj_name="up_proj",

606 + num_experts=max(self.config.moe_num_experts))

607 +

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

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

610 + for name, loaded_weight in weights:

611 + if self.config.tie_word_embeddings and name.endswith(

612 + "lm_head.weight"):

613 + loaded_params.add("lm_head.weight")

614 + continue

615 + # MTP will be supported soon.

616 + if "mtp" in name or \

617 + "vision_model" in name or \

618 + "resampler_model" in name:

619 + continue

620 +

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

622 + # Skip non-stacked layers and experts (experts handled below).

623 + if weight_name not in name:

624 + continue

625 +

626 + if (("mlp.experts." in name) and name not in params_dict):

627 + continue

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

629 + # Skip loading extra bias for GPTQ models.

630 + if ((name.endswith(".bias") or name.endswith("_bias"))

631 + and name not in params_dict):

632 + continue

633 + # Skip layers on other devices.

634 + if is_pp_missing_parameter(name, self):

635 + continue

636 +

637 + param = params_dict[name]

638 + weight_loader = param.weight_loader

639 + weight_loader(param, loaded_weight, shard_id)

640 + break

641 + else:

642 + # Distinguish between vision experts and text experts

643 + if "mlp.experts" in name:

644 + moe_offset = int(name.split(".")[-3])

645 + vision_expert_start_idx = self.config.moe_num_experts[0]

646 + is_text_expert = \

647 + moe_offset <= vision_expert_start_idx - 1

648 + if is_text_expert:

649 + name = name.replace(".experts.", ".text_experts.")

650 + else:

651 + name = name.replace(

652 + f".experts.{moe_offset}",

653 + f".vision_experts.{moe_offset-vision_expert_start_idx}"

654 + )

655 +

656 + for mapping in expert_params_mapping:

657 + param_name, weight_name, expert_id, shard_id = mapping

658 +

659 + if weight_name not in name:

660 + continue

661 +

662 + # Distinguish between vision experts and text experts

663 + moe_offset = int(name.split(".")[-3])

664 + is_text_expert = \

665 + moe_offset <= self.config.moe_num_experts[0] - 1

666 +

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

668 + if is_text_expert:

669 + name = name.replace(".experts.", ".text_experts.")

670 + else:

671 + name = name.replace(".experts.", ".vision_experts.")

672 +

673 + # Skip layers on other devices.

674 + if is_pp_missing_parameter(name, self):

675 + continue

676 +

677 + # Skip loading extra bias for GPTQ models.

678 + if ((name.endswith(".bias") or name.endswith("_bias"))

679 + and name not in params_dict):

680 + continue

681 + param = params_dict[name]

682 +

683 + weight_loader = param.weight_loader

684 + weight_loader(param,

685 + loaded_weight,

686 + name,

687 + shard_id=shard_id,

688 + expert_id=expert_id)

689 + break

690 + else:

691 + # Distinguish between vision expert gate

692 + # and text expert gate

693 + if name.endswith("mlp.gate.weight"):

694 + name = name.replace("gate.weight",

695 + "text_experts_gate.weight")

696 + loaded_weight = loaded_weight.T

697 + elif name.endswith("mlp.gate.weight_1"):

698 + name = name.replace("gate.weight_1",

699 + "vision_experts_gate.weight")

700 + loaded_weight = loaded_weight.T

701 +

702 + if "e_score_correction_bias" in name:

703 + name = name.replace(".moe_statics.", ".")

704 +

705 + # Skip loading extra bias for GPTQ models.

706 + if ((name.endswith(".bias") or name.endswith("_bias"))

707 + and name not in params_dict):

708 + continue

709 + # Skip layers on other devices.

710 + if is_pp_missing_parameter(name, self):

711 + continue

712 + # Remapping the name of FP8 kv-scale.

713 + name = maybe_remap_kv_scale_name(name, params_dict)

714 + if name is None:

715 + continue

716 +

717 + param = params_dict[name]

718 +

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

720 + default_weight_loader)

721 + weight_loader(param, loaded_weight)

722 + loaded_params.add(name)

723 + return loaded_params

vllm/model_executor/models/registry.py

@@ -206,6 +206,7 @@ _MULTIMODAL_MODELS = {

206 "ChameleonForConditionalGeneration": ("chameleon", "ChameleonForConditionalGeneration"), # noqa: E501

207 "Cohere2VisionForConditionalGeneration": ("cohere2_vision", "Cohere2VisionForConditionalGeneration"), # noqa: E501

208 "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),

209 + "Ernie4_5_VLMoeForConditionalGeneration": ("ernie45_vl", "Ernie4_5_VLMoeForConditionalGeneration"), # noqa: E501

210 "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"),

211 "Gemma3ForConditionalGeneration": ("gemma3_mm", "Gemma3ForConditionalGeneration"), # noqa: E501

212 "Gemma3nForConditionalGeneration": ("gemma3n_mm", "Gemma3nForConditionalGeneration"), # noqa: E501