Qwen2.5-Omni: squash 5 commits (add components, refactor, load thinker embedding from weights, remove unused code)
Squashed commits:\n- 4e2d812 Add Qwen2.5-Omni model components\n- 351f833 load the thinker embedding from the weights, not rely on local .pt files\n- 2c32533 Refactor Qwen2.5-Omni model files\n- 9163a60 remove unused codes\n- 192bdf8 remove
This commit is contained in:
tzhouam committed
1 parent
7a85034d9b
commit
a3a77b6cc4
8 files changed
+4770
No files matched your search
Whitespace-only changes.
@@ -0,0 +1,699 @@
|
||||
import itertools
|
||||
from typing import Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding
|
||||
from vllm.model_executor.layers.rotary_embedding.common import apply_rotary_emb_dispatch
|
||||
|
||||
def _apply_rotary_emb(
|
||||
x: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
is_neox_style: bool,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: [num_tokens, num_heads, head_size]
|
||||
cos: [num_tokens, head_size // 2]
|
||||
sin: [num_tokens, head_size // 2]
|
||||
is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
|
||||
positional embeddings.
|
||||
"""
|
||||
cos = cos.unsqueeze(-2).to(x.dtype)
|
||||
sin = sin.unsqueeze(-2).to(x.dtype)
|
||||
if is_neox_style:
|
||||
x1, x2 = torch.chunk(x, 2, dim=-1)
|
||||
else:
|
||||
x1 = x[..., ::2]
|
||||
x2 = x[..., 1::2]
|
||||
o1 = x1 * cos - x2 * sin
|
||||
o2 = x2 * cos + x1 * sin
|
||||
if is_neox_style:
|
||||
return torch.cat((o1, o2), dim=-1)
|
||||
else:
|
||||
return torch.stack((o1, o2), dim=-1).flatten(-2)
|
||||
|
||||
class MRotaryEmbedding(RotaryEmbedding):
|
||||
"""Rotary Embedding with Multimodal Sections."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
head_size: int,
|
||||
rotary_dim: int,
|
||||
max_position_embeddings: int,
|
||||
base: float,
|
||||
is_neox_style: bool,
|
||||
dtype: torch.dtype,
|
||||
mrope_section: Optional[list[int]] = None,
|
||||
) -> None:
|
||||
# In Qwen2.5-VL, the maximum index value is related to the duration of
|
||||
# the input video. We enlarge max_position_embeddings to 4 times to get
|
||||
# a larger the cos and sin cache.
|
||||
self.cache_max_position_num = max_position_embeddings * 4
|
||||
super().__init__(head_size, rotary_dim, self.cache_max_position_num,
|
||||
base, is_neox_style, dtype)
|
||||
|
||||
self.mrope_section = mrope_section
|
||||
if self.mrope_section:
|
||||
print("Warning: mrope_section check is not disabled in Qwen2.5-Omni, this may cause errors, and should be removed in the future.")
|
||||
# assert sum(self.mrope_section) == rotary_dim // 2
|
||||
|
||||
def forward(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: Optional[torch.Tensor] = None,
|
||||
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
|
||||
"""PyTorch-native implementation equivalent to forward().
|
||||
|
||||
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
|
||||
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
|
||||
|
||||
cos = torch.cat([
|
||||
m[i]
|
||||
for i, m in enumerate(cos.split(self.mrope_section, dim=-1))
|
||||
],
|
||||
dim=-1)
|
||||
sin = torch.cat([
|
||||
m[i]
|
||||
for i, m in enumerate(sin.split(self.mrope_section, dim=-1))
|
||||
],
|
||||
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_dispatch(query_rot, cos, sin,
|
||||
# self.is_neox_style)
|
||||
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_dispatch(key_rot, cos, sin,
|
||||
# self.is_neox_style)
|
||||
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
|
||||
|
||||
@classmethod
|
||||
def get_input_positions(
|
||||
cls,
|
||||
input_tokens: list[int],
|
||||
hf_config: PretrainedConfig,
|
||||
image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
|
||||
video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
|
||||
second_per_grid_ts: Optional[list[float]],
|
||||
context_len: int = 0,
|
||||
seq_len: Optional[int] = None,
|
||||
audio_feature_lengths: Optional[torch.Tensor] = None,
|
||||
use_audio_in_video: bool = False,
|
||||
) -> tuple[list[list[int]], int]:
|
||||
"""Get mrope input positions and delta value."""
|
||||
|
||||
image_grid_thw = [] if image_grid_thw is None else image_grid_thw
|
||||
video_grid_thw = [] if video_grid_thw is None else video_grid_thw
|
||||
second_per_grid_ts = [] if second_per_grid_ts is None else \
|
||||
second_per_grid_ts
|
||||
|
||||
llm_positions, mrope_position_delta = \
|
||||
cls.get_input_positions_tensor(
|
||||
input_tokens=input_tokens,
|
||||
hf_config=hf_config,
|
||||
image_grid_thw=image_grid_thw,
|
||||
video_grid_thw=video_grid_thw,
|
||||
second_per_grid_ts=second_per_grid_ts,
|
||||
context_len=context_len,
|
||||
seq_len=seq_len,
|
||||
audio_feature_lengths=audio_feature_lengths,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
)
|
||||
|
||||
return llm_positions.tolist(), mrope_position_delta
|
||||
|
||||
@classmethod
|
||||
def get_input_positions_tensor(
|
||||
cls,
|
||||
input_tokens: list[int],
|
||||
hf_config: PretrainedConfig,
|
||||
image_grid_thw: Union[list[list[int]], torch.Tensor],
|
||||
video_grid_thw: Union[list[list[int]], torch.Tensor],
|
||||
second_per_grid_ts: list[float],
|
||||
context_len: int = 0,
|
||||
seq_len: Optional[int] = None,
|
||||
audio_feature_lengths: Optional[torch.Tensor] = None,
|
||||
use_audio_in_video: bool = False,
|
||||
) -> tuple[torch.Tensor, int]:
|
||||
from vllm.transformers_utils.config import thinker_uses_mrope
|
||||
if thinker_uses_mrope(hf_config):
|
||||
return cls._omni_get_input_positions_tensor(
|
||||
input_tokens=input_tokens,
|
||||
hf_config=hf_config,
|
||||
image_grid_thw=image_grid_thw,
|
||||
video_grid_thw=video_grid_thw,
|
||||
second_per_grid_ts=second_per_grid_ts,
|
||||
context_len=context_len,
|
||||
seq_len=seq_len,
|
||||
audio_feature_lengths=audio_feature_lengths,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
)
|
||||
elif hf_config.model_type in ["glm4v", "glm4v_moe"]:
|
||||
return cls._glm4v_get_input_positions_tensor(
|
||||
input_tokens=input_tokens,
|
||||
hf_config=hf_config,
|
||||
image_grid_thw=image_grid_thw,
|
||||
video_grid_thw=video_grid_thw,
|
||||
context_len=context_len,
|
||||
seq_len=seq_len,
|
||||
)
|
||||
else:
|
||||
return cls._vl_get_input_positions_tensor(
|
||||
input_tokens=input_tokens,
|
||||
hf_config=hf_config,
|
||||
image_grid_thw=image_grid_thw,
|
||||
video_grid_thw=video_grid_thw,
|
||||
second_per_grid_ts=second_per_grid_ts,
|
||||
context_len=context_len,
|
||||
seq_len=seq_len,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _glm4v_get_input_positions_tensor(
|
||||
cls,
|
||||
input_tokens: list[int],
|
||||
hf_config: PretrainedConfig,
|
||||
image_grid_thw: Union[list[list[int]], torch.Tensor],
|
||||
video_grid_thw: Union[list[list[int]], torch.Tensor],
|
||||
context_len: int = 0,
|
||||
seq_len: Optional[int] = None,
|
||||
) -> tuple[torch.Tensor, int]:
|
||||
"""Get mrope input positions and delta value for GLM4V."""
|
||||
|
||||
image_token_id = hf_config.image_token_id
|
||||
video_start_token_id = hf_config.video_start_token_id
|
||||
video_end_token_id = hf_config.video_end_token_id
|
||||
spatial_merge_size = hf_config.vision_config.spatial_merge_size
|
||||
llm_pos_ids_list: list = []
|
||||
|
||||
if not (image_grid_thw is None and video_grid_thw is None):
|
||||
if isinstance(image_grid_thw, torch.Tensor):
|
||||
image_grid_thw = image_grid_thw.tolist()
|
||||
|
||||
input_token_type: list[str] = []
|
||||
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 (video_check_flg is False):
|
||||
input_token_type.append("image")
|
||||
elif (token == image_token_id) and (video_check_flg is True):
|
||||
input_token_type.append("video")
|
||||
else:
|
||||
input_token_type.append("text")
|
||||
|
||||
input_type_group: list[tuple[str, int, int]] = []
|
||||
for key, group_iter in itertools.groupby(
|
||||
enumerate(input_token_type), lambda x: x[1]):
|
||||
group_list = list(group_iter)
|
||||
start_index = group_list[0][0]
|
||||
end_index = group_list[-1][0] + 1
|
||||
input_type_group.append((key, start_index, end_index))
|
||||
|
||||
video_frame_num = 1
|
||||
mm_data_idx = 0
|
||||
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[mm_data_idx][0],
|
||||
image_grid_thw[mm_data_idx][1],
|
||||
image_grid_thw[mm_data_idx][2],
|
||||
)
|
||||
llm_grid_t, llm_grid_h, llm_grid_w = \
|
||||
t, h // spatial_merge_size, w // spatial_merge_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)
|
||||
mm_data_idx += 1
|
||||
|
||||
elif modality_type == "video":
|
||||
t, h, w = (
|
||||
video_frame_num,
|
||||
image_grid_thw[mm_data_idx][1],
|
||||
image_grid_thw[mm_data_idx][2],
|
||||
)
|
||||
llm_grid_t, llm_grid_h, llm_grid_w = \
|
||||
t, h // spatial_merge_size, w // spatial_merge_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)
|
||||
|
||||
mm_data_idx += 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
|
||||
|
||||
else:
|
||||
text_len = len(input_tokens)
|
||||
llm_pos_ids_list.append(
|
||||
torch.arange(text_len).view(1, -1).expand(3, -1))
|
||||
|
||||
llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
|
||||
llm_positions = llm_positions[:, context_len:seq_len]
|
||||
mrope_position_delta = (llm_positions.max() + 1 -
|
||||
len(input_tokens)).item()
|
||||
return llm_positions, mrope_position_delta
|
||||
|
||||
@classmethod
|
||||
def _vl_get_input_positions_tensor(
|
||||
cls,
|
||||
input_tokens: list[int],
|
||||
hf_config: PretrainedConfig,
|
||||
image_grid_thw: Union[list[list[int]], torch.Tensor],
|
||||
video_grid_thw: Union[list[list[int]], torch.Tensor],
|
||||
second_per_grid_ts: list[float],
|
||||
context_len: int = 0,
|
||||
seq_len: Optional[int] = None,
|
||||
) -> tuple[torch.Tensor, int]:
|
||||
"""Get mrope input positions and delta value."""
|
||||
|
||||
image_token_id = hf_config.image_token_id
|
||||
video_token_id = hf_config.video_token_id
|
||||
vision_start_token_id = hf_config.vision_start_token_id
|
||||
spatial_merge_size = hf_config.vision_config.spatial_merge_size
|
||||
tokens_per_second = getattr(hf_config.vision_config,
|
||||
"tokens_per_second", 1.0)
|
||||
|
||||
input_tokens_tensor = torch.tensor(input_tokens)
|
||||
vision_start_indices = torch.argwhere(
|
||||
input_tokens_tensor == vision_start_token_id).squeeze(1)
|
||||
vision_tokens = input_tokens_tensor[vision_start_indices + 1]
|
||||
image_nums = (vision_tokens == image_token_id).sum()
|
||||
video_nums = (vision_tokens == video_token_id).sum()
|
||||
llm_pos_ids_list: list = []
|
||||
|
||||
st = 0
|
||||
remain_images, remain_videos = image_nums, video_nums
|
||||
|
||||
image_index, video_index = 0, 0
|
||||
for _ in range(image_nums + video_nums):
|
||||
video_second_per_grid_t = 0.0
|
||||
if image_token_id in input_tokens and remain_images > 0:
|
||||
ed_image = input_tokens.index(image_token_id, st)
|
||||
else:
|
||||
ed_image = len(input_tokens) + 1
|
||||
if video_token_id in input_tokens and remain_videos > 0:
|
||||
ed_video = input_tokens.index(video_token_id, st)
|
||||
else:
|
||||
ed_video = len(input_tokens) + 1
|
||||
if ed_image < ed_video:
|
||||
t, h, w = (
|
||||
image_grid_thw[image_index][0],
|
||||
image_grid_thw[image_index][1],
|
||||
image_grid_thw[image_index][2],
|
||||
)
|
||||
image_index += 1
|
||||
remain_images -= 1
|
||||
ed = ed_image
|
||||
else:
|
||||
t, h, w = (
|
||||
video_grid_thw[video_index][0],
|
||||
video_grid_thw[video_index][1],
|
||||
video_grid_thw[video_index][2],
|
||||
)
|
||||
video_second_per_grid_t = 1.0
|
||||
if second_per_grid_ts:
|
||||
video_second_per_grid_t = second_per_grid_ts[video_index]
|
||||
video_index += 1
|
||||
remain_videos -= 1
|
||||
ed = ed_video
|
||||
|
||||
llm_grid_t, llm_grid_h, llm_grid_w = \
|
||||
t, h // spatial_merge_size, w // spatial_merge_size
|
||||
text_len = ed - st
|
||||
|
||||
st_idx = llm_pos_ids_list[-1].max() + 1 if len(
|
||||
llm_pos_ids_list) > 0 else 0
|
||||
llm_pos_ids_list.append(
|
||||
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
|
||||
|
||||
t_index = (torch.arange(llm_grid_t).view(-1, 1).expand(
|
||||
-1, llm_grid_h * llm_grid_w) * video_second_per_grid_t *
|
||||
tokens_per_second).long().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]) + text_len + st_idx)
|
||||
st = ed + llm_grid_t * llm_grid_h * llm_grid_w
|
||||
|
||||
if st < len(input_tokens):
|
||||
st_idx = llm_pos_ids_list[-1].max() + 1 if len(
|
||||
llm_pos_ids_list) > 0 else 0
|
||||
text_len = len(input_tokens) - st
|
||||
llm_pos_ids_list.append(
|
||||
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
|
||||
|
||||
llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
|
||||
mrope_position_delta = (llm_positions.max() + 1 -
|
||||
len(input_tokens)).item()
|
||||
llm_positions = llm_positions[:, context_len:seq_len]
|
||||
|
||||
return llm_positions, mrope_position_delta
|
||||
|
||||
@classmethod
|
||||
def _omni_get_input_positions_tensor(
|
||||
cls,
|
||||
input_tokens: list[int],
|
||||
hf_config: PretrainedConfig,
|
||||
image_grid_thw: Union[list[list[int]], torch.Tensor],
|
||||
video_grid_thw: Union[list[list[int]], torch.Tensor],
|
||||
second_per_grid_ts: Optional[list[float]] = None,
|
||||
context_len: int = 0,
|
||||
seq_len: Optional[int] = None,
|
||||
audio_feature_lengths: Optional[torch.Tensor] = None,
|
||||
use_audio_in_video: bool = False,
|
||||
) -> tuple[torch.Tensor, int]:
|
||||
"""Get mrope input positions and delta value (Qwen2.5-Omni version).
|
||||
|
||||
Differences from MRotaryEmbedding:
|
||||
1. Add audio support (and related `audio_feature_lengths`).
|
||||
2. Add `use_audio_in_video` option to read audio from video inputs.
|
||||
In this case, audio and vision position ids will be split into
|
||||
chunks and interleaved.
|
||||
|
||||
Example:
|
||||
|
||||
(V_i are vision position ids, A_i are audio position ids)
|
||||
|
||||
|V_1 ... V_n|A_1 ... A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|...
|
||||
|vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |...
|
||||
"""
|
||||
|
||||
# TODO(fyabc): refactor and share more code with
|
||||
# _vl_get_input_positions_tensor.
|
||||
|
||||
thinker_config = hf_config.thinker_config
|
||||
audio_token_id = thinker_config.audio_token_index
|
||||
image_token_id = thinker_config.image_token_index
|
||||
video_token_id = thinker_config.video_token_index
|
||||
audio_start_token_id = thinker_config.audio_start_token_id
|
||||
audio_end_token_id = thinker_config.audio_end_token_id
|
||||
vision_start_token_id = thinker_config.vision_start_token_id
|
||||
vision_end_token_id = thinker_config.vision_end_token_id
|
||||
seconds_per_chunk = thinker_config.seconds_per_chunk
|
||||
spatial_merge_size = thinker_config.vision_config.spatial_merge_size
|
||||
tokens_per_second = getattr(thinker_config.vision_config,
|
||||
"tokens_per_second", 25)
|
||||
|
||||
if isinstance(image_grid_thw, list):
|
||||
image_grid_thw = torch.tensor(image_grid_thw)
|
||||
if isinstance(video_grid_thw, list):
|
||||
video_grid_thw = torch.tensor(video_grid_thw)
|
||||
|
||||
src_item = input_tokens
|
||||
audio_seqlens = audio_feature_lengths
|
||||
if not second_per_grid_ts:
|
||||
second_per_grid_ts = [1] * video_grid_thw.shape[0]
|
||||
audio_idx = 0
|
||||
video_idx = 0
|
||||
image_idx = 0
|
||||
new_src_item: list[int] = []
|
||||
llm_pos_ids_list: list[torch.Tensor] = []
|
||||
|
||||
idx = 0
|
||||
while idx < len(src_item):
|
||||
new_src_item_len = len(new_src_item)
|
||||
start_idx = llm_pos_ids_list[-1].max() + 1 if len(
|
||||
llm_pos_ids_list) > 0 else 0
|
||||
if src_item[idx] not in [
|
||||
audio_token_id, video_token_id, image_token_id
|
||||
]:
|
||||
if use_audio_in_video and idx > 0:
|
||||
if src_item[idx] == vision_end_token_id and \
|
||||
src_item[idx - 1] == audio_end_token_id:
|
||||
# processing the <|audio_eos|> before <|vision_eos|>
|
||||
start_idx -= 1
|
||||
elif src_item[idx] == audio_start_token_id and \
|
||||
src_item[idx - 1] == vision_start_token_id:
|
||||
# processing the <|audio_bos|> after <|vision_eos|>
|
||||
start_idx -= 1
|
||||
new_src_item.append(src_item[idx])
|
||||
llm_pos_ids = torch.tensor([start_idx],
|
||||
dtype=torch.long).expand(3, -1)
|
||||
llm_pos_ids_list.append(llm_pos_ids)
|
||||
elif src_item[idx] == audio_token_id:
|
||||
assert audio_seqlens is not None
|
||||
audio_seqlen = audio_seqlens[audio_idx]
|
||||
place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1)
|
||||
new_src_item.extend([audio_token_id] * place_num)
|
||||
llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx
|
||||
llm_pos_ids_list.append(llm_pos_ids)
|
||||
audio_idx += 1
|
||||
elif src_item[idx] == image_token_id:
|
||||
grid_t = image_grid_thw[image_idx][0]
|
||||
grid_hs = image_grid_thw[:, 1]
|
||||
grid_ws = image_grid_thw[:, 2]
|
||||
t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long()
|
||||
llm_pos_ids = cls._get_llm_pos_ids_for_vision(
|
||||
start_idx, image_idx, spatial_merge_size, t_index, grid_hs,
|
||||
grid_ws)
|
||||
llm_pos_ids_list.append(llm_pos_ids)
|
||||
vision_seqlen = image_grid_thw[image_idx].prod() // (
|
||||
spatial_merge_size**2)
|
||||
new_src_item.extend([image_token_id] * vision_seqlen)
|
||||
image_idx += 1
|
||||
elif src_item[idx] == video_token_id and not use_audio_in_video:
|
||||
grid_t = video_grid_thw[video_idx][0]
|
||||
grid_hs = video_grid_thw[:, 1]
|
||||
grid_ws = video_grid_thw[:, 2]
|
||||
t_index = (torch.arange(grid_t) *
|
||||
second_per_grid_ts[video_idx] *
|
||||
tokens_per_second).long()
|
||||
llm_pos_ids = cls._get_llm_pos_ids_for_vision(
|
||||
start_idx, video_idx, spatial_merge_size, t_index, grid_hs,
|
||||
grid_ws)
|
||||
llm_pos_ids_list.append(llm_pos_ids)
|
||||
vision_seqlen = video_grid_thw[video_idx].prod() // (
|
||||
spatial_merge_size**2)
|
||||
new_src_item.extend([video_token_id] * vision_seqlen)
|
||||
video_idx += 1
|
||||
else:
|
||||
# read audio from video
|
||||
assert audio_seqlens is not None
|
||||
audio_seqlen = audio_seqlens[audio_idx]
|
||||
vision_seqlen = video_grid_thw[video_idx].prod() // (
|
||||
spatial_merge_size**2)
|
||||
grid_t = video_grid_thw[video_idx][0]
|
||||
grid_h = video_grid_thw[video_idx][1]
|
||||
grid_w = video_grid_thw[video_idx][2]
|
||||
grid_hs = video_grid_thw[:, 1]
|
||||
grid_ws = video_grid_thw[:, 2]
|
||||
t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
|
||||
t_index = (torch.arange(grid_t) *
|
||||
second_per_grid_ts[video_idx] *
|
||||
tokens_per_second).long()
|
||||
t_index_split_chunk = cls._split_list_into_ranges(
|
||||
t_index, t_ntoken_per_chunk)
|
||||
place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2
|
||||
pure_audio_len = place_num - 2
|
||||
added_audio_len = 0
|
||||
audio_llm_pos_ids_list: list[torch.Tensor] = []
|
||||
for t_chunk in t_index_split_chunk:
|
||||
vision_ntoken_per_chunk = len(
|
||||
t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
|
||||
new_src_item.extend([video_token_id] *
|
||||
vision_ntoken_per_chunk)
|
||||
vision_llm_pos_ids_list = cls._get_llm_pos_ids_for_vision(
|
||||
start_idx, video_idx, spatial_merge_size, t_chunk,
|
||||
grid_hs, grid_ws).split(1, dim=1)
|
||||
llm_pos_ids_list.extend(vision_llm_pos_ids_list)
|
||||
new_src_item.extend(
|
||||
min(t_ntoken_per_chunk, pure_audio_len -
|
||||
added_audio_len) * [audio_token_id])
|
||||
audio_start_idx = start_idx if len(
|
||||
audio_llm_pos_ids_list
|
||||
) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1
|
||||
if min(t_ntoken_per_chunk,
|
||||
pure_audio_len - added_audio_len) > 0:
|
||||
audio_llm_pos_ids_list = (torch.arange(
|
||||
min(t_ntoken_per_chunk, pure_audio_len -
|
||||
added_audio_len)).expand(3, -1) +
|
||||
audio_start_idx).split(1,
|
||||
dim=1)
|
||||
else:
|
||||
audio_llm_pos_ids_list = []
|
||||
added_audio_len += min(t_ntoken_per_chunk,
|
||||
pure_audio_len - added_audio_len)
|
||||
llm_pos_ids_list.extend(audio_llm_pos_ids_list)
|
||||
if added_audio_len < pure_audio_len:
|
||||
new_src_item.extend(
|
||||
(pure_audio_len - added_audio_len) * [audio_token_id])
|
||||
audio_llm_pos_ids_list = (
|
||||
torch.arange(pure_audio_len - added_audio_len).expand(
|
||||
3, -1) + llm_pos_ids_list[-1].max() + 1).split(
|
||||
1, dim=1)
|
||||
llm_pos_ids_list.extend(audio_llm_pos_ids_list)
|
||||
audio_idx += 1
|
||||
video_idx += 1
|
||||
# move to the next token
|
||||
idx += len(new_src_item) - new_src_item_len
|
||||
|
||||
llm_positions = torch.cat(llm_pos_ids_list, dim=1)
|
||||
mrope_position_delta = torch.cat(llm_pos_ids_list,
|
||||
dim=1).max() + 1 - len(src_item)
|
||||
llm_positions = llm_positions[:, context_len:seq_len]
|
||||
|
||||
return llm_positions, mrope_position_delta
|
||||
|
||||
@staticmethod
|
||||
def _get_llm_pos_ids_for_vision(
|
||||
start_idx: int,
|
||||
vision_idx: int,
|
||||
spatial_merge_size: int,
|
||||
t_index: list[int],
|
||||
grid_hs: torch.Tensor,
|
||||
grid_ws: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
llm_pos_ids_list = []
|
||||
llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
|
||||
llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
|
||||
h_index = (torch.arange(llm_grid_h).view(1, -1, 1).expand(
|
||||
len(t_index), -1, llm_grid_w).flatten())
|
||||
w_index = (torch.arange(llm_grid_w).view(1, 1, -1).expand(
|
||||
len(t_index), llm_grid_h, -1).flatten())
|
||||
t_index_tensor = torch.Tensor(t_index).to(llm_grid_h.device).view(
|
||||
-1, 1).expand(-1, llm_grid_h * llm_grid_w).long().flatten()
|
||||
_llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
|
||||
llm_pos_ids_list.append(_llm_pos_ids + start_idx)
|
||||
llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
|
||||
return llm_pos_ids
|
||||
|
||||
@staticmethod
|
||||
def _split_list_into_ranges(lst: torch.Tensor,
|
||||
interval: int) -> list[list[int]]:
|
||||
ranges: list[list[int]] = [[]
|
||||
for _ in range((max(lst) // interval) + 1)]
|
||||
for num in lst:
|
||||
index = num // interval
|
||||
ranges[index].append(num)
|
||||
return ranges
|
||||
|
||||
@staticmethod
|
||||
def get_next_input_positions(
|
||||
mrope_position_delta: int,
|
||||
context_len: int,
|
||||
seq_len: int,
|
||||
) -> list[list[int]]:
|
||||
return [
|
||||
list(
|
||||
range(context_len + mrope_position_delta,
|
||||
seq_len + mrope_position_delta)) for _ in range(3)
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def get_next_input_positions_tensor(out: np.ndarray, out_offset: int,
|
||||
mrope_position_delta: int,
|
||||
context_len: int, num_new_tokens: int):
|
||||
|
||||
values = np.arange(mrope_position_delta + context_len,
|
||||
mrope_position_delta + context_len + num_new_tokens,
|
||||
dtype=out.dtype)
|
||||
out[:, out_offset:out_offset + num_new_tokens] = values
|
||||
|
||||
@classmethod
|
||||
def omni_get_updates_use_audio_in_video(
|
||||
cls,
|
||||
thinker_config: PretrainedConfig,
|
||||
audio_len: int,
|
||||
video_grid_thw: Union[list[int], torch.Tensor],
|
||||
video_second_per_grid_t: float,
|
||||
) -> list[int]:
|
||||
"""Get video prompt updates when `use_audio_in_video` is True.
|
||||
|
||||
In this case, audio and vision update ids will be split into
|
||||
chunks and interleaved (details in `_omni_get_input_positions_tensor`).
|
||||
|
||||
<|video_bos|><|VIDEO|><|video_eos|> =>
|
||||
<|video_bos|><|audio_bos|>(... chunks ...)<|audio_eos|><|video_eos|>
|
||||
"""
|
||||
|
||||
audio_token_id = thinker_config.audio_token_index
|
||||
video_token_id = thinker_config.video_token_index
|
||||
audio_start_token_id = thinker_config.audio_start_token_id
|
||||
audio_end_token_id = thinker_config.audio_end_token_id
|
||||
seconds_per_chunk = thinker_config.seconds_per_chunk
|
||||
spatial_merge_size = thinker_config.vision_config.spatial_merge_size
|
||||
tokens_per_second = getattr(thinker_config.vision_config,
|
||||
"tokens_per_second", 25)
|
||||
|
||||
grid_t = video_grid_thw[0]
|
||||
grid_h = video_grid_thw[1]
|
||||
grid_w = video_grid_thw[2]
|
||||
t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
|
||||
t_index = (torch.arange(grid_t) * video_second_per_grid_t *
|
||||
tokens_per_second).long()
|
||||
t_index_split_chunk = cls._split_list_into_ranges(
|
||||
t_index, t_ntoken_per_chunk)
|
||||
|
||||
updates = [audio_start_token_id]
|
||||
added_audio_len = 0
|
||||
for t_chunk in t_index_split_chunk:
|
||||
vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (
|
||||
spatial_merge_size**2)
|
||||
updates.extend([video_token_id] * vision_ntoken_per_chunk)
|
||||
|
||||
audio_chunk_size = min(t_ntoken_per_chunk,
|
||||
audio_len - added_audio_len)
|
||||
updates.extend(audio_chunk_size * [audio_token_id])
|
||||
added_audio_len += audio_chunk_size
|
||||
if added_audio_len < audio_len:
|
||||
updates.extend((audio_len - added_audio_len) * [audio_token_id])
|
||||
updates.extend([audio_end_token_id])
|
||||
|
||||
return updates
|
||||
@@ -0,0 +1,706 @@
|
||||
from functools import cached_property
|
||||
from typing import Iterable, List, Optional, Set, Tuple, Union, NamedTuple, Dict
|
||||
|
||||
import os
|
||||
import glob
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
|
||||
Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig, Qwen2_5OmniTalkerConfig)
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear
|
||||
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
|
||||
from vllm.model_executor.sampling_metadata import SamplingMetadata
|
||||
from vllm.model_executor.models.qwen2_5_omni_thinker import (
|
||||
Qwen2_5OmniConditionalGenerationMixin,
|
||||
Qwen2_5OmniThinkerMultiModalProcessor,
|
||||
Qwen2_5OmniThinkerProcessingInfo,
|
||||
Qwen2_5OmniThinkerDummyInputsBuilder)
|
||||
# from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
|
||||
from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
|
||||
init_vllm_registered_model,
|
||||
maybe_prefix)
|
||||
from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf
|
||||
from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
|
||||
|
||||
class OmniOutput(NamedTuple):
|
||||
"""Output from the merged Omni model containing both text and audio."""
|
||||
text_hidden_states: torch.Tensor
|
||||
multimodal_outputs: dict = {}
|
||||
intermediate_tensors: Optional[IntermediateTensors] = None
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@MULTIMODAL_REGISTRY.register_processor(
|
||||
Qwen2_5OmniThinkerMultiModalProcessor,
|
||||
info=Qwen2_5OmniThinkerProcessingInfo,
|
||||
dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
|
||||
)
|
||||
class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
|
||||
SupportsPP,
|
||||
Qwen2_5OmniConditionalGenerationMixin):
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
self.have_multimodal_outputs = True
|
||||
config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config
|
||||
quant_config = vllm_config.quant_config
|
||||
multimodal_config = vllm_config.model_config.multimodal_config
|
||||
# keep vllm_config for later submodule init
|
||||
self.vllm_config = vllm_config
|
||||
|
||||
# Initialize thinker components
|
||||
thinker_config: Qwen2_5OmniThinkerConfig = config.thinker_config
|
||||
self.thinker_config = thinker_config
|
||||
self.multimodal_config = multimodal_config
|
||||
|
||||
# Initialize talker components
|
||||
talker_config: Qwen2_5OmniTalkerConfig = config.talker_config
|
||||
self.talker_config = talker_config
|
||||
|
||||
|
||||
|
||||
self.model_stage = vllm_config.model_config.model_stage
|
||||
if self.model_stage=="thinker":
|
||||
# Initialize thinker model (multimodal processing)
|
||||
self.thinker = init_vllm_registered_model(
|
||||
vllm_config=vllm_config,
|
||||
prefix=maybe_prefix(prefix, "thinker"),
|
||||
hf_config=thinker_config,
|
||||
# Use registry architecture key
|
||||
architectures=["Qwen2_5OmniThinkerModel"],
|
||||
)
|
||||
self.model = self.thinker
|
||||
self.talker = None
|
||||
self.token2wav = None
|
||||
|
||||
elif self.model_stage=="talker":
|
||||
self.thinker = None
|
||||
# Initialize talker model wrapper (handles projection + LM)
|
||||
self.talker = init_vllm_registered_model(
|
||||
vllm_config=vllm_config,
|
||||
prefix=maybe_prefix(prefix, "talker"),
|
||||
hf_config=talker_config,
|
||||
# Use registry architecture key
|
||||
architectures=["Qwen2_5OmniTalkerModel"],
|
||||
)
|
||||
self.talker.init_multi_modal(thinker_config)
|
||||
self.model=self.talker
|
||||
self.token2wav = None
|
||||
|
||||
elif self.model_stage=="code2wav":
|
||||
self.thinker = None
|
||||
self.talker = None
|
||||
# Initialize token2wav (code->mel->wav) like thinker/talker
|
||||
self.token2wav_config = getattr(config, 'token2wav_config', None)
|
||||
self.token2wav = None
|
||||
if self.token2wav_config is not None:
|
||||
self.token2wav = init_vllm_registered_model(
|
||||
vllm_config=vllm_config,
|
||||
prefix=maybe_prefix(prefix, "token2wav"),
|
||||
hf_config=self.token2wav_config,
|
||||
architectures=["Qwen2_5OmniToken2WavModel"],
|
||||
)
|
||||
# voice resources (loaded on demand)
|
||||
self._token2wav_conds: Dict[str, torch.Tensor] = {}
|
||||
self._token2wav_ref_mels: Dict[str, torch.Tensor] = {}
|
||||
self.model = self.token2wav
|
||||
else:
|
||||
raise ValueError("Invalid model stage")
|
||||
|
||||
# Set up intermediate tensors
|
||||
self.make_empty_intermediate_tensors = (
|
||||
self.thinker.make_empty_intermediate_tensors) if self.model_stage=="thinker" else lambda: None
|
||||
|
||||
# -------------------- Device utilities --------------------
|
||||
@staticmethod
|
||||
def _module_device(module: nn.Module) -> torch.device:
|
||||
try:
|
||||
return next(module.parameters()).device
|
||||
except StopIteration:
|
||||
# No parameters; fall back to buffers or cpu
|
||||
for _, buf in module.named_buffers(recurse=True):
|
||||
return buf.device
|
||||
return torch.device("cpu")
|
||||
|
||||
def move_submodules_to_devices(
|
||||
self,
|
||||
*,
|
||||
thinker_device: Optional[Union[str, torch.device]] = None,
|
||||
talker_device: Optional[Union[str, torch.device]] = None,
|
||||
token2wav_device: Optional[Union[str, torch.device]] = None,
|
||||
) -> None:
|
||||
"""Optionally move thinker/talker/token2wav to different devices.
|
||||
|
||||
Example:
|
||||
model.move_submodules_to_devices(
|
||||
thinker_device='cuda:0',
|
||||
talker_device='cuda:1',
|
||||
token2wav_device='cpu',
|
||||
)
|
||||
"""
|
||||
if thinker_device is not None and self.thinker is not None:
|
||||
self.thinker.to(thinker_device)
|
||||
if talker_device is not None and self.talker is not None:
|
||||
self.talker.to(talker_device)
|
||||
if token2wav_device is not None and self.token2wav is not None:
|
||||
self.token2wav.to(token2wav_device)
|
||||
|
||||
@cached_property
|
||||
def sampler(self):
|
||||
if hasattr(self.model, "sampler"):
|
||||
return self.model.sampler
|
||||
return get_sampler()
|
||||
|
||||
def get_input_embeddings(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings=None,
|
||||
) -> torch.Tensor:
|
||||
if self.model_stage == "code2wav":
|
||||
return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size())
|
||||
return self.model.get_input_embeddings(
|
||||
input_ids, multimodal_embeddings)
|
||||
|
||||
|
||||
def get_multimodal_embeddings(self, **kwargs):
|
||||
# Delegate to thinker model for multimodal processing
|
||||
return self.model.get_multimodal_embeddings(**kwargs)
|
||||
|
||||
def last_index_of(self, list, value):
|
||||
return len(list) - 1 - list[::-1].index(value)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: Optional[IntermediateTensors] = None,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
generate_audio: bool = True,
|
||||
voice_type: str = "Chelsie",
|
||||
codec: Optional[torch.Tensor] = None,
|
||||
sampling_metadata: Optional[SamplingMetadata] = None,
|
||||
logits_index: Optional[int] = None,
|
||||
sampler = None,
|
||||
additional_information: Optional[dict[str, object]] = None,
|
||||
**kwargs: object,
|
||||
) -> Union[torch.Tensor, IntermediateTensors, OmniOutput]:
|
||||
"""
|
||||
Workflow:
|
||||
1) Thinker: multimodal understanding → text hidden states.
|
||||
2) If audio requested and codec not provided, use talker to derive codec.
|
||||
3) If audio requested (or codec provided), use token2wav to synthesize waveform.
|
||||
4) Return text hidden states (and audio when applicable).
|
||||
"""
|
||||
if self.model_stage=="thinker":
|
||||
# Normalize to batched inputs if caller provides 1D/2D unbatched tensors
|
||||
added_batch_dim = False
|
||||
if input_ids is not None and input_ids.ndim == 1:
|
||||
input_ids = input_ids.unsqueeze(0)
|
||||
added_batch_dim = True
|
||||
if positions is not None and positions.ndim == 1:
|
||||
positions = positions.unsqueeze(0)
|
||||
added_batch_dim = True
|
||||
if inputs_embeds is not None and inputs_embeds.ndim == 2:
|
||||
inputs_embeds = inputs_embeds.unsqueeze(0)
|
||||
added_batch_dim = True
|
||||
thinker_dev = self._module_device(self.thinker)
|
||||
|
||||
#if input_ids is None, set it to a zero tensor, in the length of the same as the embedding seq length
|
||||
if input_ids is None:
|
||||
input_ids = torch.zeros(inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev).unsqueeze(0) #(1, 0)
|
||||
added_batch_dim = True
|
||||
|
||||
# 1) Thinker (ensure inputs on thinker's device)
|
||||
if input_ids is not None and input_ids.device != thinker_dev:
|
||||
input_ids = input_ids.to(thinker_dev)
|
||||
if positions is not None and positions.device != thinker_dev:
|
||||
positions = positions.to(thinker_dev)
|
||||
if inputs_embeds is not None and inputs_embeds.device != thinker_dev:
|
||||
inputs_embeds = inputs_embeds.to(thinker_dev)
|
||||
# Run thinker
|
||||
thinker_output = self.thinker(
|
||||
input_ids=input_ids,
|
||||
positions=positions[0],
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if isinstance(thinker_output, tuple):
|
||||
embeds, text_hidden_states = thinker_output
|
||||
else:
|
||||
text_hidden_states = thinker_output
|
||||
|
||||
# Text-only path
|
||||
return OmniOutput(
|
||||
text_hidden_states=text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states,
|
||||
multimodal_outputs=None
|
||||
)
|
||||
|
||||
# 2) Talker (if codec not provided)
|
||||
if self.model_stage=="talker":
|
||||
|
||||
if input_ids is None and additional_information is None:
|
||||
input_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device)
|
||||
additional_information = {}
|
||||
self.thinker_reply_part = torch.zeros_like(inputs_embeds)
|
||||
is_profile = True
|
||||
else:
|
||||
is_profile = False
|
||||
|
||||
if input_ids is not None and additional_information is not None and not is_profile:
|
||||
# read from additional_information dict
|
||||
thinker_result = None
|
||||
if additional_information is not None and isinstance(additional_information, dict):
|
||||
thinker_result = additional_information.get("thinker_result")
|
||||
prompt_embeds = additional_information.get("prompt_embeds")
|
||||
prompt_token_ids = additional_information.get("prompt_token_ids")
|
||||
thinker_output_token_ids = additional_information.get("thinker_output_token_ids")
|
||||
else:
|
||||
thinker_result = torch.zeros_like(inputs_embeds)
|
||||
prompt_embeds = torch.zeros_like(inputs_embeds)
|
||||
prompt_token_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.int64, device=inputs_embeds.device)
|
||||
thinker_output_token_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.int64, device=inputs_embeds.device)
|
||||
|
||||
if thinker_result is None:
|
||||
thinker_result = torch.zeros_like(inputs_embeds)
|
||||
self.thinker_reply_part = thinker_result.squeeze(0)
|
||||
if self.thinker_reply_part.shape[1] > 1:
|
||||
self.thinker_reply_part = self.thinker_reply_part[1:, :]
|
||||
input_ids, inputs_embeds = self._thinker_to_talker_prefill(
|
||||
voice_type=voice_type,
|
||||
output_prompt_embeds=thinker_result,
|
||||
output_token_ids = thinker_output_token_ids,
|
||||
thinker_prompt_embeds=prompt_embeds,
|
||||
prompt_token_ids = prompt_token_ids,
|
||||
)
|
||||
elif not is_profile:
|
||||
input_ids, inputs_embeds = \
|
||||
self._thinker_to_talker_decode_one_step(
|
||||
output_prompt_embeds=self.thinker_reply_part[:1] if self.thinker_reply_part.shape[0]>=1 \
|
||||
else torch.zeros(1, self.thinker_reply_part.shape[1]).to(self._module_device(self.model)).to(torch.bfloat16)+(-1.25*2**(-123)),
|
||||
output_token_ids=input_ids,
|
||||
)
|
||||
|
||||
if self.thinker_reply_part.shape[0] >=1:
|
||||
self.thinker_reply_part = self.thinker_reply_part[1:, :]
|
||||
|
||||
with torch.inference_mode():
|
||||
talker_hidden = self.talker(
|
||||
input_ids=input_ids,
|
||||
positions=positions[0],
|
||||
inputs_embeds=inputs_embeds,
|
||||
|
||||
)
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=talker_hidden,
|
||||
multimodal_outputs=None
|
||||
)
|
||||
|
||||
if self.model_stage=="code2wav":
|
||||
code = input_ids if input_ids is not None else torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device)
|
||||
audio_tensor = self.generate_audio(code[:-1] if code[-1]==8294 else code, voice_type)
|
||||
return OmniOutput(
|
||||
text_hidden_states = None,
|
||||
multimodal_outputs = {
|
||||
"audio": audio_tensor
|
||||
}
|
||||
)
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=torch.cat(
|
||||
[
|
||||
torch.zeros([inputs_embeds.shape[0],896], dtype=torch.bfloat16).to(self._module_device(self.model)),
|
||||
self.talker.thinker_to_talker_proj(self.talker.get_input_embeddings(torch.tensor([8294,8293]).to(torch.bfloat16).to(self._module_device(self.model))))[0]
|
||||
],
|
||||
dim=0),
|
||||
multimodal_outputs=None
|
||||
)
|
||||
|
||||
def generate_audio(self, code, voice_type):
|
||||
token2wav_dev = self._module_device(self.token2wav)
|
||||
if isinstance(code, torch.Tensor):
|
||||
code_tensor = code.to(dtype=torch.long, device=token2wav_dev)
|
||||
else:
|
||||
code_tensor = torch.as_tensor(code, dtype=torch.long, device=token2wav_dev)
|
||||
if code_tensor.ndim == 2 and code_tensor.shape[0] == 1:
|
||||
code_tensor = code_tensor.squeeze(0)
|
||||
|
||||
audio_tensor = self._codec_to_audio(code_tensor, voice_type)
|
||||
|
||||
return audio_tensor
|
||||
|
||||
def _load_talker_embedding(
|
||||
self,
|
||||
) -> torch.nn.Embedding:
|
||||
return self.talker.language_model.model.embed_tokens
|
||||
|
||||
def _init_special_tokens_embeddings(
|
||||
self,
|
||||
):
|
||||
# talker embeddings
|
||||
self.talker_embedding = self._load_talker_embedding()
|
||||
|
||||
# embed_text_bos_token
|
||||
self.tts_text_spk_token_ids = {
|
||||
# M02：我是个会说标准普通话、带部分北方口音的男声
|
||||
'm02': 151870,
|
||||
'Ethan': 151870,
|
||||
|
||||
# F030：我是你的二次元虚拟女友
|
||||
'f030': 151872,
|
||||
'Chelsie': 151872,
|
||||
}
|
||||
self.default_tts_text_spk_type = list(
|
||||
self.tts_text_spk_token_ids.keys())[0]
|
||||
self.tts_text_spk_token_ids['prefix_caching'] = 151870
|
||||
|
||||
talker_hf_config = self.talker_config
|
||||
if hasattr(talker_hf_config, 'talker_config'):
|
||||
talker_hf_config = talker_hf_config.talker_config
|
||||
|
||||
self.embed_text_bos_token = self.thinker_embedding(
|
||||
torch.tensor(
|
||||
[talker_hf_config.tts_text_start_token_id],
|
||||
dtype=torch.long,
|
||||
device=self._module_device(self.talker),
|
||||
))
|
||||
self.embed_text_spk_tokens = {
|
||||
key:
|
||||
self.thinker_embedding(
|
||||
torch.tensor(
|
||||
[value],
|
||||
dtype=torch.long,
|
||||
device=self._module_device(self.talker),
|
||||
))
|
||||
for key, value in self.tts_text_spk_token_ids.items()
|
||||
}
|
||||
self.embed_text_eos_token = self.thinker_embedding(
|
||||
torch.tensor(
|
||||
[talker_hf_config.tts_text_end_token_id],
|
||||
dtype=torch.long,
|
||||
device=self._module_device(self.talker),
|
||||
))
|
||||
self.embed_text_pad_token = self.thinker_embedding(
|
||||
torch.tensor(
|
||||
[talker_hf_config.tts_text_pad_token_id],
|
||||
dtype=torch.long,
|
||||
device=self._module_device(self.talker),
|
||||
))
|
||||
self.embed_codec_bos_token = self.talker_embedding(
|
||||
torch.tensor(
|
||||
[talker_hf_config.tts_codec_start_token_id],
|
||||
dtype=torch.long,
|
||||
device=self._module_device(self.talker),
|
||||
))
|
||||
self.embed_codec_eos_token = self.talker_embedding(
|
||||
torch.tensor(
|
||||
[talker_hf_config.tts_codec_end_token_id],
|
||||
dtype=torch.long,
|
||||
device=self._module_device(self.talker),
|
||||
))
|
||||
self.embed_codec_pad_token = self.talker_embedding(
|
||||
torch.tensor(
|
||||
[talker_hf_config.tts_codec_pad_token_id],
|
||||
dtype=torch.long,
|
||||
device=self._module_device(self.talker),
|
||||
))
|
||||
return set(["thinker_embedding.weight", "talker_embedding.weight"])
|
||||
|
||||
def _get_embed_text_spk_token(self, voice_type: str):
|
||||
if voice_type not in self.embed_text_spk_tokens:
|
||||
return self.embed_text_bos_token
|
||||
return self.embed_text_spk_tokens[voice_type]
|
||||
|
||||
def _get_text_spk_token_id(self, voice_type: str):
|
||||
talker_hf_config = self.talker_config
|
||||
if hasattr(talker_hf_config, 'talker_config'):
|
||||
talker_hf_config = talker_hf_config.talker_config
|
||||
|
||||
if voice_type not in self.tts_text_spk_token_ids:
|
||||
return talker_hf_config.tts_text_start_token_id
|
||||
return self.tts_text_spk_token_ids[voice_type]
|
||||
|
||||
def _thinker_to_talker_prefill(
|
||||
self,
|
||||
voice_type: str,
|
||||
output_prompt_embeds,
|
||||
output_token_ids,
|
||||
thinker_prompt_embeds,
|
||||
prompt_token_ids
|
||||
):
|
||||
|
||||
talker_hf_config = self.talker_config
|
||||
if hasattr(talker_hf_config, 'talker_config'):
|
||||
talker_hf_config = talker_hf_config.talker_config
|
||||
|
||||
# if len(output.outputs[0].token_ids) == 2:
|
||||
# issue request
|
||||
prompt_embeds = torch.cat([
|
||||
thinker_prompt_embeds,
|
||||
self._get_embed_text_spk_token(voice_type) +
|
||||
self.embed_codec_pad_token,
|
||||
output_prompt_embeds[:1] + self.embed_codec_bos_token,
|
||||
],
|
||||
dim=0)
|
||||
|
||||
prompt_token_ids_processed = prompt_token_ids + [
|
||||
talker_hf_config.tts_codec_pad_token_id,
|
||||
output_token_ids[0],
|
||||
|
||||
]
|
||||
input_tokens_len = len(prompt_token_ids_processed)
|
||||
# the code below is from model runner in Qwen, may need to further discuss later
|
||||
if input_tokens_len > 2:
|
||||
prompt_token_ids_processed = (
|
||||
[self.talker_config.tts_codec_mask_token_id] *
|
||||
(input_tokens_len - 2) + [
|
||||
self.talker_config.tts_codec_pad_token_id,
|
||||
self.talker_config.tts_codec_start_token_id
|
||||
])
|
||||
else:
|
||||
prompt_token_ids_processed = [
|
||||
self.talker_config.tts_codec_pad_token_id,
|
||||
self.talker_config.tts_codec_start_token_id,
|
||||
][-input_tokens_len:]
|
||||
if isinstance(prompt_token_ids_processed,list):
|
||||
prompt_token_ids_processed = torch.Tensor(prompt_token_ids_processed).to(torch.int64).to(self._module_device(self.talker))
|
||||
return prompt_token_ids_processed, prompt_embeds
|
||||
|
||||
def _thinker_to_talker_decode_one_step(
|
||||
self,
|
||||
output_prompt_embeds,
|
||||
output_token_ids,
|
||||
):
|
||||
processed_output_token_embeds = output_prompt_embeds + self.talker.get_input_embeddings(output_token_ids) #for decode
|
||||
return output_token_ids, processed_output_token_embeds
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: Union[torch.Tensor, OmniOutput],
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> Optional[torch.Tensor]:
|
||||
# Handle OmniOutput type
|
||||
if isinstance(hidden_states, OmniOutput):
|
||||
hidden_states = hidden_states.text_hidden_states
|
||||
|
||||
# Use thinker model for logits computation
|
||||
return self.model.compute_logits(hidden_states, sampling_metadata)
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> Optional[SamplerOutput]:
|
||||
# Use thinker model for sampling
|
||||
return self.model.sample(logits, sampling_metadata)
|
||||
|
||||
def generate_speech(self, text_tokens: torch.Tensor, voice_type: str = "default"):
|
||||
"""
|
||||
Generate speech from text tokens using the talker and token2wav models.
|
||||
This method is kept for backward compatibility and direct speech generation.
|
||||
|
||||
Args:
|
||||
text_tokens: Text tokens from thinker model
|
||||
voice_type: Voice type for speech generation
|
||||
|
||||
Returns:
|
||||
Audio tensor
|
||||
"""
|
||||
# Generate codec tokens using talker model
|
||||
talker_output = self.talker(
|
||||
input_ids=None,
|
||||
positions=None,
|
||||
inputs_embeds=text_tokens
|
||||
)
|
||||
|
||||
# Convert talker output to codec tokens
|
||||
codec_tokens = self._convert_to_codec_tokens(talker_output)
|
||||
|
||||
# Generate audio using token2wav model
|
||||
return self._codec_to_audio(codec_tokens, voice_type=voice_type)
|
||||
|
||||
|
||||
def _convert_to_codec_tokens(self, talker_output: torch.Tensor, sampling_metadata: SamplingMetadata) -> torch.Tensor:
|
||||
"""
|
||||
参考 HF：使用 talker 的 codec 头得到 logits，抑制 BOS，再贪心选取当前步的下一个 codec token。
|
||||
"""
|
||||
with torch.inference_mode():
|
||||
logits = self.talker.compute_logits(talker_output, None)
|
||||
if logits is None:
|
||||
return torch.zeros((talker_output.size(0), 0), dtype=torch.long, device=talker_output.device)
|
||||
|
||||
# 仅抑制 codec_bos，与 HF generate 的 suppress_tokens 行为一致
|
||||
bos_id = None
|
||||
if hasattr(self, 'talker_config') and hasattr(self.talker_config, 'tts_codec_start_token_id'):
|
||||
bos_id = int(getattr(self.talker_config, 'tts_codec_start_token_id'))
|
||||
if bos_id is not None:
|
||||
logits[..., bos_id] = -1e9
|
||||
|
||||
# 取最后一步位置的分布并贪心选取
|
||||
next_id = self.talker.sample(logits, sampling_metadata).sampled_token_ids
|
||||
return next_id.to(dtype=torch.long)
|
||||
|
||||
def _init_token2wav_model(self):
|
||||
"""Initialize speaker resources if provided; model is constructed in __init__."""
|
||||
if self.token2wav is None or self.token2wav_config is None:
|
||||
return
|
||||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
# optional speaker resources
|
||||
conds = getattr(self.token2wav_config, 'conds', None)
|
||||
ref_mels = getattr(self.token2wav_config, 'ref_mels', None)
|
||||
if isinstance(conds, dict) and isinstance(ref_mels, dict):
|
||||
self._token2wav_conds = {k: torch.as_tensor(v, device=device) for k, v in conds.items()}
|
||||
self._token2wav_ref_mels = {k: torch.as_tensor(v, device=device) for k, v in ref_mels.items()}
|
||||
# legacy: load from directory if provided
|
||||
model_path = self.vllm_config.model_config.model
|
||||
if isinstance(model_path, str) and os.path.isdir(model_path):
|
||||
spk_pt = os.path.join(model_path, 'spk_dict.pt')
|
||||
if os.path.exists(spk_pt):
|
||||
data = torch.load(spk_pt, map_location=device)
|
||||
for key, value in data.items():
|
||||
self._token2wav_conds[key] = value["cond"].to(device)
|
||||
self._token2wav_ref_mels[key] = value["ref_mel"].to(device)
|
||||
else:
|
||||
# legacy npy inputs
|
||||
for f in sorted(glob.glob(os.path.join(model_path, 'inputs', '*spk_emb.npy'))):
|
||||
key = os.path.basename(f).split('_')[0].lower()
|
||||
self._token2wav_conds[key] = torch.as_tensor(np.load(f), device=device)
|
||||
for f in sorted(glob.glob(os.path.join(model_path, 'inputs', '*ref_mel.npy'))):
|
||||
key = os.path.basename(f).split('_')[0].lower()
|
||||
self._token2wav_ref_mels[key] = torch.as_tensor(np.load(f), device=device)
|
||||
|
||||
def _codec_to_audio(self, codec_tokens: torch.Tensor, voice_type: str = "default") -> Optional[torch.Tensor]:
|
||||
if self.token2wav is None:
|
||||
self._init_token2wav_model()
|
||||
if self.token2wav is None:
|
||||
return None
|
||||
# Normalize voice type
|
||||
voice = (voice_type or 'default')
|
||||
# Resolve cond / ref_mel if provided
|
||||
cond = None
|
||||
ref_mel = None
|
||||
if voice in self._token2wav_conds and voice in self._token2wav_ref_mels:
|
||||
cond = self._token2wav_conds[voice]
|
||||
ref_mel = self._token2wav_ref_mels[voice]
|
||||
# Fallback: create dummy cond/ref_mel if not provided
|
||||
token2wav_dev = self._module_device(self.token2wav)
|
||||
if cond is None:
|
||||
cond = torch.zeros((1, self.token2wav_config.dit_config.enc_emb_dim), device=token2wav_dev, dtype=torch.float32)
|
||||
if ref_mel is None:
|
||||
ref_mel = torch.zeros((1, 300, self.token2wav_config.dit_config.mel_dim), device=token2wav_dev, dtype=torch.float32)
|
||||
|
||||
# Ensure codec is (1, T) long tensor on correct device
|
||||
if isinstance(codec_tokens, torch.Tensor):
|
||||
codec = codec_tokens.to(dtype=torch.long, device=token2wav_dev)
|
||||
if codec.ndim == 1:
|
||||
codec = codec.unsqueeze(0)
|
||||
else:
|
||||
codec = torch.as_tensor(codec_tokens, dtype=torch.long, device=token2wav_dev).unsqueeze(0)
|
||||
|
||||
# Streaming with chunked process and boundary alignment (rely on token2wav.process_chunk)
|
||||
factor = getattr(self.token2wav.token2wav.factor, 'factor', 2)
|
||||
chunk_size = 48
|
||||
mel_dim = getattr(self.token2wav.token2wav.code2wav_dit_model, 'mel_dim', self.token2wav_config.dit_config.mel_dim)
|
||||
total_mel = int(codec.shape[1] * factor)
|
||||
steps = 10
|
||||
|
||||
# Prepare initial noise for the whole sequence
|
||||
y_all = torch.randn((1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev)
|
||||
|
||||
logger.info(f"Currently, we do not use the chunked process, we only use the token2wav.process_chunk for the whole sequence.\
|
||||
The stream mode will be implemented in the future.")
|
||||
|
||||
chunk_ends = []
|
||||
for i in range(codec.shape[1]):
|
||||
chunk_code_length = i * 2 - 24
|
||||
finished = i==(codec.shape[1]-1)
|
||||
if (chunk_code_length > 0 and
|
||||
chunk_code_length % chunk_size == 0) or finished:
|
||||
chunk_ends.append(i)
|
||||
|
||||
# Number of chunks in mel domain
|
||||
prev_generated = None
|
||||
wav_chunks: list = []
|
||||
prev_id = 0
|
||||
|
||||
with torch.inference_mode():
|
||||
for n,i in enumerate([0]):
|
||||
finished = (i == codec.shape[1] - 1)
|
||||
_, audio_chunk = self.token2wav.process_chunk(
|
||||
conditioning=cond,
|
||||
reference_mel=ref_mel,
|
||||
codec_all=codec,
|
||||
y_all=y_all,
|
||||
i=n,
|
||||
steps=steps,
|
||||
prev_generated=prev_generated if prev_generated is not None else [],
|
||||
finished=True,
|
||||
)
|
||||
prev_generated = audio_chunk
|
||||
wav_chunks.append(audio_chunk.detach().cpu().numpy())
|
||||
prev_id = i
|
||||
|
||||
if len(wav_chunks) == 0:
|
||||
return torch.zeros(0, device=token2wav_dev)
|
||||
|
||||
waveform = np.concatenate(wav_chunks)
|
||||
return torch.as_tensor(waveform, device=token2wav_dev)
|
||||
|
||||
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
|
||||
"""Load weights for all components of the omni model."""
|
||||
loaded_weights = set()
|
||||
thinker_weights = []
|
||||
talker_weights = []
|
||||
token2wav_weights = []
|
||||
for k, v in weights:
|
||||
if k.startswith('thinker.'):
|
||||
thinker_weights.append((k, v))
|
||||
elif k.startswith('talker.'):
|
||||
talker_weights.append((k, v))
|
||||
elif k.startswith('token2wav.'):
|
||||
token2wav_weights.append((k, v))
|
||||
else:
|
||||
raise ValueError(f"Unknown weight prefix: {k}")
|
||||
|
||||
# Load thinker weights
|
||||
if self.thinker:
|
||||
if thinker_weights:
|
||||
thinker_loaded = self.thinker.load_weights(thinker_weights)
|
||||
else:
|
||||
thinker_loaded = set([k for k,v in thinker_weights])
|
||||
thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, 'thinker')
|
||||
loaded_weights.update(thinker_loaded)
|
||||
|
||||
|
||||
# Load talker weights
|
||||
if talker_weights and self.talker is not None:
|
||||
# Map talker weights to appropriate components
|
||||
if self.thinker is None:
|
||||
thinker_embedding_weights = [w for n, w in thinker_weights if n == "thinker.model.embed_tokens.weight"]
|
||||
if thinker_embedding_weights:
|
||||
self.thinker_embedding = nn.Embedding(thinker_embedding_weights[0].shape[0], thinker_embedding_weights[0].shape[1])
|
||||
self.thinker_embedding.weight = nn.Parameter(thinker_embedding_weights[0].to(self._module_device(self.talker)))
|
||||
talker_loaded = self.talker.load_weights(talker_weights)
|
||||
talker_loaded = add_prefix_to_loaded_weights(talker_loaded, 'talker')
|
||||
loaded_weights.update(talker_loaded)
|
||||
loaded_weights.update(self._init_special_tokens_embeddings())
|
||||
|
||||
|
||||
# Load token2wav weights (if any)
|
||||
if token2wav_weights and self.token2wav is not None:
|
||||
self._init_token2wav_model()
|
||||
t2w_loaded = self.token2wav.load_weights(token2wav_weights, os.path.join(self.vllm_config.model_config.model, "spk_dict.pt"))
|
||||
t2w_loaded = add_prefix_to_loaded_weights(t2w_loaded, 'token2wav')
|
||||
loaded_weights.update(t2w_loaded)
|
||||
|
||||
return loaded_weights
|
||||
@@ -0,0 +1,239 @@
|
||||
from functools import cached_property
|
||||
from typing import Iterable, List, Optional, Set, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniTalkerConfig
|
||||
from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
|
||||
|
||||
from vllm.attention import AttentionMetadata
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.forward_context import ForwardContext, get_forward_context
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear
|
||||
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
|
||||
from vllm.model_executor.sampling_metadata import SamplingMetadata
|
||||
from vllm.model_executor.models.qwen2_5_omni_thinker import (
|
||||
Qwen2_5OmniConditionalGenerationMixin,
|
||||
Qwen2_5OmniThinkerMultiModalProcessor,
|
||||
Qwen2_5OmniThinkerProcessingInfo,
|
||||
Qwen2_5OmniThinkerDummyInputsBuilder)
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.logger import init_logger
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP, MultiModalEmbeddings
|
||||
from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
|
||||
init_vllm_registered_model,
|
||||
maybe_prefix,
|
||||
merge_multimodal_embeddings)
|
||||
from vllm.model_executor.models.qwen2_5_vl import Qwen2_5_VisionTransformer
|
||||
|
||||
|
||||
@MULTIMODAL_REGISTRY.register_processor(
|
||||
Qwen2_5OmniThinkerMultiModalProcessor,
|
||||
info=Qwen2_5OmniThinkerProcessingInfo,
|
||||
dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
|
||||
)
|
||||
class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
|
||||
SupportsPP,
|
||||
Qwen2_5OmniConditionalGenerationMixin):
|
||||
logger = init_logger(__name__)
|
||||
# Align to thinker-style static mapper for clarity
|
||||
hf_to_vllm_mapper = WeightsMapper(
|
||||
orig_to_new_prefix={
|
||||
# text LM head/body in talker
|
||||
"talker.codec_head.": "language_model.lm_head.",
|
||||
"talker.model.": "language_model.model.",
|
||||
# projection weights
|
||||
"talker.thinker_to_talker_proj.": "thinker_to_talker_proj.",
|
||||
# fallback root
|
||||
"talker.": "",
|
||||
})
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
config: Qwen2_5OmniTalkerConfig = vllm_config.model_config.hf_config
|
||||
quant_config = vllm_config.quant_config
|
||||
|
||||
self.prefix = prefix
|
||||
self.quant_config = quant_config
|
||||
|
||||
if hasattr(config, "talker_config"):
|
||||
self.config = config.talker_config
|
||||
vllm_config.model_config.hf_text_config = \
|
||||
vllm_config.model_config.hf_config.talker_config
|
||||
else:
|
||||
self.config = config
|
||||
|
||||
self.thinker_to_talker_proj = ColumnParallelLinear(
|
||||
self.config.embedding_size,
|
||||
self.config.hidden_size,
|
||||
bias=True,
|
||||
gather_output=True,
|
||||
skip_bias_add=False,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
self.language_model = init_vllm_registered_model(
|
||||
vllm_config=vllm_config,
|
||||
prefix=maybe_prefix(prefix, "language_model"),
|
||||
hf_config=getattr(self.config, 'text_config', self.config),
|
||||
architectures=["Qwen2ForCausalLM_old"],
|
||||
)
|
||||
self.make_empty_intermediate_tensors = (
|
||||
self.language_model.make_empty_intermediate_tensors)
|
||||
|
||||
def init_multi_modal(self, thinker_config):
|
||||
self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
|
||||
self.visual = Qwen2_5_VisionTransformer(
|
||||
vision_config=thinker_config.vision_config,
|
||||
norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
|
||||
quant_config=self.quant_config,
|
||||
prefix=maybe_prefix(self.prefix, "visual"),
|
||||
)
|
||||
|
||||
def get_language_model(self) -> torch.nn.Module:
|
||||
return self.language_model
|
||||
|
||||
@cached_property
|
||||
def sampler(self):
|
||||
if hasattr(self.language_model, "sampler"):
|
||||
return self.language_model.sampler
|
||||
|
||||
return get_sampler()
|
||||
|
||||
def get_input_embeddings(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
|
||||
) -> torch.Tensor:
|
||||
inputs_embeds = self.language_model.get_input_embeddings(input_ids)
|
||||
if multimodal_embeddings is not None \
|
||||
and len(multimodal_embeddings) != 0:
|
||||
|
||||
# TODO (ywang96): support overlapping modalitiy embeddings so that
|
||||
# `use_audio_in_video` will work on V1.
|
||||
inputs_embeds = merge_multimodal_embeddings(
|
||||
input_ids, inputs_embeds, multimodal_embeddings, [
|
||||
self.config.image_token_index,
|
||||
self.config.video_token_index,
|
||||
self.config.audio_token_index
|
||||
])
|
||||
return inputs_embeds
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor = None,
|
||||
positions: torch.Tensor = None,
|
||||
intermediate_tensors: Optional[IntermediateTensors] = None,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
**kwargs: object,
|
||||
) -> Union[torch.Tensor, IntermediateTensors]:
|
||||
assert input_ids is not None or inputs_embeds is not None, "input_ids or inputs_embeds must be provided"
|
||||
forward_context: ForwardContext = get_forward_context()
|
||||
attn_metadata: AttentionMetadata = forward_context.attn_metadata
|
||||
|
||||
if intermediate_tensors is not None:
|
||||
inputs_embeds = None
|
||||
elif inputs_embeds is None:
|
||||
# for profile_run:
|
||||
inputs_embeds = self.get_input_embeddings(input_ids)
|
||||
|
||||
input_ids = None
|
||||
|
||||
# projection
|
||||
inputs_embeds, _ = self.thinker_to_talker_proj(inputs_embeds)
|
||||
|
||||
hidden_states = self.language_model.model(input_ids,
|
||||
positions,
|
||||
intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds)
|
||||
return hidden_states
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> Optional[torch.Tensor]:
|
||||
return self.language_model.compute_logits(hidden_states,
|
||||
sampling_metadata)
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> Optional[SamplerOutput]:
|
||||
return self.language_model.sample(logits, sampling_metadata)
|
||||
|
||||
def load_weights(self, weights: Iterable[Tuple[str,
|
||||
torch.Tensor]]) -> Set[str]:
|
||||
loader = AutoWeightsLoader(
|
||||
self,
|
||||
skip_prefixes=['thinker.', 'token2wav.'],
|
||||
)
|
||||
loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
|
||||
# Log load summary
|
||||
try:
|
||||
total_bytes = 0
|
||||
for name, param in self.named_parameters():
|
||||
if param is not None and param.data is not None:
|
||||
total_bytes += param.data.numel() * param.data.element_size()
|
||||
device = next(self.parameters()).device
|
||||
self.logger.info(
|
||||
"[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s",
|
||||
self.__class__.__name__, True, total_bytes / (1024**2), str(device))
|
||||
except Exception:
|
||||
pass
|
||||
multi_model_weights = set()
|
||||
for name, param in self.visual.named_parameters():
|
||||
multi_model_weights.add("visual."+name)
|
||||
for name, param in self.audio_tower.named_parameters():
|
||||
multi_model_weights.add("audio_tower."+name)
|
||||
loaded.update(multi_model_weights)
|
||||
return loaded
|
||||
|
||||
def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
|
||||
mm_input_by_modality = {}
|
||||
|
||||
# Preserve the order of modalities if there are multiple of them
|
||||
# from the order of kwargs.
|
||||
for input_key in kwargs:
|
||||
if input_key in ("pixel_values", "image_embeds"
|
||||
) and "image" not in mm_input_by_modality:
|
||||
mm_input_by_modality[
|
||||
"image"] = self._parse_and_validate_image_input(**kwargs)
|
||||
if input_key in ("pixel_values_videos", "video_embeds"
|
||||
) and "video" not in mm_input_by_modality:
|
||||
mm_input_by_modality[
|
||||
"video"] = self._parse_and_validate_video_input(**kwargs)
|
||||
if input_key in ("input_audio_features"
|
||||
) and "audio" not in mm_input_by_modality:
|
||||
mm_input_by_modality[
|
||||
"audio"] = self._parse_and_validate_audio_input(**kwargs)
|
||||
return mm_input_by_modality
|
||||
|
||||
def get_multimodal_embeddings(self,
|
||||
**kwargs: object) -> MultiModalEmbeddings:
|
||||
|
||||
mm_input_by_modality = self._parse_and_validate_multimodal_inputs(
|
||||
**kwargs)
|
||||
if not mm_input_by_modality:
|
||||
return []
|
||||
|
||||
# The result multimodal_embeddings is tuple of tensors, with each
|
||||
# tensor correspoending to a multimodal data item (image or video).
|
||||
multimodal_embeddings: tuple[torch.Tensor, ...] = ()
|
||||
|
||||
# NOTE: It is important to iterate over the keys in this dictionary
|
||||
# to preserve the order of the modalities.
|
||||
for modality in mm_input_by_modality:
|
||||
multimodal_input = mm_input_by_modality[modality]
|
||||
if modality == "image":
|
||||
vision_embeddings = self._process_image_input(multimodal_input)
|
||||
multimodal_embeddings += vision_embeddings
|
||||
if modality == "video":
|
||||
video_embeddings = self._process_video_input(multimodal_input)
|
||||
multimodal_embeddings += video_embeddings
|
||||
if modality == "audio":
|
||||
audio_embeddings = self._process_audio_input(multimodal_input)
|
||||
multimodal_embeddings += audio_embeddings
|
||||
return multimodal_embeddings
|
||||
@@ -0,0 +1,889 @@
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from copy import copy
|
||||
from functools import partial
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers.feature_extraction_utils import BatchFeature
|
||||
from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
|
||||
Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig)
|
||||
from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
|
||||
Qwen2_5OmniAudioEncoder)
|
||||
from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import (
|
||||
Qwen2_5OmniProcessor)
|
||||
from transformers.models.whisper import WhisperFeatureExtractor
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm.model_executor.models.qwen2_5_vl import (
|
||||
Qwen2_5_VisionTransformer, Qwen2_5_VLImageEmbeddingInputs,
|
||||
Qwen2_5_VLImageInputs, Qwen2_5_VLImagePixelInputs,
|
||||
Qwen2_5_VLProcessingInfo, Qwen2_5_VLVideoEmbeddingInputs,
|
||||
Qwen2_5_VLVideoInputs, Qwen2_5_VLVideoPixelInputs)
|
||||
from vllm.model_executor.models.qwen2_audio import (
|
||||
Qwen2AudioInputs, Qwen2AudioProcessingInfo,
|
||||
_get_feat_extract_output_lengths)
|
||||
from vllm.model_executor.models.qwen2_vl import Qwen2VLMultiModalDataParser
|
||||
from vllm.model_executor.sampling_metadata import SamplingMetadata
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.multimodal.inputs import (ImageItem, ModalityData,
|
||||
MultiModalDataDict, MultiModalFieldConfig,
|
||||
MultiModalKwargs, NestedTensors)
|
||||
from vllm.multimodal.parse import (AudioProcessorItems, DictEmbeddingItems,
|
||||
ModalityDataItems, MultiModalDataItems,
|
||||
MultiModalDataParser)
|
||||
from vllm.multimodal.processing import (BaseMultiModalProcessor,
|
||||
PlaceholderFeaturesInfo,
|
||||
PromptReplacement, PromptUpdate)
|
||||
from vllm.multimodal.profiling import BaseDummyInputsBuilder
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.transformers_utils.tokenizer import decode_tokens, encode_tokens
|
||||
|
||||
from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
|
||||
from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
|
||||
init_vllm_registered_model, maybe_prefix,
|
||||
merge_multimodal_embeddings)
|
||||
|
||||
from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
|
||||
|
||||
try:
|
||||
import flash_attn
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
flash_attn = None
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
|
||||
audio_feature_lengths = hf_inputs.get("audio_feature_lengths",
|
||||
torch.empty((0, )))
|
||||
|
||||
image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
|
||||
image_grid_sizes = image_grid_thw.prod(-1)
|
||||
|
||||
video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3)))
|
||||
video_grid_sizes = video_grid_thw.prod(-1)
|
||||
|
||||
return dict(
|
||||
input_audio_features=MultiModalFieldConfig.flat_from_sizes(
|
||||
"audio", audio_feature_lengths, dim=1),
|
||||
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
|
||||
audio_feature_lengths=MultiModalFieldConfig.batched("audio"),
|
||||
pixel_values=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", image_grid_sizes),
|
||||
image_embeds=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", image_grid_sizes),
|
||||
image_grid_thw=MultiModalFieldConfig.batched("image"),
|
||||
pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", video_grid_sizes),
|
||||
video_embeds=MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", video_grid_sizes),
|
||||
video_grid_thw=MultiModalFieldConfig.batched("video"),
|
||||
second_per_grid_ts=MultiModalFieldConfig.batched("video"),
|
||||
)
|
||||
|
||||
|
||||
class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
|
||||
|
||||
def _parse_audio_data(
|
||||
self,
|
||||
data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],
|
||||
) -> ModalityDataItems[Any, Any]:
|
||||
if isinstance(data, dict):
|
||||
return DictEmbeddingItems(
|
||||
data,
|
||||
modality="audio",
|
||||
required_fields={
|
||||
"input_audio_features", "audio_feature_lengths"
|
||||
},
|
||||
fields_factory=_qwen2_5_omni_thinker_field_config,
|
||||
)
|
||||
|
||||
return super()._parse_audio_data(data)
|
||||
|
||||
|
||||
class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo,
|
||||
Qwen2_5_VLProcessingInfo):
|
||||
|
||||
def get_hf_config(self):
|
||||
return self.ctx.get_hf_config(Qwen2_5OmniConfig).thinker_config
|
||||
|
||||
def get_hf_processor(self, **kwargs: object) -> Qwen2_5OmniProcessor:
|
||||
return self.ctx.get_hf_processor(
|
||||
Qwen2_5OmniProcessor,
|
||||
use_fast=kwargs.pop("use_fast", True),
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def get_feature_extractor(self, **kwargs: object):
|
||||
hf_processor = self.get_hf_processor(**kwargs)
|
||||
feature_extractor = hf_processor.feature_extractor # type: ignore
|
||||
assert isinstance(feature_extractor, WhisperFeatureExtractor)
|
||||
return feature_extractor
|
||||
|
||||
def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
|
||||
return {"audio": None, "image": None, "video": None}
|
||||
|
||||
|
||||
class Qwen2_5OmniThinkerDummyInputsBuilder(
|
||||
BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]):
|
||||
|
||||
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
|
||||
num_audios = mm_counts.get("audio", 0)
|
||||
num_images = mm_counts.get("image", 0)
|
||||
num_videos = mm_counts.get("video", 0)
|
||||
|
||||
hf_processor = self.info.get_hf_processor()
|
||||
|
||||
audio_token: str = hf_processor.audio_token
|
||||
image_token: str = hf_processor.image_token
|
||||
video_token: str = hf_processor.video_token
|
||||
|
||||
return (audio_token * num_audios + image_token * num_images +
|
||||
video_token * num_videos)
|
||||
|
||||
def get_dummy_mm_data(
|
||||
self,
|
||||
seq_len: int,
|
||||
mm_counts: Mapping[str, int],
|
||||
) -> MultiModalDataDict:
|
||||
num_audios = mm_counts.get("audio", 0)
|
||||
num_images = mm_counts.get("image", 0)
|
||||
num_videos = mm_counts.get("video", 0)
|
||||
|
||||
feature_extractor = self.info.get_feature_extractor()
|
||||
|
||||
target_audio_length = min(
|
||||
feature_extractor.chunk_length,
|
||||
30,
|
||||
) * feature_extractor.sampling_rate
|
||||
target_width, target_height = \
|
||||
self.info.get_image_size_with_most_features()
|
||||
target_num_frames = \
|
||||
self.info.get_num_frames_with_most_features(seq_len, mm_counts)
|
||||
|
||||
mm_data = {
|
||||
"audio":
|
||||
self._get_dummy_audios(length=target_audio_length,
|
||||
num_audios=num_audios),
|
||||
"image":
|
||||
self._get_dummy_images(width=target_width,
|
||||
height=target_height,
|
||||
num_images=num_images),
|
||||
"video":
|
||||
self._get_dummy_videos(width=target_width,
|
||||
height=target_height,
|
||||
num_frames=target_num_frames,
|
||||
num_videos=num_videos),
|
||||
}
|
||||
|
||||
return mm_data
|
||||
|
||||
|
||||
class Qwen2_5OmniThinkerMultiModalProcessor(
|
||||
BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]):
|
||||
|
||||
def _get_data_parser(self) -> MultiModalDataParser:
|
||||
feature_extractor = self.info.get_feature_extractor()
|
||||
return Qwen2_5OmniThinkerMultiModalDataParser(
|
||||
target_sr=feature_extractor.sampling_rate)
|
||||
|
||||
def _call_hf_processor(
|
||||
self,
|
||||
prompt: str,
|
||||
mm_data: Mapping[str, object],
|
||||
mm_kwargs: Mapping[str, object],
|
||||
tok_kwargs: Mapping[str, object],
|
||||
) -> BatchFeature:
|
||||
mm_data = dict(mm_data)
|
||||
audios = mm_data.pop("audios", [])
|
||||
|
||||
# NOTE: WhisperFeatureExtractor cannot handle empty list of audios
|
||||
if audios:
|
||||
# NOTE: Qwen2.5-Omni processor accept "audio"
|
||||
mm_data["audio"] = audios
|
||||
mm_kwargs = dict(**mm_kwargs, )
|
||||
|
||||
hf_inputs = super()._call_hf_processor(
|
||||
prompt=prompt,
|
||||
mm_data=mm_data,
|
||||
mm_kwargs=mm_kwargs,
|
||||
tok_kwargs=tok_kwargs,
|
||||
)
|
||||
|
||||
input_features = hf_inputs.pop('input_features', None)
|
||||
feature_attention_mask = hf_inputs.get('feature_attention_mask', None)
|
||||
if ('input_audio_features' not in hf_inputs
|
||||
and input_features is not None):
|
||||
if feature_attention_mask is not None:
|
||||
input_features = input_features.permute(
|
||||
0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
|
||||
hf_inputs['input_audio_features'] = input_features
|
||||
if ('audio_feature_lengths' not in hf_inputs
|
||||
and feature_attention_mask is not None):
|
||||
hf_inputs['audio_feature_lengths'] = feature_attention_mask.sum(-1)
|
||||
return hf_inputs
|
||||
|
||||
def _get_mm_fields_config(
|
||||
self,
|
||||
hf_inputs: BatchFeature,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
) -> Mapping[str, MultiModalFieldConfig]:
|
||||
return _qwen2_5_omni_thinker_field_config(hf_inputs)
|
||||
|
||||
def _maybe_apply_prompt_updates(
|
||||
self,
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
prompt_ids: list[int],
|
||||
mm_kwargs: MultiModalKwargs,
|
||||
is_update_applied: bool,
|
||||
) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]:
|
||||
"""
|
||||
Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.
|
||||
"""
|
||||
unbound_prompt_updates = self._get_prompt_updates(
|
||||
mm_items,
|
||||
hf_processor_mm_kwargs,
|
||||
mm_kwargs,
|
||||
)
|
||||
mm_prompt_updates = self._bind_and_group_updates(
|
||||
unbound_prompt_updates)
|
||||
|
||||
mm_item_counts = mm_items.get_all_counts()
|
||||
self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
|
||||
|
||||
use_audio_in_video = hf_processor_mm_kwargs.get(
|
||||
"use_audio_in_video", False)
|
||||
|
||||
if is_update_applied:
|
||||
mm_placeholders = self._find_mm_placeholders(
|
||||
mm_prompt_updates,
|
||||
prompt_ids,
|
||||
mm_item_counts,
|
||||
)
|
||||
self._validate_mm_placeholders(
|
||||
mm_placeholders,
|
||||
mm_item_counts,
|
||||
use_audio_in_video=use_audio_in_video)
|
||||
|
||||
tokenizer = self.info.get_tokenizer()
|
||||
prompt = decode_tokens(tokenizer, prompt_ids)
|
||||
else:
|
||||
(
|
||||
prompt_ids,
|
||||
prompt,
|
||||
mm_placeholders,
|
||||
) = self._apply_prompt_updates(
|
||||
prompt_ids,
|
||||
mm_prompt_updates,
|
||||
mm_item_counts,
|
||||
)
|
||||
self._validate_mm_placeholders(
|
||||
mm_placeholders,
|
||||
mm_item_counts,
|
||||
use_audio_in_video=use_audio_in_video)
|
||||
|
||||
tokenizer = self.info.get_tokenizer()
|
||||
prompt = decode_tokens(tokenizer, prompt_ids)
|
||||
|
||||
if use_audio_in_video:
|
||||
mm_kwargs["use_audio_in_video"] = True
|
||||
|
||||
return prompt_ids, prompt, mm_placeholders
|
||||
|
||||
def _get_prompt_updates(
|
||||
self,
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, Any],
|
||||
out_mm_kwargs: MultiModalKwargs,
|
||||
) -> Sequence[PromptUpdate]:
|
||||
processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
|
||||
tokenizer = self.info.get_tokenizer()
|
||||
image_processor = self.info.get_image_processor(
|
||||
**hf_processor_mm_kwargs)
|
||||
vocab = tokenizer.get_vocab()
|
||||
|
||||
audio_token = processor.audio_token
|
||||
image_token = processor.image_token
|
||||
video_token = processor.video_token
|
||||
audio_token_id = vocab[audio_token]
|
||||
image_token_id = vocab[image_token]
|
||||
video_token_id = vocab[video_token]
|
||||
|
||||
audio_feature_lengths = out_mm_kwargs.get("audio_feature_lengths")
|
||||
feature_attention_mask = out_mm_kwargs.get("feature_attention_mask")
|
||||
if audio_feature_lengths is None and feature_attention_mask is None:
|
||||
audio_output_lengths = []
|
||||
elif audio_feature_lengths is not None:
|
||||
_, audio_output_lens = _get_feat_extract_output_lengths(
|
||||
audio_feature_lengths)
|
||||
audio_output_lengths = audio_output_lens.tolist()
|
||||
elif feature_attention_mask is not None:
|
||||
assert isinstance(feature_attention_mask, torch.Tensor)
|
||||
_, audio_output_lens = _get_feat_extract_output_lengths(
|
||||
feature_attention_mask.sum(-1))
|
||||
audio_output_lengths = audio_output_lens.tolist()
|
||||
|
||||
# number of audios read from video.
|
||||
audio_in_video_item_idx = 0
|
||||
|
||||
def get_replacement_qwen2_audio(item_idx: int):
|
||||
item_idx += audio_in_video_item_idx
|
||||
|
||||
num_features = audio_output_lengths[item_idx]
|
||||
if num_features == 0:
|
||||
audios = mm_items.get_items("audio", AudioProcessorItems)
|
||||
audio = audios.get(item_idx)
|
||||
raise ValueError(
|
||||
f"The audio {audio} (len={len(audio)}) is too short "
|
||||
"to be represented inside the model")
|
||||
|
||||
return [audio_token_id] * num_features
|
||||
|
||||
def get_replacement_qwen2_vision(item_idx: int, modality: str):
|
||||
grid_thw = out_mm_kwargs[f"{modality}_grid_thw"][item_idx]
|
||||
assert isinstance(grid_thw, torch.Tensor)
|
||||
merge_length = image_processor.merge_size**2
|
||||
|
||||
token_id = image_token_id if modality == "image" else video_token_id
|
||||
return [token_id] * (int(grid_thw.prod()) // merge_length)
|
||||
|
||||
use_audio_in_video = hf_processor_mm_kwargs.get(
|
||||
"use_audio_in_video", False)
|
||||
thinker_config = self.info.get_hf_config()
|
||||
|
||||
def get_replacement_qwen2_use_audio_in_video(item_idx: int):
|
||||
nonlocal audio_in_video_item_idx
|
||||
|
||||
audio_num_features = audio_output_lengths[audio_in_video_item_idx +
|
||||
item_idx]
|
||||
video_grid_thw = out_mm_kwargs["video_grid_thw"][item_idx]
|
||||
|
||||
audio_in_video_item_idx += 1
|
||||
|
||||
second_per_grid_ts = hf_processor_mm_kwargs.get(
|
||||
"second_per_grid_ts", None)
|
||||
if second_per_grid_ts:
|
||||
video_second_per_grid_t = second_per_grid_ts[item_idx]
|
||||
else:
|
||||
video_second_per_grid_t = 1.0
|
||||
|
||||
return MRotaryEmbedding.omni_get_updates_use_audio_in_video(
|
||||
thinker_config=thinker_config,
|
||||
audio_len=audio_num_features,
|
||||
video_grid_thw=video_grid_thw,
|
||||
video_second_per_grid_t=video_second_per_grid_t,
|
||||
)
|
||||
|
||||
video_replacement_fn = (
|
||||
get_replacement_qwen2_use_audio_in_video if use_audio_in_video else
|
||||
partial(get_replacement_qwen2_vision, modality="video"))
|
||||
|
||||
return [
|
||||
PromptReplacement(
|
||||
modality="audio",
|
||||
target=audio_token,
|
||||
replacement=get_replacement_qwen2_audio,
|
||||
),
|
||||
PromptReplacement(
|
||||
modality="image",
|
||||
target=image_token,
|
||||
replacement=partial(get_replacement_qwen2_vision,
|
||||
modality="image"),
|
||||
),
|
||||
PromptReplacement(
|
||||
modality="video",
|
||||
target=video_token,
|
||||
replacement=video_replacement_fn,
|
||||
),
|
||||
]
|
||||
|
||||
def _apply_hf_processor_main(
|
||||
self,
|
||||
prompt: Union[str, list[int]],
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
tokenization_kwargs: Mapping[str, object],
|
||||
*,
|
||||
enable_hf_prompt_update: bool,
|
||||
) -> tuple[list[int], BatchFeature, bool]:
|
||||
"""
|
||||
Qwen2.5-Omni reimplements this function to handle text only.
|
||||
"""
|
||||
if isinstance(prompt, str):
|
||||
if enable_hf_prompt_update:
|
||||
return self._apply_hf_processor_text_mm(
|
||||
prompt_text=prompt,
|
||||
mm_items=mm_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
)
|
||||
tokenizer = self.info.get_tokenizer()
|
||||
prompt_ids = encode_tokens(tokenizer, prompt)
|
||||
else:
|
||||
prompt_ids = self._apply_hf_processor_tokens_only(prompt)
|
||||
|
||||
mm_processed_data = self._apply_hf_processor_mm_only(
|
||||
mm_items=mm_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
)
|
||||
|
||||
return prompt_ids, mm_processed_data, False
|
||||
|
||||
def _apply_hf_processor_mm_only(
|
||||
self,
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
tokenization_kwargs: Mapping[str, object],
|
||||
) -> BatchFeature:
|
||||
"""
|
||||
Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.
|
||||
"""
|
||||
mm_counts = mm_items.get_all_counts()
|
||||
|
||||
use_audio_in_video = hf_processor_mm_kwargs.get(
|
||||
"use_audio_in_video", False)
|
||||
if use_audio_in_video and "video" in mm_counts:
|
||||
assert "audio" in mm_counts
|
||||
mm_counts["audio"] -= mm_counts["video"]
|
||||
|
||||
_, mm_processed_data, _ = self._apply_hf_processor_text_mm(
|
||||
prompt_text=self.dummy_inputs.get_dummy_text(mm_counts),
|
||||
mm_items=mm_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
)
|
||||
|
||||
return mm_processed_data
|
||||
|
||||
def _validate_mm_placeholders(
|
||||
self,
|
||||
mm_placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
|
||||
mm_item_counts: Mapping[str, int],
|
||||
use_audio_in_video: bool = False,
|
||||
) -> None:
|
||||
if use_audio_in_video:
|
||||
mm_item_counts = copy(mm_item_counts)
|
||||
if "video" in mm_item_counts:
|
||||
assert "audio" in mm_item_counts
|
||||
mm_item_counts["audio"] -= mm_item_counts["video"]
|
||||
super()._validate_mm_placeholders(mm_placeholders, mm_item_counts)
|
||||
|
||||
|
||||
class Qwen2_5OmniConditionalGenerationMixin:
|
||||
|
||||
def _validate_and_reshape_mm_tensor(self,
|
||||
mm_input: object,
|
||||
name: str,
|
||||
dim: int = 0) -> torch.Tensor:
|
||||
if not isinstance(mm_input, (torch.Tensor, list)):
|
||||
raise ValueError(f"Incorrect type of {name}. "
|
||||
f"Got type: {type(mm_input)}")
|
||||
if isinstance(mm_input, torch.Tensor):
|
||||
return torch.concat(list(mm_input), dim=dim)
|
||||
else:
|
||||
return torch.concat(mm_input, dim=dim)
|
||||
|
||||
def _parse_and_validate_audio_input(
|
||||
self, **kwargs: object) -> Optional[Qwen2AudioInputs]:
|
||||
input_audio_features = kwargs.pop('input_audio_features', None)
|
||||
audio_feature_lengths = kwargs.pop('audio_feature_lengths', None)
|
||||
feature_attention_mask = kwargs.pop('feature_attention_mask', None)
|
||||
if input_audio_features is None:
|
||||
return None
|
||||
input_audio_features = self._validate_and_reshape_mm_tensor(
|
||||
input_audio_features, 'input_audio_features', dim=1)
|
||||
if feature_attention_mask is not None:
|
||||
feature_attention_mask = self._validate_and_reshape_mm_tensor(
|
||||
feature_attention_mask, 'feature_attention_mask')
|
||||
if not isinstance(input_audio_features, (torch.Tensor, list)):
|
||||
raise ValueError("Incorrect type of audio input features. "
|
||||
f"Got type: {type(input_audio_features)}")
|
||||
return Qwen2AudioInputs(input_features=input_audio_features,
|
||||
audio_feature_lengths=audio_feature_lengths,
|
||||
feature_attention_mask=feature_attention_mask)
|
||||
|
||||
def _parse_and_validate_image_input(
|
||||
self,
|
||||
**kwargs: dict[str, Any],
|
||||
) -> Optional[Qwen2_5_VLImageInputs]:
|
||||
pixel_values = kwargs.pop("pixel_values", None)
|
||||
image_embeds = kwargs.pop("image_embeds", None)
|
||||
image_grid_thw = kwargs.pop("image_grid_thw", None)
|
||||
|
||||
if pixel_values is None and image_embeds is None:
|
||||
return None
|
||||
|
||||
if pixel_values is not None:
|
||||
pixel_values = self._validate_and_reshape_mm_tensor(
|
||||
pixel_values, "image pixel values")
|
||||
image_grid_thw = self._validate_and_reshape_mm_tensor(
|
||||
image_grid_thw, "image grid_thw")
|
||||
|
||||
if not isinstance(pixel_values, (torch.Tensor, list)):
|
||||
raise ValueError("Incorrect type of image pixel values. "
|
||||
f"Got type: {type(pixel_values)}")
|
||||
|
||||
return Qwen2_5_VLImagePixelInputs(type="pixel_values",
|
||||
pixel_values=pixel_values,
|
||||
image_grid_thw=image_grid_thw)
|
||||
|
||||
if image_embeds is not None:
|
||||
image_embeds = self._validate_and_reshape_mm_tensor(
|
||||
image_embeds, "image embeds")
|
||||
image_grid_thw = self._validate_and_reshape_mm_tensor(
|
||||
image_grid_thw, "image grid_thw")
|
||||
|
||||
if not isinstance(image_embeds, torch.Tensor):
|
||||
raise ValueError("Incorrect type of image embeddings. "
|
||||
f"Got type: {type(image_embeds)}")
|
||||
return Qwen2_5_VLImageEmbeddingInputs(
|
||||
type="image_embeds",
|
||||
image_embeds=image_embeds,
|
||||
image_grid_thw=image_grid_thw)
|
||||
|
||||
def _parse_and_validate_video_input(
|
||||
self,
|
||||
**kwargs: dict[str, Any],
|
||||
) -> Optional[Qwen2_5_VLVideoInputs]:
|
||||
pixel_values_videos = kwargs.pop("pixel_values_videos", None)
|
||||
video_embeds = kwargs.pop("video_embeds", None)
|
||||
video_grid_thw = kwargs.pop("video_grid_thw", None)
|
||||
|
||||
if pixel_values_videos is None and video_embeds is None:
|
||||
return None
|
||||
|
||||
if pixel_values_videos is not None:
|
||||
pixel_values_videos = self._validate_and_reshape_mm_tensor(
|
||||
pixel_values_videos, "video pixel values")
|
||||
video_grid_thw = self._validate_and_reshape_mm_tensor(
|
||||
video_grid_thw, "video grid_thw")
|
||||
|
||||
return Qwen2_5_VLVideoPixelInputs(
|
||||
type="pixel_values_videos",
|
||||
pixel_values_videos=pixel_values_videos,
|
||||
video_grid_thw=video_grid_thw,
|
||||
)
|
||||
|
||||
if video_embeds is not None:
|
||||
video_embeds = self._validate_and_reshape_mm_tensor(
|
||||
video_embeds, "video embeds")
|
||||
video_grid_thw = self._validate_and_reshape_mm_tensor(
|
||||
video_grid_thw, "video grid_thw")
|
||||
|
||||
if not isinstance(video_embeds, torch.Tensor):
|
||||
raise ValueError("Incorrect type of video embeddings. "
|
||||
f"Got type: {type(video_embeds)}")
|
||||
return Qwen2_5_VLVideoEmbeddingInputs(
|
||||
type="video_embeds",
|
||||
video_embeds=video_embeds,
|
||||
video_grid_thw=video_grid_thw)
|
||||
|
||||
def _process_audio_input(
|
||||
self,
|
||||
audio_input: Qwen2AudioInputs,
|
||||
audio_hashes: list[str] = None,
|
||||
cached_audio_features: torch.Tensor = None,
|
||||
) -> torch.Tensor:
|
||||
|
||||
input_features = audio_input["input_features"]
|
||||
audio_feature_lengths = audio_input["audio_feature_lengths"]
|
||||
if input_features.ndim == 3:
|
||||
assert input_features.shape[0] == 1
|
||||
input_features = input_features.squeeze(0)
|
||||
if audio_feature_lengths.ndim == 2:
|
||||
assert audio_feature_lengths.shape[
|
||||
0] == 1 or audio_feature_lengths.shape[1] == 1
|
||||
if audio_feature_lengths.shape[0] == 1:
|
||||
audio_feature_lengths = audio_feature_lengths.squeeze(0)
|
||||
else:
|
||||
audio_feature_lengths = audio_feature_lengths.squeeze(1)
|
||||
|
||||
audio_feat_lengths, audio_output_lengths = (
|
||||
self.audio_tower._get_feat_extract_output_lengths(
|
||||
audio_feature_lengths))
|
||||
|
||||
audio_outputs = self.audio_tower(
|
||||
input_features.to(self.audio_tower.dtype),
|
||||
feature_lens=audio_feature_lengths,
|
||||
aftercnn_lens=audio_feat_lengths,
|
||||
)
|
||||
audio_features = audio_outputs.last_hidden_state
|
||||
return audio_features.split(audio_output_lengths.tolist())
|
||||
|
||||
def _process_image_input(
|
||||
self,
|
||||
image_input: Qwen2_5_VLImageInputs) -> tuple[torch.Tensor, ...]:
|
||||
if image_input["type"] == "image_embeds":
|
||||
return image_input["image_embeds"].type(self.visual.dtype)
|
||||
|
||||
grid_thw = image_input["image_grid_thw"]
|
||||
assert grid_thw.ndim == 2
|
||||
|
||||
pixel_values = image_input["pixel_values"].type(self.visual.dtype)
|
||||
image_embeds = self.visual(pixel_values, grid_thw=grid_thw)
|
||||
# Split concatenated embeddings for each image item.
|
||||
merge_size = self.visual.spatial_merge_size
|
||||
sizes = grid_thw.prod(-1) // merge_size // merge_size
|
||||
|
||||
return image_embeds.split(sizes.tolist())
|
||||
|
||||
def _process_video_input(
|
||||
self,
|
||||
video_input: Qwen2_5_VLVideoInputs,
|
||||
video_hashes: list[str] = None,
|
||||
cached_video_embeds: torch.Tensor = None) -> torch.Tensor:
|
||||
if video_input["type"] == "video_embeds":
|
||||
return video_input["video_embeds"].type(self.visual.dtype)
|
||||
|
||||
grid_thw = video_input["video_grid_thw"]
|
||||
assert grid_thw.ndim == 2
|
||||
|
||||
pixel_values_videos = video_input["pixel_values_videos"].type(
|
||||
self.visual.dtype)
|
||||
video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)
|
||||
# Split concatenated embeddings for each video item.
|
||||
merge_size = self.visual.spatial_merge_size
|
||||
sizes = grid_thw.prod(-1) // merge_size // merge_size
|
||||
|
||||
return video_embeds.split(sizes.tolist())
|
||||
|
||||
|
||||
@MULTIMODAL_REGISTRY.register_processor(
|
||||
Qwen2_5OmniThinkerMultiModalProcessor,
|
||||
info=Qwen2_5OmniThinkerProcessingInfo,
|
||||
dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
|
||||
)
|
||||
class Qwen2_5OmniThinkerForConditionalGeneration(
|
||||
nn.Module, SupportsMultiModal, SupportsPP,
|
||||
Qwen2_5OmniConditionalGenerationMixin):
|
||||
hf_to_vllm_mapper = WeightsMapper(
|
||||
orig_to_new_prefix={
|
||||
"thinker.lm_head.": "language_model.lm_head.",
|
||||
"thinker.model.": "language_model.model.",
|
||||
"thinker.": "",
|
||||
})
|
||||
|
||||
@classmethod
|
||||
def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
|
||||
if modality.startswith("image"):
|
||||
return "<|vision_start|><|IMAGE|><|vision_end|>"
|
||||
if modality.startswith("video"):
|
||||
return "<|vision_start|><|VIDEO|><|vision_end|>"
|
||||
if modality.startswith("audio"):
|
||||
return f"Audio {i}: <|audio_bos|><|AUDIO|><|audio_eos|>"
|
||||
|
||||
raise ValueError("Only image, video or audio modality is supported")
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
thinker_config: Qwen2_5OmniThinkerConfig = (
|
||||
vllm_config.model_config.hf_config)
|
||||
quant_config = vllm_config.quant_config
|
||||
multimodal_config = vllm_config.model_config.multimodal_config
|
||||
self.config = thinker_config
|
||||
self.multimodal_config = multimodal_config
|
||||
|
||||
# force "use_flash_attention_2=True" to audio tower to align
|
||||
# the results.
|
||||
if flash_attn is not None:
|
||||
audio_config = thinker_config.audio_config
|
||||
audio_config._attn_implementation_autoset = True
|
||||
audio_config._attn_implementation = "flash_attention_2"
|
||||
else:
|
||||
logger.warning(
|
||||
"flash_attn is not available, the model may not yield the "
|
||||
"exactly same result as the transformers implementation "
|
||||
"in the audio tower part.")
|
||||
|
||||
self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
|
||||
self.visual = Qwen2_5_VisionTransformer(
|
||||
vision_config=thinker_config.vision_config,
|
||||
norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
|
||||
quant_config=quant_config,
|
||||
prefix=maybe_prefix(prefix, "visual"),
|
||||
)
|
||||
self.quant_config = quant_config
|
||||
self.language_model = init_vllm_registered_model(
|
||||
vllm_config=vllm_config,
|
||||
prefix=maybe_prefix(prefix, "language_model"),
|
||||
hf_config=thinker_config.text_config,
|
||||
architectures=["Qwen2ForCausalLM"],
|
||||
)
|
||||
|
||||
self.make_empty_intermediate_tensors = (
|
||||
self.language_model.make_empty_intermediate_tensors)
|
||||
|
||||
def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
|
||||
mm_input_by_modality = {}
|
||||
|
||||
# Preserve the order of modalities if there are multiple of them
|
||||
# from the order of kwargs.
|
||||
for input_key in kwargs:
|
||||
if input_key in ("pixel_values", "image_embeds"
|
||||
) and "image" not in mm_input_by_modality:
|
||||
mm_input_by_modality[
|
||||
"image"] = self._parse_and_validate_image_input(**kwargs)
|
||||
if input_key in ("pixel_values_videos", "video_embeds"
|
||||
) and "video" not in mm_input_by_modality:
|
||||
mm_input_by_modality[
|
||||
"video"] = self._parse_and_validate_video_input(**kwargs)
|
||||
if input_key in ("input_audio_features"
|
||||
) and "audio" not in mm_input_by_modality:
|
||||
mm_input_by_modality[
|
||||
"audio"] = self._parse_and_validate_audio_input(**kwargs)
|
||||
return mm_input_by_modality
|
||||
|
||||
def get_language_model(self) -> torch.nn.Module:
|
||||
return self.language_model
|
||||
|
||||
def get_multimodal_embeddings(self,
|
||||
**kwargs: object) -> MultiModalEmbeddings:
|
||||
|
||||
mm_input_by_modality = self._parse_and_validate_multimodal_inputs(
|
||||
**kwargs)
|
||||
if not mm_input_by_modality:
|
||||
return []
|
||||
|
||||
# The result multimodal_embeddings is tuple of tensors, with each
|
||||
# tensor correspoending to a multimodal data item (image or video).
|
||||
multimodal_embeddings: tuple[torch.Tensor, ...] = ()
|
||||
|
||||
# NOTE: It is important to iterate over the keys in this dictionary
|
||||
# to preserve the order of the modalities.
|
||||
for modality in mm_input_by_modality:
|
||||
multimodal_input = mm_input_by_modality[modality]
|
||||
if modality == "image":
|
||||
vision_embeddings = self._process_image_input(multimodal_input)
|
||||
multimodal_embeddings += vision_embeddings
|
||||
if modality == "video":
|
||||
video_embeddings = self._process_video_input(multimodal_input)
|
||||
multimodal_embeddings += video_embeddings
|
||||
if modality == "audio":
|
||||
audio_embeddings = self._process_audio_input(multimodal_input)
|
||||
multimodal_embeddings += audio_embeddings
|
||||
return multimodal_embeddings
|
||||
|
||||
def get_input_embeddings(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
|
||||
) -> torch.Tensor:
|
||||
inputs_embeds = self.language_model.get_input_embeddings(input_ids)
|
||||
if multimodal_embeddings is not None \
|
||||
and len(multimodal_embeddings) != 0:
|
||||
|
||||
# TODO (ywang96): support overlapping modalitiy embeddings so that
|
||||
# `use_audio_in_video` will work on V1.
|
||||
inputs_embeds = merge_multimodal_embeddings(
|
||||
input_ids, inputs_embeds, multimodal_embeddings, [
|
||||
self.config.image_token_index,
|
||||
self.config.video_token_index,
|
||||
self.config.audio_token_index
|
||||
])
|
||||
return inputs_embeds
|
||||
|
||||
def get_multimodal_embeddings_v0(
|
||||
self, **kwargs: object) -> Optional[NestedTensors]:
|
||||
audio_input = self._parse_and_validate_audio_input(**kwargs)
|
||||
image_input = self._parse_and_validate_image_input(**kwargs)
|
||||
video_input = self._parse_and_validate_video_input(**kwargs)
|
||||
|
||||
if audio_input is None and image_input is None and video_input is None:
|
||||
return None
|
||||
|
||||
multimodal_embeddings: list[tuple[NestedTensors, str]] = []
|
||||
|
||||
if audio_input is not None:
|
||||
audio_embeds = self._process_audio_input(audio_input)
|
||||
multimodal_embeddings.append((audio_embeds, "audio"))
|
||||
if image_input is not None:
|
||||
image_embeds = self._process_image_input(image_input)
|
||||
multimodal_embeddings.append((image_embeds, "image"))
|
||||
if video_input is not None:
|
||||
video_embeds = self._process_video_input(video_input)
|
||||
multimodal_embeddings.append((video_embeds, "video"))
|
||||
return multimodal_embeddings
|
||||
|
||||
def get_input_embeddings_v0(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings: Optional[NestedTensors] = None,
|
||||
) -> torch.Tensor:
|
||||
inputs_embeds = self.language_model.get_input_embeddings(input_ids)
|
||||
if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
|
||||
return inputs_embeds
|
||||
|
||||
for embeddings, modality in multimodal_embeddings:
|
||||
if modality == "audio":
|
||||
placeholder_token_id = self.config.audio_token_index
|
||||
if modality == "image":
|
||||
placeholder_token_id = self.config.image_token_index
|
||||
if modality == "video":
|
||||
placeholder_token_id = self.config.video_token_index
|
||||
inputs_embeds = merge_multimodal_embeddings(
|
||||
input_ids, inputs_embeds, embeddings, placeholder_token_id)
|
||||
return inputs_embeds
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: Optional[IntermediateTensors] = None,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
**kwargs: object,
|
||||
) -> Union[torch.Tensor, IntermediateTensors]:
|
||||
if intermediate_tensors is not None:
|
||||
inputs_embeds = None
|
||||
|
||||
# NOTE: In v1, inputs_embeds is always generated at model runner, this
|
||||
# condition is for v0 compatibility.
|
||||
elif inputs_embeds is None:
|
||||
multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)
|
||||
inputs_embeds = self.get_input_embeddings(
|
||||
input_ids, multimodal_embeddings)
|
||||
text_inputs_embeds = self.get_input_embeddings(input_ids, [(torch.zeros_like(embeddings), modality) for embeddings, modality in multimodal_embeddings] if multimodal_embeddings is not None else None)
|
||||
input_ids = None
|
||||
else:
|
||||
text_inputs_embeds = inputs_embeds
|
||||
|
||||
hidden_states = self.language_model.model(input_ids,
|
||||
positions,
|
||||
intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds)
|
||||
return text_inputs_embeds, hidden_states.unsqueeze(0) #(1, S, D)
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> Optional[torch.Tensor]:
|
||||
return self.language_model.compute_logits(hidden_states,
|
||||
sampling_metadata)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str,
|
||||
torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(
|
||||
self,
|
||||
skip_prefixes=["talker.", "token2wav."],
|
||||
)
|
||||
loaded_weights = loader.load_weights(weights,
|
||||
mapper=self.hf_to_vllm_mapper)
|
||||
# Log load summary
|
||||
try:
|
||||
total_bytes = 0
|
||||
for name, param in self.named_parameters():
|
||||
if param is not None and param.data is not None:
|
||||
total_bytes += param.data.numel() * param.data.element_size()
|
||||
device = next(self.parameters()).device
|
||||
logger.info(
|
||||
"[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s",
|
||||
self.__class__.__name__, True, total_bytes / (1024**2), str(device))
|
||||
except Exception:
|
||||
pass
|
||||
return loaded_weights
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,539 @@
|
||||
from typing import Iterable, Optional, Set, Tuple, Union
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers import Qwen2Config
|
||||
|
||||
from vllm.attention import Attention, AttentionType
|
||||
from vllm.compilation.decorators import support_torch_compile
|
||||
from vllm.config import CacheConfig, VllmConfig
|
||||
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.activation import SiluAndMul
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear)
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.layers.pooler import Pooler, PoolingType
|
||||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
ParallelLMHead, VocabParallelEmbedding)
|
||||
from vllm.model_executor.model_loader.weight_utils import (
|
||||
default_weight_loader, maybe_remap_kv_scale_name)
|
||||
from vllm.v1.pool.metadata import PoolingMetadata
|
||||
from vllm.model_executor.sampling_metadata import SamplingMetadata
|
||||
from vllm.sequence import IntermediateTensors, PoolerOutput
|
||||
|
||||
from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
|
||||
from vllm.model_executor.models.utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper,
|
||||
is_pp_missing_parameter,
|
||||
make_empty_intermediate_tensors_factory, make_layers,
|
||||
maybe_prefix)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class Qwen2MLP(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
hidden_act: str,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.gate_up_proj = MergedColumnParallelLinear(
|
||||
hidden_size,
|
||||
[intermediate_size] * 2,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.gate_up_proj",
|
||||
)
|
||||
self.down_proj = RowParallelLinear(
|
||||
intermediate_size,
|
||||
hidden_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.down_proj",
|
||||
)
|
||||
if hidden_act != "silu":
|
||||
raise ValueError(f"Unsupported activation: {hidden_act}. "
|
||||
"Only silu is supported for now.")
|
||||
self.act_fn = SiluAndMul()
|
||||
|
||||
def forward(self, x):
|
||||
gate_up, _ = self.gate_up_proj(x)
|
||||
x = self.act_fn(gate_up)
|
||||
x, _ = self.down_proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class Qwen2Attention(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
hidden_size: int,
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
max_position: int = 4096 * 32,
|
||||
rope_theta: float = 10000,
|
||||
head_dim: Optional[int] = None,
|
||||
cache_config: Optional[CacheConfig] = None,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
rope_scaling: Optional[Tuple] = None,
|
||||
prefix: str = "",
|
||||
attn_type: str = AttentionType.DECODER) -> 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)
|
||||
self.head_dim = head_dim or (hidden_size // self.total_num_heads)
|
||||
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.qkv_proj = QKVParallelLinear(
|
||||
hidden_size,
|
||||
self.head_dim,
|
||||
self.total_num_heads,
|
||||
self.total_num_kv_heads,
|
||||
bias=True,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.qkv_proj",
|
||||
)
|
||||
self.o_proj = RowParallelLinear(
|
||||
self.total_num_heads * self.head_dim,
|
||||
hidden_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.o_proj",
|
||||
)
|
||||
|
||||
self.rotary_emb = get_rope(
|
||||
self.head_dim,
|
||||
rotary_dim=self.head_dim,
|
||||
max_position=max_position,
|
||||
base=self.rope_theta,
|
||||
rope_scaling=rope_scaling,
|
||||
)
|
||||
self.attn = Attention(self.num_heads,
|
||||
self.head_dim,
|
||||
self.scaling,
|
||||
num_kv_heads=self.num_kv_heads,
|
||||
cache_config=cache_config,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.attn",
|
||||
attn_type=attn_type)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
) -> 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)
|
||||
output, _ = self.o_proj(attn_output)
|
||||
return output
|
||||
|
||||
|
||||
class Qwen2DecoderLayer(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Qwen2Config,
|
||||
cache_config: Optional[CacheConfig] = None,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.hidden_size = config.hidden_size
|
||||
# Requires transformers > 4.32.0
|
||||
rope_theta = getattr(config, "rope_theta", 1000000)
|
||||
rope_scaling = getattr(config, "rope_scaling", None)
|
||||
|
||||
# By default, Qwen2 uses causal attention as it is a decoder-only model.
|
||||
# You can override the HF config with `is_causal=False` to enable
|
||||
# bidirectional attention, which is used in some embedding models
|
||||
# (e.g. Alibaba-NLP/gte-Qwen2-7B-instruct)
|
||||
if getattr(config, "is_causal", True):
|
||||
attn_type = AttentionType.DECODER
|
||||
else:
|
||||
attn_type = AttentionType.ENCODER_ONLY
|
||||
|
||||
self.self_attn = Qwen2Attention(
|
||||
hidden_size=self.hidden_size,
|
||||
num_heads=config.num_attention_heads,
|
||||
max_position=config.max_position_embeddings,
|
||||
num_kv_heads=config.num_key_value_heads,
|
||||
rope_theta=rope_theta,
|
||||
cache_config=cache_config,
|
||||
quant_config=quant_config,
|
||||
rope_scaling=rope_scaling,
|
||||
head_dim=getattr(config, "head_dim", None),
|
||||
prefix=f"{prefix}.self_attn",
|
||||
attn_type=attn_type,
|
||||
)
|
||||
self.mlp = Qwen2MLP(
|
||||
hidden_size=self.hidden_size,
|
||||
intermediate_size=config.intermediate_size,
|
||||
hidden_act=config.hidden_act,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.mlp",
|
||||
)
|
||||
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,
|
||||
residual: Optional[torch.Tensor],
|
||||
) -> 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,
|
||||
)
|
||||
|
||||
# Fully Connected
|
||||
hidden_states, residual = self.post_attention_layernorm(
|
||||
hidden_states, residual)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
return hidden_states, residual
|
||||
|
||||
|
||||
@support_torch_compile(
|
||||
dynamic_arg_dims={
|
||||
"input_ids": 0,
|
||||
# positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
|
||||
# otherwise (seq_len, ).
|
||||
"positions": -1,
|
||||
"intermediate_tensors": 0,
|
||||
"inputs_embeds": 0,
|
||||
})
|
||||
class Qwen2Model(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
*,
|
||||
vllm_config: VllmConfig,
|
||||
prefix: str = "",
|
||||
decoder_layer_type: type[nn.Module] = Qwen2DecoderLayer):
|
||||
super().__init__()
|
||||
|
||||
config = vllm_config.model_config.hf_config
|
||||
cache_config = vllm_config.cache_config
|
||||
quant_config = vllm_config.quant_config
|
||||
|
||||
# TODO (@robertgshaw2): see if this can be moved out
|
||||
if (cache_config.sliding_window is not None
|
||||
and hasattr(config, "max_window_layers")):
|
||||
raise ValueError("Sliding window for some but all layers is not "
|
||||
"supported. This model uses sliding window "
|
||||
"but `max_window_layers` = {} is less than "
|
||||
"`num_hidden_layers` = {}. Please open an issue "
|
||||
"to discuss this feature.".format(
|
||||
config.max_window_layers,
|
||||
config.num_hidden_layers,
|
||||
))
|
||||
|
||||
self.config = config
|
||||
self.quant_config = quant_config
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
if get_pp_group().is_first_rank or (config.tie_word_embeddings
|
||||
and get_pp_group().is_last_rank):
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
config.vocab_size,
|
||||
getattr(config, "embedding_size", config.hidden_size),
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.embed_tokens",
|
||||
)
|
||||
else:
|
||||
self.embed_tokens = PPMissingLayer()
|
||||
|
||||
# Use the provided decoder layer type or default to Qwen2DecoderLayer
|
||||
decoder_layer_type = decoder_layer_type or Qwen2DecoderLayer
|
||||
self.start_layer, self.end_layer, self.layers = make_layers(
|
||||
config.num_hidden_layers,
|
||||
lambda prefix: decoder_layer_type(config=config,
|
||||
cache_config=cache_config,
|
||||
quant_config=quant_config,
|
||||
prefix=prefix),
|
||||
prefix=f"{prefix}.layers",
|
||||
)
|
||||
|
||||
self.make_empty_intermediate_tensors = (
|
||||
make_empty_intermediate_tensors_factory(
|
||||
["hidden_states", "residual"], config.hidden_size))
|
||||
if get_pp_group().is_last_rank:
|
||||
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
else:
|
||||
self.norm = PPMissingLayer()
|
||||
|
||||
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return self.embed_tokens(input_ids)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: Optional[IntermediateTensors] = None,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
) -> Union[torch.Tensor, IntermediateTensors]:
|
||||
if get_pp_group().is_first_rank:
|
||||
if inputs_embeds is not None:
|
||||
hidden_states = inputs_embeds
|
||||
else:
|
||||
hidden_states = self.get_input_embeddings(input_ids)
|
||||
residual = None
|
||||
else:
|
||||
assert intermediate_tensors is not None
|
||||
hidden_states = intermediate_tensors["hidden_states"]
|
||||
residual = intermediate_tensors["residual"]
|
||||
for layer in self.layers[self.start_layer:self.end_layer]:
|
||||
hidden_states, residual = layer(
|
||||
positions,
|
||||
hidden_states,
|
||||
residual,
|
||||
)
|
||||
if not get_pp_group().is_last_rank:
|
||||
return IntermediateTensors({
|
||||
"hidden_states": hidden_states,
|
||||
"residual": residual
|
||||
})
|
||||
hidden_states, _ = self.norm(hidden_states, residual)
|
||||
return hidden_states
|
||||
|
||||
def load_weights(self, weights: Iterable[Tuple[str,
|
||||
torch.Tensor]]) -> Set[str]:
|
||||
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", "gate_proj", 0),
|
||||
("gate_up_proj", "up_proj", 1),
|
||||
]
|
||||
params_dict = dict(self.named_parameters(remove_duplicate=False))
|
||||
loaded_params: Set[str] = set()
|
||||
for name, loaded_weight in weights:
|
||||
if "rotary_emb.inv_freq" in name:
|
||||
continue
|
||||
if (self.quant_config is not None and
|
||||
(scale_name := self.quant_config.get_cache_scale(name))):
|
||||
# Loading kv cache quantization scales
|
||||
param = params_dict[scale_name]
|
||||
weight_loader = getattr(param, "weight_loader",
|
||||
default_weight_loader)
|
||||
loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
|
||||
loaded_weight[0])
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded_params.add(scale_name)
|
||||
continue
|
||||
for (param_name, weight_name, shard_id) in stacked_params_mapping:
|
||||
if weight_name not in name:
|
||||
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
|
||||
if is_pp_missing_parameter(name, self):
|
||||
continue
|
||||
param = params_dict[name]
|
||||
weight_loader = param.weight_loader
|
||||
weight_loader(param, loaded_weight, shard_id)
|
||||
break
|
||||
else:
|
||||
# Skip loading extra bias for GPTQ models.
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
# Remapping the name of FP8 kv-scale.
|
||||
name = maybe_remap_kv_scale_name(name, params_dict)
|
||||
if name is None:
|
||||
continue
|
||||
if is_pp_missing_parameter(name, self):
|
||||
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 Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
}
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
config = vllm_config.model_config.hf_config
|
||||
quant_config = vllm_config.quant_config
|
||||
lora_config = vllm_config.lora_config
|
||||
|
||||
self.config = config
|
||||
self.lora_config = lora_config
|
||||
|
||||
self.quant_config = quant_config
|
||||
self.model = Qwen2Model(vllm_config=vllm_config,
|
||||
prefix=maybe_prefix(prefix, "model"))
|
||||
|
||||
if get_pp_group().is_last_rank:
|
||||
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=maybe_prefix(
|
||||
prefix, "lm_head"))
|
||||
else:
|
||||
self.lm_head = PPMissingLayer()
|
||||
|
||||
self.logits_processor = LogitsProcessor(config.vocab_size)
|
||||
self.sampler = get_sampler()
|
||||
|
||||
self.make_empty_intermediate_tensors = (
|
||||
self.model.make_empty_intermediate_tensors)
|
||||
|
||||
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return self.model.get_input_embeddings(input_ids)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: Optional[IntermediateTensors] = None,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
) -> Union[torch.Tensor, IntermediateTensors]:
|
||||
hidden_states = self.model(input_ids, positions, intermediate_tensors,
|
||||
inputs_embeds)
|
||||
return hidden_states
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> Optional[torch.Tensor]:
|
||||
logits = self.logits_processor(self.lm_head, hidden_states,
|
||||
sampling_metadata)
|
||||
return logits
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> Optional[SamplerOutput]:
|
||||
next_tokens = self.sampler(logits, sampling_metadata)
|
||||
return next_tokens
|
||||
|
||||
def load_weights(self, weights: Iterable[Tuple[str,
|
||||
torch.Tensor]]) -> Set[str]:
|
||||
loader = AutoWeightsLoader(
|
||||
self,
|
||||
skip_prefixes=(["lm_head."]
|
||||
if self.config.tie_word_embeddings else None),
|
||||
)
|
||||
return loader.load_weights(weights)
|
||||
|
||||
|
||||
class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
}
|
||||
|
||||
hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
config = vllm_config.model_config.hf_config
|
||||
quant_config = vllm_config.quant_config
|
||||
lora_config = vllm_config.lora_config
|
||||
pooler_config = vllm_config.model_config.pooler_config
|
||||
|
||||
self.config = config
|
||||
self.lora_config = lora_config
|
||||
|
||||
self.quant_config = quant_config
|
||||
self.model = Qwen2Model(vllm_config=vllm_config,
|
||||
prefix=maybe_prefix(prefix, "model"))
|
||||
|
||||
# TODO: Replace this model class with as_embedding_model(
|
||||
# Qwen2ForCausalLM) after changing the default pooling method
|
||||
if pooler_config.pooling_type is None:
|
||||
logger.warning(
|
||||
"This embedding model will default to last-token pooling in "
|
||||
"an upcoming version. To avoid breaking changes, you should "
|
||||
"pass `--override-pooler-config '{\"pooling_type\": \"MEAN\"}'`"
|
||||
" explicitly.")
|
||||
|
||||
self._pooler = Pooler.from_config_with_defaults(
|
||||
pooler_config,
|
||||
pooling_type=PoolingType.MEAN,
|
||||
normalize=True,
|
||||
softmax=False)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: Optional[IntermediateTensors] = None,
|
||||
) -> torch.Tensor:
|
||||
return self.model(input_ids, positions, intermediate_tensors)
|
||||
|
||||
def pooler(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
pooling_metadata: PoolingMetadata,
|
||||
) -> Optional[PoolerOutput]:
|
||||
return self._pooler(hidden_states, pooling_metadata)
|
||||
|
||||
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
|
||||
weights = self.hf_to_vllm_mapper.apply(weights)
|
||||
weights = ((name, data) for name, data in weights
|
||||
if not name.startswith("lm_head."))
|
||||
self.model.load_weights(weights)
|
||||
@@ -0,0 +1,8 @@
|
||||
from vllm.model_executor.models.utils import maybe_prefix
|
||||
|
||||
|
||||
def add_prefix_to_loaded_weights(weights: set[str], prefix: str) -> set[str]:
|
||||
"""
|
||||
Add a prefix to the names of the loaded weights.
|
||||
"""
|
||||
return {maybe_prefix(prefix, name) for name in weights}
|
||||
Reference in new issue
Block a user