[Model] Add Ernie4.5 VL model support (#15679)
Signed-off-by: CSWYF3634076 <wangyafeng@baidu.com> Signed-off-by: wangyafeng <wangyafeng@baidu.com>
This commit is contained in:
CSWYF3634076
@@ -2285,6 +2285,177 @@ class MRotaryEmbedding(RotaryEmbedding):
|
||||
|
||||
return position_ids, mrope_position_deltas
|
||||
|
||||
@staticmethod
|
||||
def get_rope_index_ernie45(
|
||||
input_ids: torch.Tensor,
|
||||
hf_config: Any,
|
||||
image_grid_thw: Union[list[list[int]], torch.Tensor],
|
||||
video_grid_thw: Union[list[list[int]], torch.Tensor],
|
||||
**kwargs,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Get mrope input positions and delta value for Ernie VL."""
|
||||
|
||||
image_token_id = hf_config.im_patch_id
|
||||
video_start_token_id = hf_config.video_start_token_id
|
||||
video_end_token_id = hf_config.video_end_token_id
|
||||
spatial_conv_size = hf_config.spatial_conv_size
|
||||
temporal_conv_size = hf_config.temporal_conv_size
|
||||
|
||||
mrope_position_deltas = []
|
||||
if input_ids is not None and (
|
||||
image_grid_thw is not None or video_grid_thw is not None
|
||||
):
|
||||
total_input_ids = input_ids
|
||||
position_ids = torch.ones(
|
||||
3,
|
||||
input_ids.shape[0],
|
||||
input_ids.shape[1],
|
||||
dtype=input_ids.dtype,
|
||||
device=input_ids.device,
|
||||
)
|
||||
image_index, video_index = 0, 0
|
||||
for i, input_ids in enumerate(total_input_ids):
|
||||
input_tokens = input_ids.tolist()
|
||||
|
||||
input_token_type = []
|
||||
video_check_flg = False
|
||||
for token in input_tokens:
|
||||
if token == video_start_token_id:
|
||||
video_check_flg = True
|
||||
elif token == video_end_token_id:
|
||||
video_check_flg = False
|
||||
|
||||
if token == image_token_id and not video_check_flg:
|
||||
input_token_type.append("image")
|
||||
elif token == image_token_id and video_check_flg:
|
||||
input_token_type.append("video")
|
||||
else:
|
||||
input_token_type.append("text")
|
||||
|
||||
input_type_group = []
|
||||
for key, group in itertools.groupby(
|
||||
enumerate(input_token_type), lambda x: x[1]
|
||||
):
|
||||
group = list(group)
|
||||
start_index = group[0][0]
|
||||
end_index = group[-1][0] + 1
|
||||
input_type_group.append((key, start_index, end_index))
|
||||
|
||||
llm_pos_ids_list = []
|
||||
video_frame_num = 1
|
||||
for modality_type, start_idx, end_idx in input_type_group:
|
||||
st_idx = (
|
||||
llm_pos_ids_list[-1].max() + 1
|
||||
if len(llm_pos_ids_list) > 0
|
||||
else 0
|
||||
)
|
||||
|
||||
if modality_type == "image":
|
||||
t, h, w = (
|
||||
image_grid_thw[image_index][0],
|
||||
image_grid_thw[image_index][1],
|
||||
image_grid_thw[image_index][2],
|
||||
)
|
||||
llm_grid_t, llm_grid_h, llm_grid_w = (
|
||||
t.item(),
|
||||
h.item() // spatial_conv_size,
|
||||
w.item() // spatial_conv_size,
|
||||
)
|
||||
|
||||
t_index = (
|
||||
torch.arange(llm_grid_t)
|
||||
.view(-1, 1)
|
||||
.expand(-1, llm_grid_h * llm_grid_w)
|
||||
.flatten()
|
||||
)
|
||||
h_index = (
|
||||
torch.arange(llm_grid_h)
|
||||
.view(1, -1, 1)
|
||||
.expand(llm_grid_t, -1, llm_grid_w)
|
||||
.flatten()
|
||||
)
|
||||
w_index = (
|
||||
torch.arange(llm_grid_w)
|
||||
.view(1, 1, -1)
|
||||
.expand(llm_grid_t, llm_grid_h, -1)
|
||||
.flatten()
|
||||
)
|
||||
llm_pos_ids_list.append(
|
||||
torch.stack([t_index, h_index, w_index]) + st_idx
|
||||
)
|
||||
|
||||
image_index += 1
|
||||
video_frame_num = 1
|
||||
|
||||
elif modality_type == "video":
|
||||
t, h, w = (
|
||||
video_grid_thw[video_index][0],
|
||||
video_grid_thw[video_index][1],
|
||||
video_grid_thw[video_index][2],
|
||||
)
|
||||
|
||||
llm_grid_t, llm_grid_h, llm_grid_w = (
|
||||
t.item() // temporal_conv_size,
|
||||
h.item() // spatial_conv_size,
|
||||
w.item() // spatial_conv_size,
|
||||
)
|
||||
|
||||
for t_idx in range(llm_grid_t):
|
||||
t_index = (
|
||||
torch.tensor(t_idx)
|
||||
.view(-1, 1)
|
||||
.expand(-1, llm_grid_h * llm_grid_w)
|
||||
.flatten()
|
||||
)
|
||||
|
||||
h_index = (
|
||||
torch.arange(llm_grid_h)
|
||||
.view(1, -1, 1)
|
||||
.expand(1, -1, llm_grid_w)
|
||||
.flatten()
|
||||
)
|
||||
w_index = (
|
||||
torch.arange(llm_grid_w)
|
||||
.view(1, 1, -1)
|
||||
.expand(1, llm_grid_h, -1)
|
||||
.flatten()
|
||||
)
|
||||
llm_pos_ids_list.append(
|
||||
torch.stack([t_index, h_index, w_index]) + st_idx
|
||||
)
|
||||
|
||||
video_index += 1
|
||||
video_frame_num += 1
|
||||
|
||||
else:
|
||||
text_len = end_idx - start_idx
|
||||
llm_pos_ids_list.append(
|
||||
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
|
||||
)
|
||||
|
||||
video_frame_num = 1
|
||||
|
||||
llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
|
||||
position_ids[..., i, :] = llm_positions.to(position_ids.device)
|
||||
mrope_position_deltas.append(
|
||||
llm_positions.max() + 1 - len(total_input_ids[i])
|
||||
)
|
||||
mrope_position_deltas = torch.tensor(
|
||||
mrope_position_deltas, device=input_ids.device
|
||||
).unsqueeze(1)
|
||||
return position_ids, mrope_position_deltas
|
||||
else:
|
||||
s = input_ids.shape[1]
|
||||
position_ids = torch.arange(s)
|
||||
position_ids = (
|
||||
position_ids.unsqueeze(0).expand(3, -1, -1).to(input_ids.device)
|
||||
)
|
||||
max_position_ids = position_ids.max(0, keepdim=False)[0].max(
|
||||
-1, keepdim=True
|
||||
)[0]
|
||||
mrope_position_deltas = max_position_ids + 1 - s
|
||||
return position_ids, mrope_position_deltas
|
||||
|
||||
# For qwen3-omni
|
||||
@staticmethod
|
||||
def _get_feat_extract_output_lengths(input_lengths):
|
||||
@@ -2324,6 +2495,91 @@ class MRotaryEmbedding(RotaryEmbedding):
|
||||
return llm_pos_ids
|
||||
|
||||
|
||||
class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):
|
||||
"""3D rotary positional embedding. [h w h w h w h w... t t t...]"""
|
||||
|
||||
def forward_native( # type: ignore[override]
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
assert positions.ndim == 1 or positions.ndim == 2
|
||||
assert key is not None
|
||||
|
||||
num_tokens = positions.shape[-1]
|
||||
cos_sin = self.cos_sin_cache[positions]
|
||||
cos, sin = cos_sin.chunk(2, dim=-1)
|
||||
if positions.ndim == 2:
|
||||
assert self.mrope_section
|
||||
|
||||
section_h = self.mrope_section[0] # 22
|
||||
section_w = self.mrope_section[1] # 22
|
||||
section_t = self.mrope_section[2] # 20
|
||||
assert section_h == section_w
|
||||
# Split according to [h w h w h w h w... t t t...]
|
||||
section_cos_t = cos[..., -section_t:]
|
||||
section_cos_h = cos[..., : section_h + section_w : 2]
|
||||
section_cos_w = cos[..., 1 : section_h + section_w : 2]
|
||||
|
||||
cos_t, cos_h, cos_w = section_cos_t[0], section_cos_h[1], section_cos_w[2]
|
||||
cos_hw = torch.stack([cos_h, cos_w], dim=-1).reshape(
|
||||
cos_h.shape[:-1] + (cos_h.shape[-1] * 2,)
|
||||
)
|
||||
cos = torch.cat([cos_hw, cos_t], dim=-1)
|
||||
|
||||
section_sin_t = sin[..., -section_t:]
|
||||
section_sin_h = sin[..., : section_h + section_w : 2]
|
||||
section_sin_w = sin[..., 1 : section_h + section_w : 2]
|
||||
|
||||
sin_t, sin_h, sin_w = section_sin_t[0], section_sin_h[1], section_sin_w[2]
|
||||
sin_hw = torch.stack([sin_h, sin_w], dim=-1).reshape(
|
||||
sin_h.shape[:-1] + (sin_h.shape[-1] * 2,)
|
||||
)
|
||||
sin = torch.cat([sin_hw, sin_t], dim=-1)
|
||||
|
||||
query_shape = query.shape
|
||||
query = query.view(num_tokens, -1, self.head_size)
|
||||
query_rot = query[..., : self.rotary_dim]
|
||||
query_pass = query[..., self.rotary_dim :]
|
||||
query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
|
||||
query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
|
||||
|
||||
key_shape = key.shape
|
||||
key = key.view(num_tokens, -1, self.head_size)
|
||||
key_rot = key[..., : self.rotary_dim]
|
||||
key_pass = key[..., self.rotary_dim :]
|
||||
key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
|
||||
key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
|
||||
return query, key
|
||||
|
||||
def forward_cuda( # type: ignore[override]
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
return self.forward_native(positions, query, key)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Forward pass with optional Triton kernel acceleration.
|
||||
Args:
|
||||
positions:
|
||||
[num_tokens,] (text only) or
|
||||
[3, num_tokens] (T/H/W positions with multimodal inputs)
|
||||
query: [num_tokens, num_heads * head_size]
|
||||
key: [num_tokens, num_kv_heads * head_size]
|
||||
"""
|
||||
assert positions.ndim == 1 or positions.ndim == 2
|
||||
return self.forward_native(positions, query, key)
|
||||
|
||||
|
||||
class DualChunkRotaryEmbedding(MultiPlatformOp):
|
||||
"""Rotary positional embedding for Dual Chunk Attention."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user