[Model] Add Ernie4.5 VL model support (#15679)
Signed-off-by: CSWYF3634076 <wangyafeng@baidu.com> Signed-off-by: wangyafeng <wangyafeng@baidu.com>
This commit is contained in:
@@ -46,6 +46,7 @@ in the GitHub search bar.
|
|||||||
| **DotsVLM-OCR** | `rednote-hilab/dots.ocr` | Specialized OCR variant of DotsVLM optimized for optical character recognition tasks with enhanced text extraction and document understanding capabilities. | Don't use `--trust-remote-code` |
|
| **DotsVLM-OCR** | `rednote-hilab/dots.ocr` | Specialized OCR variant of DotsVLM optimized for optical character recognition tasks with enhanced text extraction and document understanding capabilities. | Don't use `--trust-remote-code` |
|
||||||
| **NVILA** (8B, 15B, Lite-2B, Lite-8B, Lite-15B) | `Efficient-Large-Model/NVILA-8B` | `chatml` | NVILA explores the full stack efficiency of multi-modal design, achieving cheaper training, faster deployment and better performance. |
|
| **NVILA** (8B, 15B, Lite-2B, Lite-8B, Lite-15B) | `Efficient-Large-Model/NVILA-8B` | `chatml` | NVILA explores the full stack efficiency of multi-modal design, achieving cheaper training, faster deployment and better performance. |
|
||||||
| **NVIDIA Nemotron Nano 2.0 VL** | `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` | NVIDIA Nemotron Nano v2 VL enables multi-image reasoning and video understanding, along with strong document intelligence, visual Q&A and summarization capabilities. It builds on Nemotron Nano V2, a hybrid Mamba-Transformer LLM, in order to achieve higher inference throughput in long document and video scenarios. | Use `--trust-remote-code`. You may need to adjust `--max-mamba-cache-size` [default is 512] to fit memory constraints. |
|
| **NVIDIA Nemotron Nano 2.0 VL** | `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` | NVIDIA Nemotron Nano v2 VL enables multi-image reasoning and video understanding, along with strong document intelligence, visual Q&A and summarization capabilities. It builds on Nemotron Nano V2, a hybrid Mamba-Transformer LLM, in order to achieve higher inference throughput in long document and video scenarios. | Use `--trust-remote-code`. You may need to adjust `--max-mamba-cache-size` [default is 512] to fit memory constraints. |
|
||||||
|
| **Ernie4.5-VL** | `baidu/ERNIE-4.5-VL-28B-A3B-PT` | Baidu's vision-language models(28B,424B). Support image and video comprehension, and also support thinking. | |
|
||||||
| **JetVLM** | | JetVLM is an vision-language model designed for high-performance multimodal understanding and generation tasks built upon Jet-Nemotron. | Coming soon |
|
| **JetVLM** | | JetVLM is an vision-language model designed for high-performance multimodal understanding and generation tasks built upon Jet-Nemotron. | Coming soon |
|
||||||
|
|
||||||
## Video Input Support
|
## Video Input Support
|
||||||
|
|||||||
@@ -1123,6 +1123,7 @@ def is_generation_model(model_architectures: List[str], is_embedding: bool = Fal
|
|||||||
multimodal_model_archs = [
|
multimodal_model_archs = [
|
||||||
"CLIPModel",
|
"CLIPModel",
|
||||||
"DeepseekVL2ForCausalLM",
|
"DeepseekVL2ForCausalLM",
|
||||||
|
"Ernie4_5_VLMoeForConditionalGeneration",
|
||||||
"Gemma3ForConditionalGeneration",
|
"Gemma3ForConditionalGeneration",
|
||||||
"Gemma3nForConditionalGeneration",
|
"Gemma3nForConditionalGeneration",
|
||||||
"Glm4vForConditionalGeneration",
|
"Glm4vForConditionalGeneration",
|
||||||
|
|||||||
@@ -2285,6 +2285,177 @@ class MRotaryEmbedding(RotaryEmbedding):
|
|||||||
|
|
||||||
return position_ids, mrope_position_deltas
|
return position_ids, mrope_position_deltas
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def get_rope_index_ernie45(
|
||||||
|
input_ids: torch.Tensor,
|
||||||
|
hf_config: Any,
|
||||||
|
image_grid_thw: Union[list[list[int]], torch.Tensor],
|
||||||
|
video_grid_thw: Union[list[list[int]], torch.Tensor],
|
||||||
|
**kwargs,
|
||||||
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||||
|
"""Get mrope input positions and delta value for Ernie VL."""
|
||||||
|
|
||||||
|
image_token_id = hf_config.im_patch_id
|
||||||
|
video_start_token_id = hf_config.video_start_token_id
|
||||||
|
video_end_token_id = hf_config.video_end_token_id
|
||||||
|
spatial_conv_size = hf_config.spatial_conv_size
|
||||||
|
temporal_conv_size = hf_config.temporal_conv_size
|
||||||
|
|
||||||
|
mrope_position_deltas = []
|
||||||
|
if input_ids is not None and (
|
||||||
|
image_grid_thw is not None or video_grid_thw is not None
|
||||||
|
):
|
||||||
|
total_input_ids = input_ids
|
||||||
|
position_ids = torch.ones(
|
||||||
|
3,
|
||||||
|
input_ids.shape[0],
|
||||||
|
input_ids.shape[1],
|
||||||
|
dtype=input_ids.dtype,
|
||||||
|
device=input_ids.device,
|
||||||
|
)
|
||||||
|
image_index, video_index = 0, 0
|
||||||
|
for i, input_ids in enumerate(total_input_ids):
|
||||||
|
input_tokens = input_ids.tolist()
|
||||||
|
|
||||||
|
input_token_type = []
|
||||||
|
video_check_flg = False
|
||||||
|
for token in input_tokens:
|
||||||
|
if token == video_start_token_id:
|
||||||
|
video_check_flg = True
|
||||||
|
elif token == video_end_token_id:
|
||||||
|
video_check_flg = False
|
||||||
|
|
||||||
|
if token == image_token_id and not video_check_flg:
|
||||||
|
input_token_type.append("image")
|
||||||
|
elif token == image_token_id and video_check_flg:
|
||||||
|
input_token_type.append("video")
|
||||||
|
else:
|
||||||
|
input_token_type.append("text")
|
||||||
|
|
||||||
|
input_type_group = []
|
||||||
|
for key, group in itertools.groupby(
|
||||||
|
enumerate(input_token_type), lambda x: x[1]
|
||||||
|
):
|
||||||
|
group = list(group)
|
||||||
|
start_index = group[0][0]
|
||||||
|
end_index = group[-1][0] + 1
|
||||||
|
input_type_group.append((key, start_index, end_index))
|
||||||
|
|
||||||
|
llm_pos_ids_list = []
|
||||||
|
video_frame_num = 1
|
||||||
|
for modality_type, start_idx, end_idx in input_type_group:
|
||||||
|
st_idx = (
|
||||||
|
llm_pos_ids_list[-1].max() + 1
|
||||||
|
if len(llm_pos_ids_list) > 0
|
||||||
|
else 0
|
||||||
|
)
|
||||||
|
|
||||||
|
if modality_type == "image":
|
||||||
|
t, h, w = (
|
||||||
|
image_grid_thw[image_index][0],
|
||||||
|
image_grid_thw[image_index][1],
|
||||||
|
image_grid_thw[image_index][2],
|
||||||
|
)
|
||||||
|
llm_grid_t, llm_grid_h, llm_grid_w = (
|
||||||
|
t.item(),
|
||||||
|
h.item() // spatial_conv_size,
|
||||||
|
w.item() // spatial_conv_size,
|
||||||
|
)
|
||||||
|
|
||||||
|
t_index = (
|
||||||
|
torch.arange(llm_grid_t)
|
||||||
|
.view(-1, 1)
|
||||||
|
.expand(-1, llm_grid_h * llm_grid_w)
|
||||||
|
.flatten()
|
||||||
|
)
|
||||||
|
h_index = (
|
||||||
|
torch.arange(llm_grid_h)
|
||||||
|
.view(1, -1, 1)
|
||||||
|
.expand(llm_grid_t, -1, llm_grid_w)
|
||||||
|
.flatten()
|
||||||
|
)
|
||||||
|
w_index = (
|
||||||
|
torch.arange(llm_grid_w)
|
||||||
|
.view(1, 1, -1)
|
||||||
|
.expand(llm_grid_t, llm_grid_h, -1)
|
||||||
|
.flatten()
|
||||||
|
)
|
||||||
|
llm_pos_ids_list.append(
|
||||||
|
torch.stack([t_index, h_index, w_index]) + st_idx
|
||||||
|
)
|
||||||
|
|
||||||
|
image_index += 1
|
||||||
|
video_frame_num = 1
|
||||||
|
|
||||||
|
elif modality_type == "video":
|
||||||
|
t, h, w = (
|
||||||
|
video_grid_thw[video_index][0],
|
||||||
|
video_grid_thw[video_index][1],
|
||||||
|
video_grid_thw[video_index][2],
|
||||||
|
)
|
||||||
|
|
||||||
|
llm_grid_t, llm_grid_h, llm_grid_w = (
|
||||||
|
t.item() // temporal_conv_size,
|
||||||
|
h.item() // spatial_conv_size,
|
||||||
|
w.item() // spatial_conv_size,
|
||||||
|
)
|
||||||
|
|
||||||
|
for t_idx in range(llm_grid_t):
|
||||||
|
t_index = (
|
||||||
|
torch.tensor(t_idx)
|
||||||
|
.view(-1, 1)
|
||||||
|
.expand(-1, llm_grid_h * llm_grid_w)
|
||||||
|
.flatten()
|
||||||
|
)
|
||||||
|
|
||||||
|
h_index = (
|
||||||
|
torch.arange(llm_grid_h)
|
||||||
|
.view(1, -1, 1)
|
||||||
|
.expand(1, -1, llm_grid_w)
|
||||||
|
.flatten()
|
||||||
|
)
|
||||||
|
w_index = (
|
||||||
|
torch.arange(llm_grid_w)
|
||||||
|
.view(1, 1, -1)
|
||||||
|
.expand(1, llm_grid_h, -1)
|
||||||
|
.flatten()
|
||||||
|
)
|
||||||
|
llm_pos_ids_list.append(
|
||||||
|
torch.stack([t_index, h_index, w_index]) + st_idx
|
||||||
|
)
|
||||||
|
|
||||||
|
video_index += 1
|
||||||
|
video_frame_num += 1
|
||||||
|
|
||||||
|
else:
|
||||||
|
text_len = end_idx - start_idx
|
||||||
|
llm_pos_ids_list.append(
|
||||||
|
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
|
||||||
|
)
|
||||||
|
|
||||||
|
video_frame_num = 1
|
||||||
|
|
||||||
|
llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
|
||||||
|
position_ids[..., i, :] = llm_positions.to(position_ids.device)
|
||||||
|
mrope_position_deltas.append(
|
||||||
|
llm_positions.max() + 1 - len(total_input_ids[i])
|
||||||
|
)
|
||||||
|
mrope_position_deltas = torch.tensor(
|
||||||
|
mrope_position_deltas, device=input_ids.device
|
||||||
|
).unsqueeze(1)
|
||||||
|
return position_ids, mrope_position_deltas
|
||||||
|
else:
|
||||||
|
s = input_ids.shape[1]
|
||||||
|
position_ids = torch.arange(s)
|
||||||
|
position_ids = (
|
||||||
|
position_ids.unsqueeze(0).expand(3, -1, -1).to(input_ids.device)
|
||||||
|
)
|
||||||
|
max_position_ids = position_ids.max(0, keepdim=False)[0].max(
|
||||||
|
-1, keepdim=True
|
||||||
|
)[0]
|
||||||
|
mrope_position_deltas = max_position_ids + 1 - s
|
||||||
|
return position_ids, mrope_position_deltas
|
||||||
|
|
||||||
# For qwen3-omni
|
# For qwen3-omni
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _get_feat_extract_output_lengths(input_lengths):
|
def _get_feat_extract_output_lengths(input_lengths):
|
||||||
@@ -2324,6 +2495,91 @@ class MRotaryEmbedding(RotaryEmbedding):
|
|||||||
return llm_pos_ids
|
return llm_pos_ids
|
||||||
|
|
||||||
|
|
||||||
|
class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):
|
||||||
|
"""3D rotary positional embedding. [h w h w h w h w... t t t...]"""
|
||||||
|
|
||||||
|
def forward_native( # type: ignore[override]
|
||||||
|
self,
|
||||||
|
positions: torch.Tensor,
|
||||||
|
query: torch.Tensor,
|
||||||
|
key: torch.Tensor | None = None,
|
||||||
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||||
|
assert positions.ndim == 1 or positions.ndim == 2
|
||||||
|
assert key is not None
|
||||||
|
|
||||||
|
num_tokens = positions.shape[-1]
|
||||||
|
cos_sin = self.cos_sin_cache[positions]
|
||||||
|
cos, sin = cos_sin.chunk(2, dim=-1)
|
||||||
|
if positions.ndim == 2:
|
||||||
|
assert self.mrope_section
|
||||||
|
|
||||||
|
section_h = self.mrope_section[0] # 22
|
||||||
|
section_w = self.mrope_section[1] # 22
|
||||||
|
section_t = self.mrope_section[2] # 20
|
||||||
|
assert section_h == section_w
|
||||||
|
# Split according to [h w h w h w h w... t t t...]
|
||||||
|
section_cos_t = cos[..., -section_t:]
|
||||||
|
section_cos_h = cos[..., : section_h + section_w : 2]
|
||||||
|
section_cos_w = cos[..., 1 : section_h + section_w : 2]
|
||||||
|
|
||||||
|
cos_t, cos_h, cos_w = section_cos_t[0], section_cos_h[1], section_cos_w[2]
|
||||||
|
cos_hw = torch.stack([cos_h, cos_w], dim=-1).reshape(
|
||||||
|
cos_h.shape[:-1] + (cos_h.shape[-1] * 2,)
|
||||||
|
)
|
||||||
|
cos = torch.cat([cos_hw, cos_t], dim=-1)
|
||||||
|
|
||||||
|
section_sin_t = sin[..., -section_t:]
|
||||||
|
section_sin_h = sin[..., : section_h + section_w : 2]
|
||||||
|
section_sin_w = sin[..., 1 : section_h + section_w : 2]
|
||||||
|
|
||||||
|
sin_t, sin_h, sin_w = section_sin_t[0], section_sin_h[1], section_sin_w[2]
|
||||||
|
sin_hw = torch.stack([sin_h, sin_w], dim=-1).reshape(
|
||||||
|
sin_h.shape[:-1] + (sin_h.shape[-1] * 2,)
|
||||||
|
)
|
||||||
|
sin = torch.cat([sin_hw, sin_t], dim=-1)
|
||||||
|
|
||||||
|
query_shape = query.shape
|
||||||
|
query = query.view(num_tokens, -1, self.head_size)
|
||||||
|
query_rot = query[..., : self.rotary_dim]
|
||||||
|
query_pass = query[..., self.rotary_dim :]
|
||||||
|
query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
|
||||||
|
query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
|
||||||
|
|
||||||
|
key_shape = key.shape
|
||||||
|
key = key.view(num_tokens, -1, self.head_size)
|
||||||
|
key_rot = key[..., : self.rotary_dim]
|
||||||
|
key_pass = key[..., self.rotary_dim :]
|
||||||
|
key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
|
||||||
|
key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
|
||||||
|
return query, key
|
||||||
|
|
||||||
|
def forward_cuda( # type: ignore[override]
|
||||||
|
self,
|
||||||
|
positions: torch.Tensor,
|
||||||
|
query: torch.Tensor,
|
||||||
|
key: torch.Tensor | None = None,
|
||||||
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||||
|
return self.forward_native(positions, query, key)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
positions: torch.Tensor,
|
||||||
|
query: torch.Tensor,
|
||||||
|
key: torch.Tensor,
|
||||||
|
fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
|
||||||
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||||
|
"""Forward pass with optional Triton kernel acceleration.
|
||||||
|
Args:
|
||||||
|
positions:
|
||||||
|
[num_tokens,] (text only) or
|
||||||
|
[3, num_tokens] (T/H/W positions with multimodal inputs)
|
||||||
|
query: [num_tokens, num_heads * head_size]
|
||||||
|
key: [num_tokens, num_kv_heads * head_size]
|
||||||
|
"""
|
||||||
|
assert positions.ndim == 1 or positions.ndim == 2
|
||||||
|
return self.forward_native(positions, query, key)
|
||||||
|
|
||||||
|
|
||||||
class DualChunkRotaryEmbedding(MultiPlatformOp):
|
class DualChunkRotaryEmbedding(MultiPlatformOp):
|
||||||
"""Rotary positional embedding for Dual Chunk Attention."""
|
"""Rotary positional embedding for Dual Chunk Attention."""
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,552 @@
|
|||||||
|
# Copyright 2023-2025 SGLang Team
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
# ==============================================================================
|
||||||
|
|
||||||
|
""" Inference-only Ernie4.5 VL model compatible with baidu/ERNIE-4.5-VL-*-PT weights. """
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from itertools import islice
|
||||||
|
from typing import Any, Dict, Optional, Tuple, Union
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch import nn
|
||||||
|
from transformers import PretrainedConfig
|
||||||
|
|
||||||
|
from sglang.srt.distributed import (
|
||||||
|
get_pp_group,
|
||||||
|
get_tensor_model_parallel_world_size,
|
||||||
|
tensor_model_parallel_all_reduce,
|
||||||
|
)
|
||||||
|
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
|
||||||
|
from sglang.srt.layers.layernorm import RMSNorm
|
||||||
|
from sglang.srt.layers.linear import (
|
||||||
|
QKVParallelLinear,
|
||||||
|
ReplicatedLinear,
|
||||||
|
RowParallelLinear,
|
||||||
|
)
|
||||||
|
from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
|
||||||
|
from sglang.srt.layers.moe.topk import TopK
|
||||||
|
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||||
|
from sglang.srt.layers.radix_attention import RadixAttention
|
||||||
|
from sglang.srt.layers.rotary_embedding import Ernie4_5_VLRotaryEmbedding
|
||||||
|
from sglang.srt.layers.utils import PPMissingLayer
|
||||||
|
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
|
||||||
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
|
||||||
|
from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Ernie4_5_VLMoeMLP
|
||||||
|
from sglang.srt.utils import add_prefix, make_layers
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class Ernie4_5_VLMoeAttention(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config: PretrainedConfig,
|
||||||
|
hidden_size: int,
|
||||||
|
num_heads: int,
|
||||||
|
num_kv_heads: int,
|
||||||
|
layer_id: int = 0,
|
||||||
|
rope_theta: float = 10000,
|
||||||
|
rope_scaling: Optional[Dict[str, Any]] = None,
|
||||||
|
rope_is_neox_style: bool = True,
|
||||||
|
freq_allocation: int = 20,
|
||||||
|
max_position_embeddings: int = 8192,
|
||||||
|
quant_config: Optional[QuantizationConfig] = None,
|
||||||
|
prefix: str = "",
|
||||||
|
bias: bool = False,
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.hidden_size = hidden_size
|
||||||
|
tp_size = get_tensor_model_parallel_world_size()
|
||||||
|
self.total_num_heads = num_heads
|
||||||
|
assert self.total_num_heads % tp_size == 0
|
||||||
|
self.num_heads = self.total_num_heads // tp_size
|
||||||
|
self.total_num_kv_heads = num_kv_heads
|
||||||
|
if self.total_num_kv_heads >= tp_size:
|
||||||
|
# Number of KV heads is greater than TP size, so we partition
|
||||||
|
# the KV heads across multiple tensor parallel GPUs.
|
||||||
|
assert self.total_num_kv_heads % tp_size == 0
|
||||||
|
else:
|
||||||
|
# Number of KV heads is less than TP size, so we replicate
|
||||||
|
# the KV heads across multiple tensor parallel GPUs.
|
||||||
|
assert tp_size % self.total_num_kv_heads == 0
|
||||||
|
self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
|
||||||
|
# MistralConfig has an optional head_dim introduced by Mistral-Nemo
|
||||||
|
self.head_dim = getattr(
|
||||||
|
config, "head_dim", self.hidden_size // self.total_num_heads
|
||||||
|
)
|
||||||
|
partial_rotary_factor = getattr(config, "partial_rotary_factor", 1)
|
||||||
|
self.rotary_dim = int(partial_rotary_factor * self.head_dim)
|
||||||
|
self.q_size = self.num_heads * self.head_dim
|
||||||
|
self.kv_size = self.num_kv_heads * self.head_dim
|
||||||
|
self.scaling = self.head_dim**-0.5
|
||||||
|
self.rope_theta = rope_theta
|
||||||
|
self.max_position_embeddings = max_position_embeddings
|
||||||
|
|
||||||
|
self.qkv_proj = QKVParallelLinear(
|
||||||
|
hidden_size,
|
||||||
|
self.head_dim,
|
||||||
|
self.total_num_heads,
|
||||||
|
self.total_num_kv_heads,
|
||||||
|
bias=bias,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("qkv_proj", prefix),
|
||||||
|
)
|
||||||
|
self.o_proj = RowParallelLinear(
|
||||||
|
self.total_num_heads * self.head_dim,
|
||||||
|
hidden_size,
|
||||||
|
bias=bias,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("o_proj", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
# 3D rope
|
||||||
|
t_rope = freq_allocation
|
||||||
|
h_rope = (self.head_dim // 2 - freq_allocation) // 2
|
||||||
|
w_rope = (self.head_dim // 2 - freq_allocation) // 2
|
||||||
|
|
||||||
|
self.rotary_emb = Ernie4_5_VLRotaryEmbedding(
|
||||||
|
head_size=self.head_dim,
|
||||||
|
rotary_dim=self.head_dim,
|
||||||
|
max_position_embeddings=max_position_embeddings,
|
||||||
|
base=rope_theta,
|
||||||
|
is_neox_style=False,
|
||||||
|
dtype=torch.get_default_dtype(),
|
||||||
|
mrope_section=[h_rope, w_rope, t_rope],
|
||||||
|
)
|
||||||
|
self.attn = RadixAttention(
|
||||||
|
self.num_heads,
|
||||||
|
self.head_dim,
|
||||||
|
self.scaling,
|
||||||
|
num_kv_heads=self.num_kv_heads,
|
||||||
|
layer_id=layer_id,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("attn", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
positions: torch.Tensor,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
forward_batch: ForwardBatch,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
qkv, _ = self.qkv_proj(hidden_states)
|
||||||
|
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
||||||
|
q, k = self.rotary_emb(positions, q, k)
|
||||||
|
attn_output = self.attn(q, k, v, forward_batch)
|
||||||
|
output, _ = self.o_proj(attn_output)
|
||||||
|
return output
|
||||||
|
|
||||||
|
|
||||||
|
class Ernie4_5_VLMoeMoE(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config: PretrainedConfig,
|
||||||
|
layer_id: int,
|
||||||
|
quant_config: Optional[QuantizationConfig] = None,
|
||||||
|
prefix: str = "",
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.layer_id = layer_id
|
||||||
|
self.tp_size = get_tensor_model_parallel_world_size()
|
||||||
|
self.moe_num_shared_experts = getattr(config, "moe_num_shared_experts", 0)
|
||||||
|
self.hidden_size = config.hidden_size
|
||||||
|
|
||||||
|
moe_num_experts = config.moe_num_experts
|
||||||
|
max_moe_num_experts = max(moe_num_experts)
|
||||||
|
|
||||||
|
if self.tp_size > max_moe_num_experts:
|
||||||
|
raise ValueError(
|
||||||
|
f"Tensor parallel size {self.tp_size} is greater than "
|
||||||
|
f"the number of experts {moe_num_experts}."
|
||||||
|
)
|
||||||
|
|
||||||
|
moe_layer_start_index = config.moe_layer_start_index
|
||||||
|
text_moe_layer_start_index = moe_layer_start_index[0]
|
||||||
|
vision_moe_layer_start_index = moe_layer_start_index[1]
|
||||||
|
moe_layer_end_index = config.moe_layer_end_index
|
||||||
|
moe_layer_end_index = getattr(
|
||||||
|
config,
|
||||||
|
"moe_layer_end_index",
|
||||||
|
[config.num_hidden_layers - 1, config.num_hidden_layers - 1],
|
||||||
|
)
|
||||||
|
text_moe_layer_end_index = moe_layer_end_index[0]
|
||||||
|
vision_moe_layer_end_index = moe_layer_end_index[1]
|
||||||
|
|
||||||
|
assert config.moe_num_experts[0] == config.moe_num_experts[1]
|
||||||
|
self.e_score_correction_bias = nn.Parameter(
|
||||||
|
torch.empty(2, config.moe_num_experts[0], dtype=torch.float32)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert text_moe_layer_start_index <= text_moe_layer_end_index
|
||||||
|
|
||||||
|
if (
|
||||||
|
layer_id >= text_moe_layer_start_index
|
||||||
|
and layer_id <= text_moe_layer_end_index
|
||||||
|
):
|
||||||
|
self.text_experts_gate = ReplicatedLinear(
|
||||||
|
config.hidden_size,
|
||||||
|
config.moe_num_experts[0],
|
||||||
|
bias=False,
|
||||||
|
params_dtype=torch.float32,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("text_experts_gate", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.text_experts_topk = TopK(
|
||||||
|
top_k=config.moe_k,
|
||||||
|
renormalize=True,
|
||||||
|
use_grouped_topk=False,
|
||||||
|
correction_bias=self.e_score_correction_bias[0],
|
||||||
|
)
|
||||||
|
|
||||||
|
self.text_experts = get_moe_impl_class(quant_config)(
|
||||||
|
num_experts=config.moe_num_experts[0],
|
||||||
|
top_k=config.moe_k,
|
||||||
|
hidden_size=config.hidden_size,
|
||||||
|
intermediate_size=config.moe_intermediate_size[0],
|
||||||
|
layer_id=self.layer_id,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("text_experts", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert vision_moe_layer_start_index <= vision_moe_layer_end_index
|
||||||
|
if (
|
||||||
|
layer_id >= vision_moe_layer_start_index
|
||||||
|
and layer_id <= vision_moe_layer_end_index
|
||||||
|
):
|
||||||
|
|
||||||
|
self.vision_experts_gate = ReplicatedLinear(
|
||||||
|
config.hidden_size,
|
||||||
|
config.moe_num_experts[1],
|
||||||
|
bias=False,
|
||||||
|
params_dtype=torch.float32,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("vision_experts_gate", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.vision_experts_topk = TopK(
|
||||||
|
top_k=config.moe_k,
|
||||||
|
renormalize=True,
|
||||||
|
use_grouped_topk=False,
|
||||||
|
correction_bias=self.e_score_correction_bias[1],
|
||||||
|
)
|
||||||
|
|
||||||
|
self.vision_experts = get_moe_impl_class(quant_config)(
|
||||||
|
num_experts=config.moe_num_experts[1],
|
||||||
|
top_k=config.moe_k,
|
||||||
|
hidden_size=config.hidden_size,
|
||||||
|
intermediate_size=config.moe_intermediate_size[1],
|
||||||
|
layer_id=self.layer_id,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("vision_experts", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
if self.moe_num_shared_experts > 0:
|
||||||
|
intermediate_size = (
|
||||||
|
config.moe_intermediate_size[0] * config.moe_num_shared_experts
|
||||||
|
)
|
||||||
|
self.shared_experts = Ernie4_5_VLMoeMLP(
|
||||||
|
hidden_size=config.hidden_size,
|
||||||
|
intermediate_size=intermediate_size,
|
||||||
|
hidden_act=config.hidden_act,
|
||||||
|
quant_config=quant_config,
|
||||||
|
reduce_results=False,
|
||||||
|
prefix=add_prefix("shared_experts", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
visual_token_mask: torch.Tensor,
|
||||||
|
**kwargs: object,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
shared_output = (
|
||||||
|
self.shared_experts(hidden_states)
|
||||||
|
if self.moe_num_shared_experts > 0
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
|
||||||
|
orig_shape = hidden_states.shape
|
||||||
|
hidden_dim = hidden_states.shape[-1]
|
||||||
|
hidden_states = hidden_states.view(-1, hidden_dim)
|
||||||
|
|
||||||
|
capturing = torch.cuda.is_current_stream_capturing()
|
||||||
|
|
||||||
|
if visual_token_mask is not None and not capturing:
|
||||||
|
all_visual = visual_token_mask.all()
|
||||||
|
any_visual = visual_token_mask.any()
|
||||||
|
else:
|
||||||
|
# During CUDA Graph capture, all set false
|
||||||
|
all_visual = False
|
||||||
|
any_visual = False
|
||||||
|
|
||||||
|
if all_visual:
|
||||||
|
# vision modal input processing directly
|
||||||
|
vision_router_logits, _ = self.vision_experts_gate(
|
||||||
|
hidden_states.to(dtype=torch.float32)
|
||||||
|
)
|
||||||
|
vision_topk_output = self.vision_experts_topk(
|
||||||
|
hidden_states, vision_router_logits
|
||||||
|
)
|
||||||
|
final_hidden_states = self.vision_experts(
|
||||||
|
hidden_states=hidden_states, topk_output=vision_topk_output
|
||||||
|
)
|
||||||
|
elif any_visual:
|
||||||
|
visual_token_mask = visual_token_mask.repeat(1, self.hidden_size).bool()
|
||||||
|
text_token_mask = ~visual_token_mask
|
||||||
|
final_hidden_states = torch.zeros_like(hidden_states)
|
||||||
|
|
||||||
|
text_hidden_states = hidden_states[text_token_mask].reshape(
|
||||||
|
-1, self.hidden_size
|
||||||
|
)
|
||||||
|
vision_hidden_states = hidden_states[visual_token_mask].reshape(
|
||||||
|
-1, self.hidden_size
|
||||||
|
)
|
||||||
|
|
||||||
|
text_router_logits, _ = self.text_experts_gate(
|
||||||
|
text_hidden_states.to(dtype=torch.float32)
|
||||||
|
)
|
||||||
|
text_topk_output = self.text_experts_topk(
|
||||||
|
text_hidden_states, text_router_logits
|
||||||
|
)
|
||||||
|
final_hidden_states[text_token_mask] = self.text_experts(
|
||||||
|
hidden_states=text_hidden_states, topk_output=text_topk_output
|
||||||
|
).flatten()
|
||||||
|
|
||||||
|
vision_router_logits, _ = self.vision_experts_gate(
|
||||||
|
vision_hidden_states.to(dtype=torch.float32)
|
||||||
|
)
|
||||||
|
vision_topk_output = self.vision_experts_topk(
|
||||||
|
vision_hidden_states, vision_router_logits
|
||||||
|
)
|
||||||
|
final_hidden_states[visual_token_mask] = self.vision_experts(
|
||||||
|
hidden_states=vision_hidden_states, topk_output=vision_topk_output
|
||||||
|
).flatten()
|
||||||
|
|
||||||
|
else:
|
||||||
|
# text modal input processing directly
|
||||||
|
text_router_logits, _ = self.text_experts_gate(
|
||||||
|
hidden_states.to(dtype=torch.float32)
|
||||||
|
)
|
||||||
|
topk_output = self.text_experts_topk(hidden_states, text_router_logits)
|
||||||
|
final_hidden_states = self.text_experts(
|
||||||
|
hidden_states=hidden_states, topk_output=topk_output
|
||||||
|
)
|
||||||
|
|
||||||
|
if shared_output is not None:
|
||||||
|
final_hidden_states = final_hidden_states + shared_output
|
||||||
|
|
||||||
|
if self.tp_size > 1:
|
||||||
|
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
|
||||||
|
|
||||||
|
return final_hidden_states.view(orig_shape)
|
||||||
|
|
||||||
|
|
||||||
|
class Ernie4_5_VLMoeDecoderLayer(nn.Module):
|
||||||
|
"""A single transformer layer.
|
||||||
|
|
||||||
|
Transformer layer takes input with size [s, b, h] and returns an
|
||||||
|
output of the same size.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config,
|
||||||
|
layer_id: int,
|
||||||
|
quant_config: Optional[QuantizationConfig] = None,
|
||||||
|
prefix: str = "",
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
rope_theta = getattr(config, "rope_theta", 500000)
|
||||||
|
rope_scaling = getattr(config, "rope_scaling", None)
|
||||||
|
rope_is_neox_style = getattr(config, "rope_is_neox_style", False)
|
||||||
|
freq_allocation = getattr(config, "freq_allocation", 20)
|
||||||
|
max_position_embeddings = getattr(config, "max_position_embeddings", 131072)
|
||||||
|
# Self attention.
|
||||||
|
self.self_attn = Ernie4_5_VLMoeAttention(
|
||||||
|
config=config,
|
||||||
|
hidden_size=config.hidden_size,
|
||||||
|
num_heads=config.num_attention_heads,
|
||||||
|
num_kv_heads=config.num_key_value_heads,
|
||||||
|
layer_id=layer_id,
|
||||||
|
rope_theta=rope_theta,
|
||||||
|
rope_scaling=rope_scaling,
|
||||||
|
rope_is_neox_style=rope_is_neox_style,
|
||||||
|
freq_allocation=freq_allocation,
|
||||||
|
max_position_embeddings=config.max_position_embeddings,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("self_attn", prefix),
|
||||||
|
bias=config.use_bias,
|
||||||
|
)
|
||||||
|
|
||||||
|
# MoE
|
||||||
|
moe_layer_start_index = config.moe_layer_start_index
|
||||||
|
min_moe_layer_start_index = min(moe_layer_start_index)
|
||||||
|
moe_layer_end_index = getattr(
|
||||||
|
config,
|
||||||
|
"moe_layer_end_index",
|
||||||
|
[config.num_hidden_layers - 1, config.num_hidden_layers - 1],
|
||||||
|
)
|
||||||
|
max_moe_layer_end_index = max(moe_layer_end_index)
|
||||||
|
assert min_moe_layer_start_index <= max_moe_layer_end_index
|
||||||
|
moe_num_experts = config.moe_num_experts
|
||||||
|
max_moe_num_experts = max(moe_num_experts)
|
||||||
|
moe_layer_interval = getattr(config, "moe_layer_interval", 1)
|
||||||
|
use_moe = getattr(config, "use_moe", max_moe_num_experts > 0)
|
||||||
|
# MLP
|
||||||
|
if (
|
||||||
|
use_moe
|
||||||
|
and ((layer_id + 1) % moe_layer_interval == 0)
|
||||||
|
and layer_id >= min_moe_layer_start_index
|
||||||
|
and layer_id <= max_moe_layer_end_index
|
||||||
|
):
|
||||||
|
self.mlp = Ernie4_5_VLMoeMoE(
|
||||||
|
config=config,
|
||||||
|
layer_id=layer_id,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("mlp", prefix),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.mlp = Ernie4_5_VLMoeMLP(
|
||||||
|
hidden_size=config.hidden_size,
|
||||||
|
intermediate_size=config.intermediate_size,
|
||||||
|
hidden_act=config.hidden_act,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("mlp", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||||
|
self.post_attention_layernorm = RMSNorm(
|
||||||
|
config.hidden_size, eps=config.rms_norm_eps
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
positions: torch.Tensor,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
forward_batch: ForwardBatch,
|
||||||
|
residual: Optional[torch.Tensor],
|
||||||
|
visual_token_mask: torch.Tensor | None,
|
||||||
|
**kwargs: object,
|
||||||
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||||
|
# Self Attention
|
||||||
|
if residual is None:
|
||||||
|
residual = hidden_states
|
||||||
|
hidden_states = self.input_layernorm(hidden_states)
|
||||||
|
else:
|
||||||
|
hidden_states, residual = self.input_layernorm(hidden_states, residual)
|
||||||
|
hidden_states = self.self_attn(
|
||||||
|
positions=positions,
|
||||||
|
hidden_states=hidden_states,
|
||||||
|
forward_batch=forward_batch,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Fully Connected
|
||||||
|
hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
|
||||||
|
if isinstance(self.mlp, Ernie4_5_VLMoeMoE):
|
||||||
|
hidden_states = self.mlp(hidden_states, visual_token_mask, **kwargs)
|
||||||
|
else:
|
||||||
|
hidden_states = self.mlp(hidden_states)
|
||||||
|
|
||||||
|
return hidden_states, residual
|
||||||
|
|
||||||
|
|
||||||
|
# only used as text backbone for ernie4.5 vl
|
||||||
|
class Ernie4_5_VLMoeModel(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config: PretrainedConfig,
|
||||||
|
quant_config: Optional[QuantizationConfig] = None,
|
||||||
|
prefix: str = "",
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.config = config
|
||||||
|
self.pp_group = get_pp_group()
|
||||||
|
|
||||||
|
if self.pp_group.is_first_rank:
|
||||||
|
self.embed_tokens = VocabParallelEmbedding(
|
||||||
|
config.vocab_size,
|
||||||
|
config.hidden_size,
|
||||||
|
enable_tp=not is_dp_attention_enabled(),
|
||||||
|
prefix=add_prefix("embed_tokens", prefix),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.embed_tokens = PPMissingLayer()
|
||||||
|
|
||||||
|
self.layers, self.start_layer, self.end_layer = make_layers(
|
||||||
|
config.num_hidden_layers,
|
||||||
|
lambda idx, prefix: Ernie4_5_VLMoeDecoderLayer(
|
||||||
|
layer_id=idx,
|
||||||
|
config=config,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=prefix,
|
||||||
|
),
|
||||||
|
pp_rank=self.pp_group.rank_in_group,
|
||||||
|
pp_size=self.pp_group.world_size,
|
||||||
|
prefix=add_prefix("layers", prefix),
|
||||||
|
)
|
||||||
|
if self.pp_group.is_last_rank:
|
||||||
|
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||||
|
else:
|
||||||
|
self.norm = PPMissingLayer(return_tuple=True)
|
||||||
|
|
||||||
|
def get_input_embeddings(self) -> torch.Tensor:
|
||||||
|
return self.embed_tokens
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
input_ids: torch.Tensor,
|
||||||
|
positions: torch.Tensor,
|
||||||
|
forward_batch: ForwardBatch,
|
||||||
|
input_embeds: torch.Tensor = None,
|
||||||
|
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
||||||
|
visual_token_mask: torch.Tensor | None = None,
|
||||||
|
) -> Union[torch.Tensor, PPProxyTensors]:
|
||||||
|
|
||||||
|
if self.pp_group.is_first_rank:
|
||||||
|
if input_embeds is None:
|
||||||
|
hidden_states = self.embed_tokens(input_ids)
|
||||||
|
else:
|
||||||
|
hidden_states = input_embeds
|
||||||
|
residual = None
|
||||||
|
else:
|
||||||
|
assert pp_proxy_tensors is not None
|
||||||
|
hidden_states = pp_proxy_tensors["hidden_states"]
|
||||||
|
residual = pp_proxy_tensors["residual"]
|
||||||
|
|
||||||
|
for layer in islice(self.layers, self.start_layer, self.end_layer):
|
||||||
|
hidden_states, residual = layer(
|
||||||
|
positions,
|
||||||
|
hidden_states,
|
||||||
|
forward_batch,
|
||||||
|
residual,
|
||||||
|
visual_token_mask,
|
||||||
|
)
|
||||||
|
|
||||||
|
if not self.pp_group.is_last_rank:
|
||||||
|
return PPProxyTensors(
|
||||||
|
{
|
||||||
|
"hidden_states": hidden_states,
|
||||||
|
"residual": residual,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
if hidden_states.shape[0] != 0:
|
||||||
|
if residual is None:
|
||||||
|
hidden_states = self.norm(hidden_states)
|
||||||
|
else:
|
||||||
|
hidden_states, _ = self.norm(hidden_states, residual)
|
||||||
|
|
||||||
|
return hidden_states
|
||||||
@@ -0,0 +1,845 @@
|
|||||||
|
# Copyright 2023-2025 SGLang Team
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
"""Inference-only Ernie45-VL model compatible with HuggingFace weights."""
|
||||||
|
import logging
|
||||||
|
from functools import lru_cache, partial
|
||||||
|
from typing import Iterable, List, Optional, Tuple, Type
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from einops import rearrange
|
||||||
|
from transformers import PretrainedConfig
|
||||||
|
|
||||||
|
from sglang.srt.layers.activation import QuickGELU
|
||||||
|
from sglang.srt.layers.attention.vision import VisionAttention
|
||||||
|
from sglang.srt.layers.layernorm import RMSNorm
|
||||||
|
from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
|
||||||
|
from sglang.srt.layers.logits_processor import LogitsProcessor
|
||||||
|
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
|
||||||
|
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||||
|
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
|
||||||
|
from sglang.srt.managers.mm_utils import (
|
||||||
|
MultiModalityDataPaddingPatternMultimodalTokens,
|
||||||
|
general_mm_embed_routine,
|
||||||
|
)
|
||||||
|
from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs
|
||||||
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||||
|
from sglang.srt.model_loader.weight_utils import default_weight_loader
|
||||||
|
from sglang.srt.models.ernie45_moe_vl import Ernie4_5_VLMoeModel
|
||||||
|
from sglang.srt.utils import add_prefix
|
||||||
|
from sglang.srt.utils.hf_transformers_utils import get_processor
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
# === Vision Encoder === #
|
||||||
|
|
||||||
|
|
||||||
|
class Ernie4_5_VisionMLP(nn.Module):
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
in_features: int,
|
||||||
|
hidden_features: int = None,
|
||||||
|
act_layer: Type[nn.Module] = QuickGELU,
|
||||||
|
quant_config: Optional[QuantizationConfig] = None,
|
||||||
|
prefix: str = "",
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.fc1 = ColumnParallelLinear(
|
||||||
|
in_features,
|
||||||
|
hidden_features,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("fc1", prefix),
|
||||||
|
)
|
||||||
|
self.act = act_layer()
|
||||||
|
self.fc2 = RowParallelLinear(
|
||||||
|
hidden_features,
|
||||||
|
in_features,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("fc2", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
x_parallel, _ = self.fc1(x)
|
||||||
|
x_parallel = self.act(x_parallel)
|
||||||
|
x, _ = self.fc2(x_parallel)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class Ernie4_5_VisionBlock(nn.Module):
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim: int,
|
||||||
|
num_heads: int,
|
||||||
|
mlp_ratio: float,
|
||||||
|
act_layer: Type[nn.Module] = QuickGELU,
|
||||||
|
norm_layer: Type[nn.Module] = None,
|
||||||
|
quant_config: Optional[QuantizationConfig] = None,
|
||||||
|
prefix: str = "",
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
if norm_layer is None:
|
||||||
|
norm_layer = partial(nn.LayerNorm, eps=1e-6)
|
||||||
|
self.norm1 = norm_layer(dim)
|
||||||
|
self.norm2 = norm_layer(dim)
|
||||||
|
mlp_hidden_dim = int(dim * mlp_ratio)
|
||||||
|
|
||||||
|
self.attn = VisionAttention(
|
||||||
|
embed_dim=dim,
|
||||||
|
num_heads=num_heads,
|
||||||
|
projection_size=dim,
|
||||||
|
use_qkv_parallel=True,
|
||||||
|
flatten_batch=True,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("attn", prefix),
|
||||||
|
)
|
||||||
|
self.mlp = Ernie4_5_VisionMLP(
|
||||||
|
dim,
|
||||||
|
mlp_hidden_dim,
|
||||||
|
act_layer=act_layer,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("mlp", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
x: torch.Tensor,
|
||||||
|
cu_seqlens: torch.Tensor,
|
||||||
|
position_embeddings: torch.Tensor,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
hidden_states = self.norm1(x)
|
||||||
|
hidden_states = rearrange(hidden_states, "s b ... -> b s ...")
|
||||||
|
attn = self.attn(
|
||||||
|
hidden_states,
|
||||||
|
cu_seqlens=cu_seqlens,
|
||||||
|
position_embeddings=position_embeddings,
|
||||||
|
)
|
||||||
|
attn = rearrange(attn, "b s ... -> s b ...")
|
||||||
|
x = x + attn
|
||||||
|
x = x + self.mlp(self.norm2(x))
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class Ernie4_5_VisionPatchEmbed(nn.Module):
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
patch_size: int = 14,
|
||||||
|
in_chans: int = 3,
|
||||||
|
embed_dim: int = 1280,
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.patch_size = patch_size
|
||||||
|
self.in_channels = in_chans
|
||||||
|
self.embed_dim = embed_dim
|
||||||
|
|
||||||
|
self.proj = nn.Linear(in_chans * patch_size * patch_size, embed_dim, bias=False)
|
||||||
|
|
||||||
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||||
|
target_dtype = self.proj.weight.dtype
|
||||||
|
hidden_states = hidden_states.to(target_dtype)
|
||||||
|
hidden_states = self.proj(hidden_states)
|
||||||
|
|
||||||
|
return hidden_states
|
||||||
|
|
||||||
|
|
||||||
|
class VariableResolutionResamplerModel(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
in_dim,
|
||||||
|
out_dim,
|
||||||
|
spatial_conv_size,
|
||||||
|
temporal_conv_size,
|
||||||
|
config,
|
||||||
|
prefix: str = "",
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.in_dim = in_dim
|
||||||
|
self.out_dim = out_dim
|
||||||
|
self.config = config
|
||||||
|
self.spatial_conv_size = spatial_conv_size
|
||||||
|
self.temporal_conv_size = temporal_conv_size
|
||||||
|
self.use_temporal_conv = config.use_temporal_conv
|
||||||
|
|
||||||
|
# compress 2d conv(picture) to 1d
|
||||||
|
self.spatial_dim = self.in_dim * self.spatial_conv_size * self.spatial_conv_size
|
||||||
|
# compress 3d conv(video) to 1d
|
||||||
|
self.temporal_dim = (
|
||||||
|
self.in_dim
|
||||||
|
* self.spatial_conv_size
|
||||||
|
* self.spatial_conv_size
|
||||||
|
* self.temporal_conv_size
|
||||||
|
)
|
||||||
|
|
||||||
|
self.spatial_linear1 = ColumnParallelLinear(
|
||||||
|
self.spatial_dim,
|
||||||
|
self.spatial_dim,
|
||||||
|
bias=True,
|
||||||
|
gather_output=True,
|
||||||
|
quant_config=getattr(config, "quant_config", None),
|
||||||
|
prefix=f"{prefix}.spatial_linear1",
|
||||||
|
)
|
||||||
|
|
||||||
|
self.spatial_gelu = nn.GELU()
|
||||||
|
|
||||||
|
self.spatial_linear2 = ColumnParallelLinear(
|
||||||
|
self.spatial_dim,
|
||||||
|
self.spatial_dim,
|
||||||
|
bias=True,
|
||||||
|
gather_output=True,
|
||||||
|
quant_config=getattr(config, "quant_config", None),
|
||||||
|
prefix=f"{prefix}.spatial_linear2",
|
||||||
|
)
|
||||||
|
|
||||||
|
self.spatial_norm = nn.LayerNorm(self.spatial_dim, eps=1e-6)
|
||||||
|
|
||||||
|
if self.use_temporal_conv:
|
||||||
|
self.temporal_linear1 = ColumnParallelLinear(
|
||||||
|
self.temporal_dim,
|
||||||
|
self.spatial_dim,
|
||||||
|
bias=True,
|
||||||
|
gather_output=True,
|
||||||
|
quant_config=getattr(config, "quant_config", None),
|
||||||
|
prefix=f"{prefix}.temporal_linear1",
|
||||||
|
)
|
||||||
|
|
||||||
|
self.temporal_gelu = nn.GELU()
|
||||||
|
|
||||||
|
self.temporal_linear2 = ColumnParallelLinear(
|
||||||
|
self.spatial_dim,
|
||||||
|
self.spatial_dim,
|
||||||
|
bias=True,
|
||||||
|
gather_output=True,
|
||||||
|
quant_config=getattr(config, "quant_config", None),
|
||||||
|
prefix=f"{prefix}.temporal_linear2",
|
||||||
|
)
|
||||||
|
|
||||||
|
self.temporal_norm = nn.LayerNorm(self.spatial_dim, eps=1e-6)
|
||||||
|
|
||||||
|
self.mlp = ColumnParallelLinear(
|
||||||
|
self.spatial_dim,
|
||||||
|
self.out_dim,
|
||||||
|
bias=True,
|
||||||
|
gather_output=True,
|
||||||
|
quant_config=getattr(config, "quant_config", None),
|
||||||
|
prefix=f"{prefix}.mlp",
|
||||||
|
)
|
||||||
|
|
||||||
|
self.after_norm = RMSNorm(
|
||||||
|
hidden_size=out_dim, eps=getattr(config, "rms_norm_eps", 1e-6)
|
||||||
|
)
|
||||||
|
|
||||||
|
def spatial_conv_reshape(self, x, spatial_conv_size):
|
||||||
|
S, C = x.shape
|
||||||
|
x = x.reshape([-1, C * (spatial_conv_size**2)])
|
||||||
|
return x
|
||||||
|
|
||||||
|
def forward(self, x, grid_thw):
|
||||||
|
def fwd_spatial(x):
|
||||||
|
x = self.spatial_conv_reshape(x, self.spatial_conv_size)
|
||||||
|
|
||||||
|
x, _ = self.spatial_linear1(x)
|
||||||
|
x = self.spatial_gelu(x)
|
||||||
|
x, _ = self.spatial_linear2(x)
|
||||||
|
x = self.spatial_norm(x)
|
||||||
|
|
||||||
|
return x
|
||||||
|
|
||||||
|
def fwd_placeholder(x, grid_thw, to_tensor=False):
|
||||||
|
grid_thw_cpu = grid_thw.cpu().numpy()
|
||||||
|
grid_t, grid_hw = grid_thw_cpu[:, 0], grid_thw_cpu[:, 1:]
|
||||||
|
grid_hw_after_conv = grid_hw.prod(-1) // (self.spatial_conv_size**2)
|
||||||
|
|
||||||
|
tokens_per_img_or_vid = grid_thw_cpu.prod(-1) // (self.spatial_conv_size**2)
|
||||||
|
batch_offset = np.empty(
|
||||||
|
tokens_per_img_or_vid.size, dtype=tokens_per_img_or_vid.dtype
|
||||||
|
)
|
||||||
|
batch_offset[0] = 0
|
||||||
|
batch_offset[1:] = tokens_per_img_or_vid.cumsum()[:-1]
|
||||||
|
|
||||||
|
slice_offsets = []
|
||||||
|
for temporoal_size, spatial_size, b_offset in zip(
|
||||||
|
grid_t, grid_hw_after_conv, batch_offset
|
||||||
|
):
|
||||||
|
for temp_offset in range(0, temporoal_size, 2):
|
||||||
|
slice_offsets.append(
|
||||||
|
np.arange(
|
||||||
|
b_offset + (temp_offset) * spatial_size,
|
||||||
|
b_offset + (temp_offset + 1) * spatial_size,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
slice_offsets = torch.tensor(np.concatenate(slice_offsets, axis=-1)).to(
|
||||||
|
x.device
|
||||||
|
)
|
||||||
|
|
||||||
|
slice_offsets2 = []
|
||||||
|
for temporoal_size, spatial_size, b_offset in zip(
|
||||||
|
grid_t, grid_hw_after_conv, batch_offset
|
||||||
|
):
|
||||||
|
for temp_offset in range(
|
||||||
|
1 if temporoal_size > 1 else 0, temporoal_size, 2
|
||||||
|
):
|
||||||
|
slice_offsets2.append(
|
||||||
|
np.arange(
|
||||||
|
b_offset + (temp_offset) * spatial_size,
|
||||||
|
b_offset + (temp_offset + 1) * spatial_size,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
slice_offsets2 = torch.tensor(np.concatenate(slice_offsets2, axis=-1)).to(
|
||||||
|
x.device
|
||||||
|
)
|
||||||
|
|
||||||
|
x_timestep_1 = torch.index_select(x, dim=0, index=slice_offsets)
|
||||||
|
x_timestep_2 = torch.index_select(x, dim=0, index=slice_offsets2)
|
||||||
|
x = torch.concat([x_timestep_1, x_timestep_2], dim=-1)
|
||||||
|
return x
|
||||||
|
|
||||||
|
def fwd_temporal(x):
|
||||||
|
x, _ = self.temporal_linear1(x)
|
||||||
|
x = self.temporal_gelu(x)
|
||||||
|
x, _ = self.temporal_linear2(x)
|
||||||
|
x = self.temporal_norm(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
def fwd_mlp(x):
|
||||||
|
x, _ = self.mlp(x)
|
||||||
|
x = self.after_norm(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
x = fwd_spatial(x)
|
||||||
|
if self.use_temporal_conv:
|
||||||
|
x = fwd_placeholder(x, grid_thw)
|
||||||
|
x = fwd_temporal(x)
|
||||||
|
x = fwd_mlp(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||||
|
params_dict = dict(self.named_parameters(remove_duplicate=False))
|
||||||
|
loaded_params: set[str] = set()
|
||||||
|
|
||||||
|
for name, loaded_weight in weights:
|
||||||
|
if name not in params_dict:
|
||||||
|
continue
|
||||||
|
param = params_dict[name]
|
||||||
|
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||||
|
weight_loader(param, loaded_weight)
|
||||||
|
loaded_params.add(name)
|
||||||
|
return loaded_params
|
||||||
|
|
||||||
|
|
||||||
|
class Ernie4_5_VisionRotaryEmbedding(nn.Module):
|
||||||
|
|
||||||
|
def __init__(self, dim: int, theta: float = 10000.0) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.inv_freq = 1.0 / theta ** (
|
||||||
|
torch.arange(start=0, end=dim, step=2, dtype=torch.float32) / dim
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, seqlen: int) -> torch.Tensor:
|
||||||
|
seq = torch.arange(
|
||||||
|
seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype
|
||||||
|
)
|
||||||
|
freqs = torch.outer(input=seq, vec2=self.inv_freq)
|
||||||
|
return freqs
|
||||||
|
|
||||||
|
|
||||||
|
class Ernie4_5_VisionTransformer(nn.Module):
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
vision_config: PretrainedConfig,
|
||||||
|
norm_eps: float = 1e-6,
|
||||||
|
quant_config: Optional[QuantizationConfig] = None,
|
||||||
|
prefix: str = "",
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
patch_size: int = vision_config.patch_size
|
||||||
|
spatial_merge_size: int = vision_config.spatial_merge_size
|
||||||
|
in_chans: int = vision_config.in_chans
|
||||||
|
hidden_size: int = vision_config.hidden_size
|
||||||
|
embed_dim: int = vision_config.embed_dim
|
||||||
|
depth: int = vision_config.depth
|
||||||
|
num_heads: int = vision_config.num_heads
|
||||||
|
mlp_ratio: float = vision_config.mlp_ratio
|
||||||
|
|
||||||
|
self.spatial_merge_size = spatial_merge_size
|
||||||
|
|
||||||
|
self.patch_embed = Ernie4_5_VisionPatchEmbed(
|
||||||
|
patch_size=patch_size,
|
||||||
|
in_chans=in_chans,
|
||||||
|
embed_dim=embed_dim,
|
||||||
|
)
|
||||||
|
|
||||||
|
norm_layer = partial(nn.LayerNorm, eps=norm_eps)
|
||||||
|
head_dim = embed_dim // num_heads
|
||||||
|
self.rotary_pos_emb = Ernie4_5_VisionRotaryEmbedding(head_dim // 2)
|
||||||
|
self.blocks = nn.ModuleList(
|
||||||
|
[
|
||||||
|
Ernie4_5_VisionBlock(
|
||||||
|
dim=embed_dim,
|
||||||
|
num_heads=num_heads,
|
||||||
|
mlp_ratio=mlp_ratio,
|
||||||
|
norm_layer=norm_layer,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix(f"blocks.{i}", prefix),
|
||||||
|
)
|
||||||
|
for i in range(depth)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
self.ln = nn.LayerNorm(hidden_size, eps=1e-6)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def dtype(self) -> torch.dtype:
|
||||||
|
return self.patch_embed.proj.weight.dtype
|
||||||
|
|
||||||
|
@property
|
||||||
|
def device(self) -> torch.device:
|
||||||
|
return self.blocks[0].mlp.fc2.weight.device
|
||||||
|
|
||||||
|
def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
|
||||||
|
pos_ids = []
|
||||||
|
for i in range(grid_thw.size(0)):
|
||||||
|
t, h, w = grid_thw[i].tolist()
|
||||||
|
hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
|
||||||
|
wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
|
||||||
|
hpos_ids = (
|
||||||
|
hpos_ids.reshape(
|
||||||
|
h // self.spatial_merge_size,
|
||||||
|
self.spatial_merge_size,
|
||||||
|
w // self.spatial_merge_size,
|
||||||
|
self.spatial_merge_size,
|
||||||
|
)
|
||||||
|
.permute(0, 2, 1, 3)
|
||||||
|
.flatten()
|
||||||
|
)
|
||||||
|
wpos_ids = (
|
||||||
|
wpos_ids.reshape(
|
||||||
|
h // self.spatial_merge_size,
|
||||||
|
self.spatial_merge_size,
|
||||||
|
w // self.spatial_merge_size,
|
||||||
|
self.spatial_merge_size,
|
||||||
|
)
|
||||||
|
.permute(0, 2, 1, 3)
|
||||||
|
.flatten()
|
||||||
|
)
|
||||||
|
pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
|
||||||
|
pos_ids = torch.cat(pos_ids, dim=0)
|
||||||
|
max_grid_size = grid_thw[:, 1:].max()
|
||||||
|
rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
|
||||||
|
rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
|
||||||
|
return rotary_pos_emb
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
x: torch.Tensor,
|
||||||
|
grid_thw: torch.Tensor,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
# patchify
|
||||||
|
x = x.to(device=self.device, dtype=self.dtype)
|
||||||
|
x = self.patch_embed(x)
|
||||||
|
|
||||||
|
# compute position embedding
|
||||||
|
rotary_pos_emb = self.rot_pos_emb(grid_thw)
|
||||||
|
emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
|
||||||
|
position_embeddings = (emb.cos(), emb.sin())
|
||||||
|
# compute cu_seqlens
|
||||||
|
cu_seqlens = torch.repeat_interleave(
|
||||||
|
grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
|
||||||
|
).cumsum(dim=0, dtype=torch.int32)
|
||||||
|
cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
|
||||||
|
|
||||||
|
# transformers
|
||||||
|
x = x.unsqueeze(1)
|
||||||
|
for blk in self.blocks:
|
||||||
|
x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings)
|
||||||
|
|
||||||
|
final_output = self.ln(x)
|
||||||
|
|
||||||
|
if final_output.ndim == 3:
|
||||||
|
final_output = final_output.squeeze(dim=1)
|
||||||
|
|
||||||
|
return final_output
|
||||||
|
|
||||||
|
|
||||||
|
cached_get_processor = lru_cache(get_processor)
|
||||||
|
|
||||||
|
|
||||||
|
class Ernie4_5_VLMoeForConditionalGeneration(nn.Module):
|
||||||
|
# BitandBytes specific attributes
|
||||||
|
default_bitsandbytes_target_modules = [
|
||||||
|
".gate_proj.",
|
||||||
|
".down_proj.",
|
||||||
|
".up_proj.",
|
||||||
|
".q_proj.",
|
||||||
|
".k_proj.",
|
||||||
|
".v_proj.",
|
||||||
|
".o_proj.",
|
||||||
|
]
|
||||||
|
bitsandbytes_stacked_params_mapping = {
|
||||||
|
# shard_name, weight_name, index
|
||||||
|
"q_proj": ("qkv_proj", 0),
|
||||||
|
"k_proj": ("qkv_proj", 1),
|
||||||
|
"v_proj": ("qkv_proj", 2),
|
||||||
|
"gate_proj": ("gate_up_proj", 0),
|
||||||
|
"up_proj": ("gate_up_proj", 1),
|
||||||
|
}
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config: PretrainedConfig,
|
||||||
|
quant_config: Optional[QuantizationConfig] = None,
|
||||||
|
prefix: str = "",
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
self.config = config
|
||||||
|
self.vision_model = Ernie4_5_VisionTransformer(
|
||||||
|
config.vision_config,
|
||||||
|
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("vision_model", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.model = Ernie4_5_VLMoeModel(
|
||||||
|
config, quant_config, prefix=add_prefix("model", prefix)
|
||||||
|
)
|
||||||
|
|
||||||
|
self.resampler_model = VariableResolutionResamplerModel(
|
||||||
|
self.config.pixel_hidden_size,
|
||||||
|
self.config.hidden_size,
|
||||||
|
self.config.spatial_conv_size,
|
||||||
|
self.config.temporal_conv_size,
|
||||||
|
config=self.config,
|
||||||
|
prefix=add_prefix("resampler_model", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
if config.tie_word_embeddings:
|
||||||
|
self.lm_head = self.model.embed_tokens
|
||||||
|
else:
|
||||||
|
self.lm_head = ParallelLMHead(
|
||||||
|
config.vocab_size,
|
||||||
|
config.hidden_size,
|
||||||
|
quant_config=quant_config,
|
||||||
|
prefix=add_prefix("lm_head", prefix),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.is_mrope_enabled = "mrope_section" in self.config.rope_scaling
|
||||||
|
self.logits_processor = LogitsProcessor(config)
|
||||||
|
|
||||||
|
if getattr(self.config, "im_patch_id", None):
|
||||||
|
visual_token_ids = [
|
||||||
|
token_id
|
||||||
|
for token_id in [
|
||||||
|
self.config.im_patch_id,
|
||||||
|
getattr(self.config, "image_start_token_id", None),
|
||||||
|
getattr(self.config, "image_end_token_id", None),
|
||||||
|
getattr(self.config, "video_start_token_id", None),
|
||||||
|
getattr(self.config, "video_end_token_id", None),
|
||||||
|
]
|
||||||
|
if token_id is not None
|
||||||
|
]
|
||||||
|
self._visual_token_ids_tensor_cache = torch.tensor(
|
||||||
|
visual_token_ids, dtype=torch.long
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self._visual_token_ids_tensor_cache = None
|
||||||
|
|
||||||
|
def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
|
||||||
|
pattern = MultiModalityDataPaddingPatternMultimodalTokens()
|
||||||
|
return pattern.pad_input_tokens(input_ids, mm_inputs)
|
||||||
|
|
||||||
|
def _vision_forward(
|
||||||
|
self,
|
||||||
|
pixel_values: torch.Tensor,
|
||||||
|
grid_thw: torch.Tensor,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
if grid_thw is not None:
|
||||||
|
grid_thw = grid_thw[grid_thw > 0]
|
||||||
|
if grid_thw.numel() % 3 != 0:
|
||||||
|
raise ValueError(
|
||||||
|
f"grid_thw has {grid_thw.numel()} elements after filtering,"
|
||||||
|
"which is not divisible by 3."
|
||||||
|
)
|
||||||
|
grid_thw = grid_thw.reshape(-1, 3)
|
||||||
|
# example: [[1,64,64],[2,80,80]] -> [[1,64,64],[1,80,80],[1,80,80]]
|
||||||
|
grid_thw = F.pad(
|
||||||
|
torch.repeat_interleave(grid_thw[:, 1:], grid_thw[:, 0], 0),
|
||||||
|
[1, 0, 0, 0],
|
||||||
|
value=1,
|
||||||
|
)
|
||||||
|
image_features = self.vision_model(pixel_values, grid_thw)
|
||||||
|
return image_features
|
||||||
|
|
||||||
|
def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
|
||||||
|
# in qwen-vl, last dim is the same
|
||||||
|
pixel_values = torch.cat([item.feature for item in items], dim=0).type(
|
||||||
|
self.vision_model.dtype
|
||||||
|
)
|
||||||
|
image_grid_thw = torch.concat([item.image_grid_thw for item in items], dim=0)
|
||||||
|
assert pixel_values.dim() == 2, pixel_values.dim()
|
||||||
|
assert image_grid_thw.dim() == 2, image_grid_thw.dim()
|
||||||
|
image_feature = self._vision_forward(pixel_values, grid_thw=image_grid_thw)
|
||||||
|
image_embeds = self.resampler_model(image_feature, image_grid_thw)
|
||||||
|
return image_embeds
|
||||||
|
|
||||||
|
def get_video_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
|
||||||
|
# in qwen-vl, last dim is the same
|
||||||
|
pixel_values = torch.cat([item.feature for item in items], dim=0).type(
|
||||||
|
self.vision_model.dtype
|
||||||
|
)
|
||||||
|
video_grid_thw = torch.concat([item.video_grid_thw for item in items], dim=0)
|
||||||
|
assert pixel_values.dim() == 2, pixel_values.dim()
|
||||||
|
assert video_grid_thw.dim() == 2, video_grid_thw.dim()
|
||||||
|
video_feature = self._vision_forward(pixel_values, grid_thw=video_grid_thw)
|
||||||
|
video_embeds = self.resampler_model(video_feature, video_grid_thw)
|
||||||
|
return video_embeds
|
||||||
|
|
||||||
|
def _set_visual_token_mask(
|
||||||
|
self, input_ids: torch.Tensor, forward_batch: ForwardBatch
|
||||||
|
) -> None:
|
||||||
|
"""Set mask for visual tokens (image/video patches and delimiters)."""
|
||||||
|
if self._visual_token_ids_tensor_cache is None:
|
||||||
|
self.visual_token_mask = None
|
||||||
|
return
|
||||||
|
# Create tensor on the correct device
|
||||||
|
visual_token_ids_tensor = self._visual_token_ids_tensor_cache.to(
|
||||||
|
device=input_ids.device,
|
||||||
|
dtype=input_ids.dtype,
|
||||||
|
)
|
||||||
|
|
||||||
|
pad_values = []
|
||||||
|
if hasattr(forward_batch, "mm_inputs") and forward_batch.mm_inputs is not None:
|
||||||
|
for mm_input in forward_batch.mm_inputs:
|
||||||
|
if mm_input is None:
|
||||||
|
continue
|
||||||
|
for item in mm_input.mm_items:
|
||||||
|
pad_values.append(item.pad_value)
|
||||||
|
placeholder_tensor = torch.as_tensor(
|
||||||
|
pad_values,
|
||||||
|
device=input_ids.device,
|
||||||
|
)
|
||||||
|
pad_visual_token_ids_tensor = torch.cat(
|
||||||
|
[visual_token_ids_tensor, placeholder_tensor], dim=0
|
||||||
|
)
|
||||||
|
self.visual_token_mask = torch.isin(
|
||||||
|
input_ids, pad_visual_token_ids_tensor
|
||||||
|
).reshape(-1, 1)
|
||||||
|
|
||||||
|
def get_input_embeddings(self):
|
||||||
|
return self.model.embed_tokens
|
||||||
|
|
||||||
|
def should_apply_lora(self, module_name: str) -> bool:
|
||||||
|
# skip vision_model
|
||||||
|
return not module_name.startswith("vision_model")
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
input_ids: torch.Tensor,
|
||||||
|
positions: torch.Tensor,
|
||||||
|
forward_batch: ForwardBatch,
|
||||||
|
get_embedding: bool = False,
|
||||||
|
):
|
||||||
|
"""Run forward pass for Ernie45-VL.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_ids: Flattened (concatenated) input_ids corresponding to a
|
||||||
|
batch.
|
||||||
|
positions: Flattened (concatenated) position ids corresponding to a
|
||||||
|
batch.
|
||||||
|
**NOTE**: If mrope is enabled (default setting for Qwen2-VL
|
||||||
|
opensource models), the shape will be `(3, seq_len)`,
|
||||||
|
otherwise it will be `(seq_len,).
|
||||||
|
(Use input_metadata.mrope_positions to replace it)
|
||||||
|
"""
|
||||||
|
if self.is_mrope_enabled:
|
||||||
|
positions = forward_batch.mrope_positions
|
||||||
|
|
||||||
|
if not (
|
||||||
|
forward_batch.forward_mode.is_decode()
|
||||||
|
or not forward_batch.contains_image_inputs()
|
||||||
|
):
|
||||||
|
if self.is_mrope_enabled:
|
||||||
|
assert positions.ndim == 2 and positions.size(0) == 3, (
|
||||||
|
"multimodal section rotary embedding requires "
|
||||||
|
f"(3, seq_len) positions, but got {positions.size()}"
|
||||||
|
)
|
||||||
|
|
||||||
|
self._set_visual_token_mask(input_ids, forward_batch)
|
||||||
|
|
||||||
|
assert (
|
||||||
|
input_ids.numel() == positions.shape[-1]
|
||||||
|
), f"input_ids {input_ids.shape} and position_ids {positions.shape} should have the same length"
|
||||||
|
|
||||||
|
hidden_states = general_mm_embed_routine(
|
||||||
|
input_ids=input_ids,
|
||||||
|
forward_batch=forward_batch,
|
||||||
|
language_model=self.model,
|
||||||
|
multimodal_model=self,
|
||||||
|
positions=positions,
|
||||||
|
visual_token_mask=self.visual_token_mask,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.visual_token_mask = None
|
||||||
|
|
||||||
|
return self.logits_processor(
|
||||||
|
input_ids, hidden_states, self.lm_head, forward_batch
|
||||||
|
)
|
||||||
|
|
||||||
|
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
|
||||||
|
stacked_params_mapping = [
|
||||||
|
# (param_name, shard_name, shard_id)
|
||||||
|
("qkv_proj", "q_proj", "q"),
|
||||||
|
("qkv_proj", "k_proj", "k"),
|
||||||
|
("qkv_proj", "v_proj", "v"),
|
||||||
|
("gate_up_proj", "up_proj", 1),
|
||||||
|
("gate_up_proj", "gate_proj", 0),
|
||||||
|
]
|
||||||
|
|
||||||
|
# resampler_weight_mappings
|
||||||
|
resampler_weight_mapping = {
|
||||||
|
"spatial_linear.0.": "spatial_linear1.",
|
||||||
|
"spatial_linear.2.": "spatial_linear2.",
|
||||||
|
"spatial_linear.3.": "spatial_norm.",
|
||||||
|
"temporal_linear.0.": "temporal_linear1.",
|
||||||
|
"temporal_linear.2.": "temporal_linear2.",
|
||||||
|
"temporal_linear.3.": "temporal_norm.",
|
||||||
|
}
|
||||||
|
|
||||||
|
expert_params_mapping = FusedMoE.make_expert_params_mapping(
|
||||||
|
ckpt_gate_proj_name="gate_proj",
|
||||||
|
ckpt_down_proj_name="down_proj",
|
||||||
|
ckpt_up_proj_name="up_proj",
|
||||||
|
num_experts=max(self.config.moe_num_experts),
|
||||||
|
)
|
||||||
|
params_dict = dict(self.named_parameters(remove_duplicate=False))
|
||||||
|
for name, loaded_weight in weights:
|
||||||
|
if "rotary_emb.inv_freq" in name:
|
||||||
|
continue
|
||||||
|
if self.config.tie_word_embeddings and "lm_head.weight" in name:
|
||||||
|
continue
|
||||||
|
|
||||||
|
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||||
|
if weight_name not in name:
|
||||||
|
continue
|
||||||
|
|
||||||
|
if ("mlp.experts." in name) and name not in params_dict:
|
||||||
|
continue
|
||||||
|
name = name.replace(weight_name, param_name)
|
||||||
|
|
||||||
|
# Skip loading extra bias for GPTQ models.
|
||||||
|
if name.endswith(".bias") and name not in params_dict:
|
||||||
|
continue
|
||||||
|
param = params_dict[name]
|
||||||
|
weight_loader = param.weight_loader
|
||||||
|
weight_loader(param, loaded_weight, shard_id)
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
if "vision_model" in name:
|
||||||
|
# adapt to VisionAttention
|
||||||
|
name = name.replace(r"attn.qkv.", r"attn.qkv_proj.")
|
||||||
|
if name.startswith("model.resampler_model"):
|
||||||
|
name = name.replace("model.resampler_model", "resampler_model")
|
||||||
|
|
||||||
|
for (
|
||||||
|
old_weight_name,
|
||||||
|
new_weight_name,
|
||||||
|
) in resampler_weight_mapping.items():
|
||||||
|
if old_weight_name in name:
|
||||||
|
name = name.replace(old_weight_name, new_weight_name, 1)
|
||||||
|
break
|
||||||
|
|
||||||
|
# Distinguish between vision experts and text experts
|
||||||
|
if "mlp.experts" in name:
|
||||||
|
moe_offset = int(name.split(".")[-3])
|
||||||
|
vision_expert_start_idx = self.config.moe_num_experts[0]
|
||||||
|
is_text_expert = moe_offset <= vision_expert_start_idx - 1
|
||||||
|
if is_text_expert:
|
||||||
|
name = name.replace(".experts.", ".text_experts.")
|
||||||
|
else:
|
||||||
|
name = name.replace(
|
||||||
|
f".experts.{moe_offset}",
|
||||||
|
f".vision_experts.{moe_offset - vision_expert_start_idx}",
|
||||||
|
)
|
||||||
|
|
||||||
|
for mapping in expert_params_mapping:
|
||||||
|
param_name, weight_name, expert_id, shard_id = mapping
|
||||||
|
if weight_name not in name:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Distinguish between vision experts and text experts
|
||||||
|
moe_offset = int(name.split(".")[-3])
|
||||||
|
is_text_expert = moe_offset <= self.config.moe_num_experts[0] - 1
|
||||||
|
|
||||||
|
name = name.replace(weight_name, param_name)
|
||||||
|
if is_text_expert:
|
||||||
|
name = name.replace(".experts.", ".text_experts.")
|
||||||
|
else:
|
||||||
|
name = name.replace(".experts.", ".vision_experts.")
|
||||||
|
|
||||||
|
# Skip loading extra bias for GPTQ models.
|
||||||
|
if (
|
||||||
|
name.endswith(".bias") or name.endswith("_bias")
|
||||||
|
) and name not in params_dict:
|
||||||
|
continue
|
||||||
|
|
||||||
|
if name in params_dict.keys():
|
||||||
|
param = params_dict[name]
|
||||||
|
weight_loader = param.weight_loader
|
||||||
|
weight_loader(
|
||||||
|
param,
|
||||||
|
loaded_weight,
|
||||||
|
name,
|
||||||
|
shard_id=shard_id,
|
||||||
|
expert_id=expert_id,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
logger.warning(f"Parameter {name} not found in params_dict")
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
# Distinguish between vision expert gate
|
||||||
|
# and text expert gate
|
||||||
|
if name.endswith("mlp.gate.weight"):
|
||||||
|
name = name.replace("gate.weight", "text_experts_gate.weight")
|
||||||
|
loaded_weight = loaded_weight.T
|
||||||
|
elif name.endswith("mlp.gate.weight_1"):
|
||||||
|
name = name.replace(
|
||||||
|
"gate.weight_1", "vision_experts_gate.weight"
|
||||||
|
)
|
||||||
|
loaded_weight = loaded_weight.T
|
||||||
|
|
||||||
|
if "e_score_correction_bias" in name:
|
||||||
|
name = name.replace(".moe_statics.", ".")
|
||||||
|
|
||||||
|
# Skip loading extra bias for GPTQ models.
|
||||||
|
if (
|
||||||
|
name.endswith(".bias") or name.endswith("_bias")
|
||||||
|
) and name not in params_dict:
|
||||||
|
continue
|
||||||
|
|
||||||
|
if name in params_dict.keys():
|
||||||
|
param = params_dict[name]
|
||||||
|
weight_loader = getattr(
|
||||||
|
param, "weight_loader", default_weight_loader
|
||||||
|
)
|
||||||
|
weight_loader(param, loaded_weight)
|
||||||
|
else:
|
||||||
|
logger.warning(f"Parameter {name} not found in params_dict")
|
||||||
|
|
||||||
|
|
||||||
|
EntryClass = [Ernie4_5_VLMoeForConditionalGeneration]
|
||||||
@@ -0,0 +1,417 @@
|
|||||||
|
import math
|
||||||
|
import os
|
||||||
|
from typing import List, Union
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
import torchvision
|
||||||
|
from PIL import Image
|
||||||
|
from torchvision.transforms import InterpolationMode
|
||||||
|
from transformers import BaseImageProcessorFast
|
||||||
|
|
||||||
|
from sglang.srt.environ import envs
|
||||||
|
from sglang.srt.layers.rotary_embedding import MRotaryEmbedding
|
||||||
|
from sglang.srt.models.ernie45_vl import Ernie4_5_VLMoeForConditionalGeneration
|
||||||
|
from sglang.srt.multimodal.processors.base_processor import (
|
||||||
|
BaseMultimodalProcessor as SGLangBaseProcessor,
|
||||||
|
)
|
||||||
|
from sglang.srt.multimodal.processors.base_processor import MultimodalSpecialTokens
|
||||||
|
from sglang.srt.utils import get_bool_env_var, is_npu, logger
|
||||||
|
|
||||||
|
_is_npu = is_npu()
|
||||||
|
|
||||||
|
SGL_USE_CUDA_IPC = get_bool_env_var("SGLANG_USE_CUDA_IPC_TRANSPORT")
|
||||||
|
|
||||||
|
|
||||||
|
IMAGE_FACTOR = 28
|
||||||
|
MIN_PIXELS = 4 * 28 * 28
|
||||||
|
# MAX_PIXELS = envs.SGLANG_IMAGE_MAX_PIXELS.get()
|
||||||
|
MAX_PIXELS = 16384 * 28 * 28
|
||||||
|
MAX_RATIO = 200
|
||||||
|
RESIZE_RESAMPLE = getattr(Image, envs.SGLANG_RESIZE_RESAMPLE.get(), None)
|
||||||
|
if envs.SGLANG_RESIZE_RESAMPLE.is_set() and RESIZE_RESAMPLE is None:
|
||||||
|
logger.warning(
|
||||||
|
f"Invalid RESIZE_RESAMPLE value: '{envs.SGLANG_RESIZE_RESAMPLE.get()}'. "
|
||||||
|
f"Ignoring and using default."
|
||||||
|
)
|
||||||
|
VIDEO_TOTAL_PIXELS = int(
|
||||||
|
float(os.environ.get("VIDEO_MAX_PIXELS", 128000 * 28 * 28 * 0.9))
|
||||||
|
)
|
||||||
|
|
||||||
|
VIDEO_MIN_PIXELS = 299 * 28 * 28
|
||||||
|
VIDEO_MAX_PIXELS = 1196 * 28 * 28
|
||||||
|
FRAME_FACTOR = 2
|
||||||
|
FPS = 2.0
|
||||||
|
FPS_MIN_FRAMES = 16
|
||||||
|
FPS_MAX_FRAMES = 180
|
||||||
|
|
||||||
|
|
||||||
|
def smart_resize(
|
||||||
|
height: int,
|
||||||
|
width: int,
|
||||||
|
factor: int = IMAGE_FACTOR,
|
||||||
|
min_pixels: int = MIN_PIXELS,
|
||||||
|
max_pixels: int = MAX_PIXELS,
|
||||||
|
):
|
||||||
|
if max(height, width) / min(height, width) > MAX_RATIO:
|
||||||
|
if height > width:
|
||||||
|
new_width = max(factor, round_by_factor(width, factor))
|
||||||
|
new_height = floor_by_factor(new_width * MAX_RATIO, factor)
|
||||||
|
else:
|
||||||
|
new_height = max(factor, round_by_factor(height, factor))
|
||||||
|
new_width = floor_by_factor(new_height * MAX_RATIO, factor)
|
||||||
|
|
||||||
|
height = new_height
|
||||||
|
width = new_width
|
||||||
|
|
||||||
|
h_bar = max(factor, round_by_factor(height, factor))
|
||||||
|
w_bar = max(factor, round_by_factor(width, factor))
|
||||||
|
if h_bar * w_bar > max_pixels:
|
||||||
|
beta = math.sqrt((height * width) / max_pixels)
|
||||||
|
h_bar = floor_by_factor(height / beta, factor)
|
||||||
|
w_bar = floor_by_factor(width / beta, factor)
|
||||||
|
elif h_bar * w_bar < min_pixels:
|
||||||
|
beta = math.sqrt(min_pixels / (height * width))
|
||||||
|
h_bar = ceil_by_factor(height * beta, factor)
|
||||||
|
w_bar = ceil_by_factor(width * beta, factor)
|
||||||
|
|
||||||
|
if min_pixels > h_bar * w_bar or h_bar * w_bar > max_pixels:
|
||||||
|
raise ValueError(f"encounter invalid h_bar: {h_bar}, w_bar: {w_bar}")
|
||||||
|
|
||||||
|
return h_bar, w_bar
|
||||||
|
|
||||||
|
|
||||||
|
def resize_image(
|
||||||
|
image,
|
||||||
|
min_pixels: int = MIN_PIXELS,
|
||||||
|
max_pixels: int = MAX_PIXELS,
|
||||||
|
size_factor: int = IMAGE_FACTOR,
|
||||||
|
) -> Image.Image:
|
||||||
|
width, height = image.size
|
||||||
|
min_pixels = min_pixels
|
||||||
|
max_pixels = max_pixels
|
||||||
|
resized_height, resized_width = smart_resize(
|
||||||
|
height,
|
||||||
|
width,
|
||||||
|
factor=size_factor,
|
||||||
|
min_pixels=min_pixels,
|
||||||
|
max_pixels=max_pixels,
|
||||||
|
)
|
||||||
|
image = image.resize((resized_width, resized_height), resample=RESIZE_RESAMPLE)
|
||||||
|
return image
|
||||||
|
|
||||||
|
|
||||||
|
def round_by_factor(number: int | float, factor: int) -> int:
|
||||||
|
return round(number / factor) * factor
|
||||||
|
|
||||||
|
|
||||||
|
def ceil_by_factor(number: int | float, factor: int) -> int:
|
||||||
|
return math.ceil(number / factor) * factor
|
||||||
|
|
||||||
|
|
||||||
|
def floor_by_factor(number: int | float, factor: int) -> int:
|
||||||
|
return math.floor(number / factor) * factor
|
||||||
|
|
||||||
|
|
||||||
|
async def resize_image_async(
|
||||||
|
image,
|
||||||
|
min_pixels: int = MIN_PIXELS,
|
||||||
|
max_pixels: int = MAX_PIXELS,
|
||||||
|
size_factor: int = IMAGE_FACTOR,
|
||||||
|
):
|
||||||
|
return resize_image(image, min_pixels, max_pixels, size_factor)
|
||||||
|
|
||||||
|
|
||||||
|
def smart_nframes(
|
||||||
|
ele: dict,
|
||||||
|
total_frames: int,
|
||||||
|
video_fps: int | float,
|
||||||
|
) -> int:
|
||||||
|
"""calculate the number of frames for video used for model inputs.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
ele (dict): a dict contains the configuration of video.
|
||||||
|
support either `fps` or `nframes`:
|
||||||
|
- nframes: the number of frames to extract for model inputs.
|
||||||
|
- fps: the fps to extract frames for model inputs.
|
||||||
|
- min_frames: the minimum number of frames of the video, only used when fps is provided.
|
||||||
|
- max_frames: the maximum number of frames of the video, only used when fps is provided.
|
||||||
|
total_frames (int): the original total number of frames of the video.
|
||||||
|
video_fps (int | float): the original fps of the video.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
int: the number of frames for video used for model inputs.
|
||||||
|
"""
|
||||||
|
assert not (
|
||||||
|
"fps" in ele and "nframes" in ele
|
||||||
|
), "Only accept either `fps` or `nframes`"
|
||||||
|
if "nframes" in ele:
|
||||||
|
nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
|
||||||
|
else:
|
||||||
|
fps = ele.get("fps", FPS)
|
||||||
|
min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
|
||||||
|
max_frames = floor_by_factor(
|
||||||
|
ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), FRAME_FACTOR
|
||||||
|
)
|
||||||
|
nframes = total_frames / video_fps * fps
|
||||||
|
if nframes > total_frames:
|
||||||
|
logger.warning(
|
||||||
|
f"smart_nframes: nframes[{nframes}] > total_frames[{total_frames}]"
|
||||||
|
)
|
||||||
|
nframes = min(min(max(nframes, min_frames), max_frames), total_frames)
|
||||||
|
nframes = floor_by_factor(nframes, FRAME_FACTOR)
|
||||||
|
if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
|
||||||
|
raise ValueError(
|
||||||
|
f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}."
|
||||||
|
)
|
||||||
|
return nframes
|
||||||
|
|
||||||
|
|
||||||
|
# process video, qwen-specific
|
||||||
|
async def preprocess_video(
|
||||||
|
vr,
|
||||||
|
image_factor: int = IMAGE_FACTOR,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
|
||||||
|
total_frames, video_fps = len(vr), vr.get_avg_fps()
|
||||||
|
nframes = smart_nframes({}, total_frames=total_frames, video_fps=video_fps)
|
||||||
|
idx = np.linspace(0, total_frames - 1, num=nframes, dtype=np.int64)
|
||||||
|
idx = np.unique(idx)
|
||||||
|
video_np = vr.get_batch(idx).asnumpy()
|
||||||
|
video = torch.from_numpy(video_np).pin_memory()
|
||||||
|
video = video.permute(0, 3, 1, 2) # Convert to TCHW format
|
||||||
|
nframes, _, height, width = video.shape
|
||||||
|
min_pixels = VIDEO_MIN_PIXELS
|
||||||
|
total_pixels = VIDEO_TOTAL_PIXELS
|
||||||
|
max_pixels = max(
|
||||||
|
min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR),
|
||||||
|
int(min_pixels * 1.05),
|
||||||
|
)
|
||||||
|
|
||||||
|
resized_height, resized_width = smart_resize(
|
||||||
|
height,
|
||||||
|
width,
|
||||||
|
factor=image_factor,
|
||||||
|
min_pixels=min_pixels,
|
||||||
|
max_pixels=max_pixels,
|
||||||
|
)
|
||||||
|
video = torchvision.transforms.functional.resize(
|
||||||
|
video,
|
||||||
|
[resized_height, resized_width],
|
||||||
|
interpolation=InterpolationMode.BILINEAR,
|
||||||
|
)
|
||||||
|
|
||||||
|
video = video.permute(0, 2, 3, 1)
|
||||||
|
video = video.pin_memory()
|
||||||
|
video_metadata = {
|
||||||
|
"fps": video_fps,
|
||||||
|
"duration": total_frames / video_fps,
|
||||||
|
"total_num_frames": total_frames,
|
||||||
|
"frames_indices": idx,
|
||||||
|
"video_backend": "torchvision",
|
||||||
|
}
|
||||||
|
|
||||||
|
return video, video_metadata
|
||||||
|
|
||||||
|
|
||||||
|
# Compatible with Ernie-VL Series
|
||||||
|
class Ernie4_5_VLImageProcessor(SGLangBaseProcessor):
|
||||||
|
models = [Ernie4_5_VLMoeForConditionalGeneration]
|
||||||
|
|
||||||
|
def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
|
||||||
|
super().__init__(hf_config, server_args, _processor, *args, **kwargs)
|
||||||
|
self.hf_config = hf_config
|
||||||
|
self.model_type = hf_config.model_type
|
||||||
|
self.image_start_token_id = hf_config.image_start_token_id
|
||||||
|
self.image_end_token_id = hf_config.image_end_token_id
|
||||||
|
self.video_start_token_id = hf_config.video_start_token_id
|
||||||
|
self.video_end_token_id = hf_config.video_end_token_id
|
||||||
|
|
||||||
|
self.IMAGE_FACTOR = 28
|
||||||
|
self.MIN_PIXELS = 4 * 28 * 28
|
||||||
|
self.MAX_PIXELS = 16384 * 28 * 28
|
||||||
|
self.MAX_RATIO = 200
|
||||||
|
self.mm_tokens = MultimodalSpecialTokens(
|
||||||
|
image_token="<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>",
|
||||||
|
video_token="<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>",
|
||||||
|
image_token_id=hf_config.im_patch_id,
|
||||||
|
video_token_id=hf_config.im_patch_id, # image and video use the same token_id
|
||||||
|
).build(_processor)
|
||||||
|
|
||||||
|
self.tokenizer = self._processor.tokenizer
|
||||||
|
self.image_processor = self._processor.image_processor
|
||||||
|
|
||||||
|
def _pixel_values_norm(
|
||||||
|
self,
|
||||||
|
pixel_values: torch.Tensor,
|
||||||
|
mm_kwargs: object,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
hf_config = self.hf_config
|
||||||
|
vision_config = hf_config.vision_config
|
||||||
|
image_processor = self.image_processor
|
||||||
|
image_mean_tensor = torch.tensor(
|
||||||
|
image_processor.image_mean, dtype=torch.float32
|
||||||
|
).reshape([1, 3, 1, 1])
|
||||||
|
image_std_tensor = torch.tensor(
|
||||||
|
image_processor.image_std, dtype=torch.float32
|
||||||
|
).reshape([1, 3, 1, 1])
|
||||||
|
rescale_factor = torch.tensor(
|
||||||
|
image_processor.rescale_factor, dtype=torch.float32
|
||||||
|
)
|
||||||
|
patch_size_squared = vision_config.patch_size**2
|
||||||
|
|
||||||
|
image_mean_tensor = image_mean_tensor.squeeze([-2, -1]).repeat_interleave(
|
||||||
|
patch_size_squared, -1
|
||||||
|
)
|
||||||
|
image_std_tensor = image_std_tensor.squeeze([-2, -1]).repeat_interleave(
|
||||||
|
patch_size_squared, -1
|
||||||
|
)
|
||||||
|
|
||||||
|
if not image_mean_tensor.is_contiguous():
|
||||||
|
image_mean_tensor = image_mean_tensor.contiguous()
|
||||||
|
if not image_std_tensor.is_contiguous():
|
||||||
|
image_std_tensor = image_std_tensor.contiguous()
|
||||||
|
|
||||||
|
pixel_values = (
|
||||||
|
rescale_factor * pixel_values.to(torch.float32) - image_mean_tensor
|
||||||
|
) / image_std_tensor
|
||||||
|
pixel_values = pixel_values.to(hf_config.dtype)
|
||||||
|
return pixel_values
|
||||||
|
|
||||||
|
def process_mm_data(
|
||||||
|
self, input_text, images=None, videos=None, audios=None, **kwargs
|
||||||
|
) -> dict:
|
||||||
|
"""
|
||||||
|
process multimodal data with transformers AutoProcessor
|
||||||
|
"""
|
||||||
|
if images:
|
||||||
|
kwargs["images"] = images
|
||||||
|
if videos:
|
||||||
|
kwargs["videos"] = videos
|
||||||
|
|
||||||
|
processor = self._processor
|
||||||
|
if (
|
||||||
|
hasattr(processor, "image_processor")
|
||||||
|
and isinstance(processor.image_processor, BaseImageProcessorFast)
|
||||||
|
and not self.server_args.disable_fast_image_processor
|
||||||
|
):
|
||||||
|
if not _is_npu:
|
||||||
|
kwargs["device"] = "cuda"
|
||||||
|
|
||||||
|
result = processor.__call__(
|
||||||
|
text=[input_text],
|
||||||
|
padding=True,
|
||||||
|
return_tensors="pt",
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Divide the processor_output into two modalities: image and video.
|
||||||
|
if result is not None:
|
||||||
|
pixel_values = result["images"]
|
||||||
|
if pixel_values is not None:
|
||||||
|
result["images"] = self._pixel_values_norm(pixel_values, kwargs)
|
||||||
|
for key in list(result.keys()):
|
||||||
|
if result[key] is None:
|
||||||
|
del result[key]
|
||||||
|
continue
|
||||||
|
if key == "grid_thw":
|
||||||
|
grid_thw = result["grid_thw"]
|
||||||
|
pixel_values_all = result["images"]
|
||||||
|
# Identify elements where the first
|
||||||
|
# dimension is greater than 1 and
|
||||||
|
# treat them as the video modality
|
||||||
|
mask = grid_thw[:, 0] > 1
|
||||||
|
result["video_grid_thw"] = grid_thw[mask]
|
||||||
|
result["image_grid_thw"] = grid_thw[~mask]
|
||||||
|
image_patch_num = result["image_grid_thw"].prod(dim=1).sum()
|
||||||
|
result["pixel_values"] = pixel_values_all[:image_patch_num]
|
||||||
|
result["pixel_values_videos"] = pixel_values_all[image_patch_num:]
|
||||||
|
del result["images"]
|
||||||
|
del result["grid_thw"]
|
||||||
|
|
||||||
|
# del empty result
|
||||||
|
if result["image_grid_thw"].numel() == 0:
|
||||||
|
del result["image_grid_thw"]
|
||||||
|
if result["pixel_values"].numel() == 0:
|
||||||
|
del result["pixel_values"]
|
||||||
|
if result["video_grid_thw"].numel() == 0:
|
||||||
|
del result["video_grid_thw"]
|
||||||
|
if result["pixel_values_videos"].numel() == 0:
|
||||||
|
del result["pixel_values_videos"]
|
||||||
|
|
||||||
|
if not self.server_args.keep_mm_feature_on_device:
|
||||||
|
# move feature tensors to cpu
|
||||||
|
for feature_name in self.FEATURE_NAMES:
|
||||||
|
if SGL_USE_CUDA_IPC:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
if feature_name in result and isinstance(
|
||||||
|
result[feature_name], torch.Tensor
|
||||||
|
):
|
||||||
|
result[feature_name] = result[feature_name].to("cpu")
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
async def process_mm_data_async(
|
||||||
|
self,
|
||||||
|
image_data: List[Union[str, bytes]],
|
||||||
|
input_text,
|
||||||
|
request_obj,
|
||||||
|
*args,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
base_output = self.load_mm_data(
|
||||||
|
prompt=input_text,
|
||||||
|
image_data=image_data,
|
||||||
|
video_data=request_obj.video_data,
|
||||||
|
audio_data=request_obj.audio_data,
|
||||||
|
multimodal_tokens=self.mm_tokens,
|
||||||
|
)
|
||||||
|
|
||||||
|
# resize images if they are raw Image objects
|
||||||
|
resized_images = []
|
||||||
|
if base_output.images and isinstance(base_output.images[0], Image.Image):
|
||||||
|
for image in base_output.images:
|
||||||
|
resized_image = resize_image(image)
|
||||||
|
resized_images.append(resized_image)
|
||||||
|
base_output.images = resized_images
|
||||||
|
|
||||||
|
if base_output.videos:
|
||||||
|
videos_processed = [
|
||||||
|
await preprocess_video(video) for video in base_output.videos
|
||||||
|
]
|
||||||
|
base_output.videos, _ = map(list, zip(*videos_processed))
|
||||||
|
|
||||||
|
mm_items, input_ids, ret = self.process_and_combine_mm_data(
|
||||||
|
base_output, self.mm_tokens
|
||||||
|
)
|
||||||
|
|
||||||
|
input_ids = input_ids.flatten()
|
||||||
|
|
||||||
|
mrope_positions, mrope_position_delta = MRotaryEmbedding.get_rope_index_ernie45(
|
||||||
|
input_ids=input_ids.unsqueeze(0),
|
||||||
|
hf_config=self.hf_config,
|
||||||
|
image_grid_thw=getattr(ret, "image_grid_thw", None),
|
||||||
|
video_grid_thw=getattr(ret, "video_grid_thw", None),
|
||||||
|
)
|
||||||
|
mrope_positions = mrope_positions.squeeze(1)
|
||||||
|
|
||||||
|
assert (
|
||||||
|
input_ids.shape[0] == mrope_positions.shape[-1]
|
||||||
|
), "input_ids and mrope_positions should have the same length"
|
||||||
|
|
||||||
|
mm_inputs = {
|
||||||
|
"input_ids": input_ids.tolist(),
|
||||||
|
"mm_items": mm_items,
|
||||||
|
"im_start_id": self.image_start_token_id,
|
||||||
|
"im_end_id": self.image_end_token_id,
|
||||||
|
"im_token_id": self.mm_tokens.image_token_id,
|
||||||
|
"video_token_id": self.mm_tokens.video_token_id,
|
||||||
|
"mrope_positions": mrope_positions,
|
||||||
|
"mrope_position_delta": mrope_position_delta,
|
||||||
|
}
|
||||||
|
|
||||||
|
return mm_inputs
|
||||||
Reference in New Issue
Block a user