wehub-resource-sync
1298 lines
62 KiB
Python
1298 lines
62 KiB
Python
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
|
|
# Copyright (c) 2024 The Qwen Team and The HuggingFace Inc. team.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# coding: utf-8
|
|
|
|
|
|
from dataclasses import dataclass
|
|
from functools import partial
|
|
from typing import List, Optional, Tuple
|
|
from einops import rearrange
|
|
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn.attention import SDPBackend, sdpa_kernel
|
|
from torch.nn.attention.flex_attention import flex_attention
|
|
from torch.nn.functional import scaled_dot_product_attention
|
|
from transformers.utils import ModelOutput
|
|
|
|
from flash_attn import flash_attn_varlen_func
|
|
from modeling.qwen2.modeling_qwen2 import (
|
|
Qwen2Attention,
|
|
Qwen2MLP,
|
|
Qwen2PreTrainedModel,
|
|
Qwen2RMSNorm,
|
|
Qwen2RotaryEmbedding,
|
|
apply_rotary_pos_emb,
|
|
)
|
|
from modeling.qwen2_5_vl.modeling_qwen2_5_vl import (
|
|
Qwen2_5_VLRotaryEmbedding,
|
|
apply_multimodal_rotary_pos_emb,
|
|
)
|
|
from modeling.qwen2.configuration_qwen2 import Qwen2Config
|
|
|
|
torch._dynamo.config.cache_size_limit = 512
|
|
torch._dynamo.config.accumulated_cache_size_limit = 4096
|
|
flex_attention = torch.compile(flex_attention)
|
|
|
|
class NaiveCache:
|
|
def __init__(self, num_layers):
|
|
self.key_cache = {k: None for k in range(num_layers)}
|
|
self.value_cache = {k: None for k in range(num_layers)}
|
|
|
|
@property
|
|
def num_layers(self):
|
|
return len(self.key_cache)
|
|
|
|
@property
|
|
def seq_lens(self):
|
|
if self.key_cache[0] is not None:
|
|
return self.key_cache[0].shape[0]
|
|
else:
|
|
return 0
|
|
|
|
|
|
@dataclass
|
|
class BaseNavitOutputWithPast(ModelOutput):
|
|
packed_query_sequence: torch.FloatTensor = None
|
|
past_key_values: Optional[NaiveCache] = None
|
|
|
|
|
|
def pad_sequence(tensor, pad_size):
|
|
H, L, D = tensor.shape
|
|
pad_tensor = tensor.new_zeros((H, pad_size, D))
|
|
return torch.cat([tensor, pad_tensor], dim=1)
|
|
|
|
|
|
class PackedAttention(Qwen2Attention):
|
|
# TODO: currently unused; QK norm logic has not been updated for this path.
|
|
def __init__(self, config, layer_idx: Optional[int] = None):
|
|
super().__init__(config, layer_idx)
|
|
if self.config.qk_norm:
|
|
self.q_norm = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
|
self.k_norm = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
|
else:
|
|
self.q_norm = nn.Identity()
|
|
self.k_norm = nn.Identity()
|
|
|
|
def forward(self, *args, **kwargs):
|
|
if self.training or kwargs.get("mode_forward") == "validation":
|
|
return self.forward_train(*args, **kwargs)
|
|
else:
|
|
return self.forward_inference(*args, **kwargs)
|
|
|
|
def forward_train(
|
|
self,
|
|
packed_sequence: torch.Tensor,
|
|
sample_lens: List[int],
|
|
attention_mask: List[torch.Tensor],
|
|
packed_position_embeddings: Tuple[torch.Tensor, torch.Tensor],
|
|
**kwargs
|
|
):
|
|
packed_query_states = self.q_proj(packed_sequence).view(-1, self.num_heads, self.head_dim)
|
|
packed_key_states = self.k_proj(packed_sequence).view(-1, self.num_key_value_heads, self.head_dim)
|
|
packed_value_states = self.v_proj(packed_sequence).view(-1, self.num_key_value_heads, self.head_dim)
|
|
|
|
packed_query_states = self.q_norm(packed_query_states)
|
|
packed_key_states = self.k_norm(packed_key_states)
|
|
|
|
packed_cos, packed_sin = packed_position_embeddings
|
|
if kwargs.get("apply_qwen_2_5_vl_pos_emb"): # kwargs.get("vit_type") == 'qwen_2_5_vl_original':
|
|
packed_query_states = rearrange(packed_query_states, "l (b h) d -> b h l d", h=self.num_heads)
|
|
packed_key_states = rearrange(packed_key_states, "l (b h) d -> b h l d", h=self.num_key_value_heads)
|
|
packed_query_states, packed_key_states = apply_multimodal_rotary_pos_emb(
|
|
packed_query_states, packed_key_states, packed_cos, packed_sin, self.config.rope_scaling["mrope_section"]
|
|
)
|
|
packed_query_states = rearrange(packed_query_states, "b h l d -> l (b h) d")
|
|
packed_key_states = rearrange(packed_key_states, "b h l d -> l (b h) d")
|
|
else:
|
|
packed_query_states, packed_key_states = apply_rotary_pos_emb(packed_query_states, packed_key_states, packed_cos, packed_sin, unsqueeze_dim=1)
|
|
|
|
if isinstance(attention_mask, List):
|
|
packed_key_states = packed_key_states[:, :, None, :].repeat(1, 1, self.num_key_value_groups, 1)
|
|
packed_key_states = packed_key_states.reshape(-1, self.num_heads, self.head_dim)
|
|
packed_value_states = packed_value_states[:, :, None, :].repeat(1, 1, self.num_key_value_groups, 1)
|
|
packed_value_states = packed_value_states.reshape(-1, self.num_heads, self.head_dim)
|
|
|
|
unpacked_query_states = packed_query_states.transpose(0, 1).split(sample_lens, dim=1)
|
|
unpacked_key_states = packed_key_states.transpose(0, 1).split(sample_lens, dim=1)
|
|
unpacked_value_states = packed_value_states.transpose(0, 1).split(sample_lens, dim=1)
|
|
upacked_attn_output = []
|
|
for query_states, key_states, value_states, attention_mask_per_sample in zip(unpacked_query_states, unpacked_key_states, unpacked_value_states, attention_mask):
|
|
with sdpa_kernel(backends=[SDPBackend.EFFICIENT_ATTENTION]):
|
|
attn_output = scaled_dot_product_attention(
|
|
query_states.to(torch.bfloat16).unsqueeze(0),
|
|
key_states.to(torch.bfloat16).unsqueeze(0),
|
|
value_states.to(torch.bfloat16).unsqueeze(0),
|
|
attention_mask_per_sample.to(torch.bfloat16).unsqueeze(0),
|
|
)
|
|
upacked_attn_output.append(attn_output.squeeze(0))
|
|
packed_attn_output = torch.cat(upacked_attn_output, dim=1)
|
|
else:
|
|
pad_size = sum(sample_lens) - packed_query_states.shape[0]
|
|
packed_query_states = pad_sequence(packed_query_states.permute(1, 0, 2), pad_size)
|
|
packed_key_states = pad_sequence(packed_key_states.permute(1, 0, 2), pad_size)
|
|
packed_value_states = pad_sequence(packed_value_states.permute(1, 0, 2), pad_size)
|
|
packed_attn_output = flex_attention(
|
|
packed_query_states.unsqueeze(0),
|
|
packed_key_states.unsqueeze(0),
|
|
packed_value_states.unsqueeze(0),
|
|
enable_gqa=True,
|
|
block_mask=attention_mask,
|
|
)
|
|
end_index = packed_attn_output.shape[2] - pad_size
|
|
packed_attn_output = packed_attn_output[0, :, :end_index, :]
|
|
|
|
packed_attn_output = packed_attn_output.transpose(0, 1).reshape(-1, self.hidden_size)
|
|
packed_attn_output = self.o_proj(packed_attn_output)
|
|
|
|
return packed_attn_output
|
|
|
|
def forward_inference(
|
|
self,
|
|
packed_query_sequence: torch.Tensor,
|
|
query_lens: torch.Tensor,
|
|
packed_query_position_embeddings: torch.Tensor,
|
|
packed_query_indexes: torch.Tensor,
|
|
past_key_values: Optional[NaiveCache] = None,
|
|
key_values_lens: Optional[torch.Tensor] = None,
|
|
packed_key_value_indexes: Optional[torch.Tensor] = None,
|
|
update_past_key_values=True,
|
|
is_causal=True,
|
|
**kwargs
|
|
):
|
|
packed_query_states = self.q_proj(packed_query_sequence).view(-1, self.num_heads, self.head_dim)
|
|
packed_key_states = self.k_proj(packed_query_sequence).view(-1, self.num_key_value_heads, self.head_dim)
|
|
packed_value_states = self.v_proj(packed_query_sequence).view(-1, self.num_key_value_heads, self.head_dim)
|
|
|
|
packed_query_states = self.q_norm(packed_query_states)
|
|
packed_key_states = self.k_norm(packed_key_states)
|
|
|
|
packed_cos, packed_sin = packed_query_position_embeddings
|
|
packed_query_states, packed_key_states = apply_rotary_pos_emb(packed_query_states, packed_key_states, packed_cos, packed_sin, unsqueeze_dim=1)
|
|
|
|
packed_query_states = packed_query_states.to(torch.bfloat16)
|
|
packed_key_states = packed_key_states.to(torch.bfloat16)
|
|
packed_value_states = packed_value_states.to(torch.bfloat16)
|
|
|
|
if past_key_values is not None and past_key_values.key_cache[self.layer_idx] is not None:
|
|
past_key_states = past_key_values.key_cache[self.layer_idx]
|
|
past_value_states = past_key_values.value_cache[self.layer_idx]
|
|
|
|
seqlens = sum(query_lens) + sum(key_values_lens)
|
|
merged_key_states = past_key_states.new_zeros((seqlens, self.num_key_value_heads, self.head_dim))
|
|
merged_value_states = past_key_states.new_zeros((seqlens, self.num_key_value_heads, self.head_dim))
|
|
merged_key_states[packed_query_indexes] = packed_key_states
|
|
merged_key_states[packed_key_value_indexes] = past_key_states
|
|
merged_value_states[packed_query_indexes] = packed_value_states
|
|
merged_value_states[packed_key_value_indexes] = past_value_states
|
|
key_values_lens = key_values_lens + query_lens
|
|
else:
|
|
merged_key_states = packed_key_states
|
|
merged_value_states = packed_value_states
|
|
key_values_lens = query_lens
|
|
|
|
cu_seqlens_q = torch.nn.functional.pad(torch.cumsum(query_lens, dim=0), (1, 0))
|
|
cu_seqlens_k = torch.nn.functional.pad(torch.cumsum(key_values_lens, dim=0), (1, 0))
|
|
|
|
packed_attn_output = flash_attn_varlen_func(
|
|
q=packed_query_states,
|
|
k=merged_key_states,
|
|
v=merged_value_states,
|
|
cu_seqlens_q=cu_seqlens_q.to(torch.int32),
|
|
cu_seqlens_k=cu_seqlens_k.to(torch.int32),
|
|
max_seqlen_q=max(query_lens).item(),
|
|
max_seqlen_k=max(key_values_lens).item(),
|
|
causal=is_causal,
|
|
)
|
|
packed_attn_output = packed_attn_output.reshape(-1, self.hidden_size)
|
|
packed_attn_output = self.o_proj(packed_attn_output)
|
|
|
|
if update_past_key_values:
|
|
past_key_values.key_cache[self.layer_idx] = merged_key_states
|
|
past_key_values.value_cache[self.layer_idx] = merged_value_states
|
|
|
|
return packed_attn_output, past_key_values
|
|
|
|
|
|
class PackedAttentionMoT(Qwen2Attention):
|
|
def __init__(self, config, layer_idx: Optional[int] = None):
|
|
super().__init__(config, layer_idx)
|
|
if self.config.qk_norm_und or self.config.qk_norm_gen:
|
|
# NOTE: initialize understanding and generation norms separately.
|
|
# Understanding.
|
|
self.q_norm = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps) if self.config.qk_norm_und else nn.Identity()
|
|
self.k_norm = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps) if self.config.qk_norm_und else nn.Identity()
|
|
|
|
# Generation.
|
|
self.q_norm_moe_gen = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps) if self.config.qk_norm_gen else nn.Identity()
|
|
self.k_norm_moe_gen = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps) if self.config.qk_norm_gen else nn.Identity()
|
|
else:
|
|
# NOTE: use the shared initialization path.
|
|
if self.config.qk_norm:
|
|
self.q_norm = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
|
self.k_norm = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
|
self.q_norm_moe_gen = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
|
self.k_norm_moe_gen = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
|
else:
|
|
self.q_norm = nn.Identity()
|
|
self.k_norm = nn.Identity()
|
|
self.q_norm_moe_gen = nn.Identity()
|
|
self.k_norm_moe_gen = nn.Identity()
|
|
|
|
self.q_proj_moe_gen = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)
|
|
self.k_proj_moe_gen = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
|
|
self.v_proj_moe_gen = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
|
|
self.o_proj_moe_gen = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
|
|
|
|
self.layer_idx = layer_idx
|
|
|
|
def forward(self, *args, **kwargs):
|
|
if self.training or kwargs.get("mode_forward") == "validation":
|
|
return self.forward_train(*args, **kwargs)
|
|
else:
|
|
return self.forward_inference(*args, **kwargs)
|
|
|
|
def forward_train(
|
|
self,
|
|
packed_sequence: torch.Tensor,
|
|
sample_lens: List[int],
|
|
attention_mask,
|
|
packed_position_embeddings: Tuple[torch.Tensor, torch.Tensor],
|
|
packed_und_token_indexes: torch.LongTensor,
|
|
packed_gen_token_indexes: torch.LongTensor,
|
|
mode=None,
|
|
**kwargs,
|
|
):
|
|
packed_query_states = packed_sequence.new_zeros((packed_sequence.shape[0], self.num_heads * self.head_dim))
|
|
packed_key_states = packed_sequence.new_zeros((packed_sequence.shape[0], self.num_key_value_heads * self.head_dim))
|
|
packed_value_states = packed_sequence.new_zeros((packed_sequence.shape[0], self.num_key_value_heads * self.head_dim))
|
|
|
|
packed_sequence_und = packed_sequence[packed_und_token_indexes]
|
|
packed_sequence_gen = packed_sequence[packed_gen_token_indexes]
|
|
|
|
packed_query_states[packed_und_token_indexes] = self.q_proj(packed_sequence_und)
|
|
packed_query_states[packed_gen_token_indexes] = self.q_proj_moe_gen(packed_sequence_gen)
|
|
|
|
packed_key_states[packed_und_token_indexes] = self.k_proj(packed_sequence_und)
|
|
packed_key_states[packed_gen_token_indexes] = self.k_proj_moe_gen(packed_sequence_gen)
|
|
|
|
packed_value_states[packed_und_token_indexes] = self.v_proj(packed_sequence_und)
|
|
packed_value_states[packed_gen_token_indexes] = self.v_proj_moe_gen(packed_sequence_gen)
|
|
|
|
|
|
packed_query_states = packed_query_states.view(-1, self.num_heads, self.head_dim)
|
|
packed_key_states = packed_key_states.view(-1, self.num_key_value_heads, self.head_dim)
|
|
packed_value_states = packed_value_states.view(-1, self.num_key_value_heads, self.head_dim)
|
|
if self.config.freeze_und:
|
|
packed_value_states[packed_und_token_indexes] = packed_value_states[packed_und_token_indexes].detach()
|
|
|
|
packed_query_states_ = packed_query_states.new_zeros(packed_query_states.shape)
|
|
packed_key_states_ = packed_key_states.new_zeros(packed_key_states.shape)
|
|
|
|
packed_query_states_[packed_und_token_indexes] = self.q_norm(packed_query_states[packed_und_token_indexes])
|
|
if self.config.freeze_und:
|
|
packed_query_states_[packed_und_token_indexes] = packed_query_states_[packed_und_token_indexes].detach()
|
|
packed_query_states_[packed_gen_token_indexes] = self.q_norm_moe_gen(packed_query_states[packed_gen_token_indexes])
|
|
|
|
packed_key_states_[packed_und_token_indexes] = self.k_norm(packed_key_states[packed_und_token_indexes])
|
|
if self.config.freeze_und:
|
|
packed_key_states_[packed_und_token_indexes] = packed_key_states_[packed_und_token_indexes].detach()
|
|
packed_key_states_[packed_gen_token_indexes] = self.k_norm_moe_gen(packed_key_states[packed_gen_token_indexes])
|
|
|
|
packed_cos, packed_sin = packed_position_embeddings
|
|
if kwargs.get("apply_qwen_2_5_vl_pos_emb"): # kwargs.get("vit_type") == 'qwen_2_5_vl_original':
|
|
packed_query_states_ = rearrange(packed_query_states_, "l (b h) d -> b h l d", h=self.num_heads)
|
|
packed_key_states_ = rearrange(packed_key_states_, "l (b h) d -> b h l d", h=self.num_key_value_heads)
|
|
packed_query_states_, packed_key_states_ = apply_multimodal_rotary_pos_emb(
|
|
packed_query_states_, packed_key_states_, packed_cos, packed_sin, self.config.rope_scaling["mrope_section"]
|
|
)
|
|
packed_query_states_ = rearrange(packed_query_states_, "b h l d -> l (b h) d")
|
|
packed_key_states_ = rearrange(packed_key_states_, "b h l d -> l (b h) d")
|
|
else:
|
|
packed_query_states_, packed_key_states_ = apply_rotary_pos_emb(packed_query_states_, packed_key_states_, packed_cos, packed_sin, unsqueeze_dim=1)
|
|
|
|
|
|
if isinstance(attention_mask, List):
|
|
packed_key_states_ = packed_key_states_[:, :, None, :].repeat(1, 1, self.num_key_value_groups, 1)
|
|
packed_key_states_ = packed_key_states_.reshape(-1, self.num_heads, self.head_dim)
|
|
packed_value_states = packed_value_states[:, :, None, :].repeat(1, 1, self.num_key_value_groups, 1)
|
|
packed_value_states = packed_value_states.reshape(-1, self.num_heads, self.head_dim)
|
|
|
|
unpacked_query_states = packed_query_states_.transpose(0, 1).split(sample_lens, dim=1)
|
|
unpacked_key_states = packed_key_states_.transpose(0, 1).split(sample_lens, dim=1)
|
|
unpacked_value_states = packed_value_states.transpose(0, 1).split(sample_lens, dim=1)
|
|
upacked_attn_output = []
|
|
for query_states, key_states, value_states, attention_mask_per_sample in zip(unpacked_query_states, unpacked_key_states, unpacked_value_states, attention_mask):
|
|
with sdpa_kernel(backends=[SDPBackend.EFFICIENT_ATTENTION]):
|
|
attn_output = scaled_dot_product_attention(
|
|
query_states.to(torch.bfloat16).unsqueeze(0),
|
|
key_states.to(torch.bfloat16).unsqueeze(0),
|
|
value_states.to(torch.bfloat16).unsqueeze(0),
|
|
attention_mask_per_sample.to(torch.bfloat16).unsqueeze(0),
|
|
)
|
|
upacked_attn_output.append(attn_output.squeeze(0))
|
|
packed_attn_output = torch.cat(upacked_attn_output, dim=1)
|
|
|
|
packed_attn_output = packed_attn_output.transpose(0, 1).reshape(-1, self.num_heads * self.head_dim)
|
|
packed_attn_output_ = packed_attn_output.new_zeros(packed_attn_output.shape)
|
|
packed_attn_output_[packed_und_token_indexes] = self.o_proj(packed_attn_output[packed_und_token_indexes])
|
|
packed_attn_output_[packed_gen_token_indexes] = self.o_proj_moe_gen(packed_attn_output[packed_gen_token_indexes])
|
|
else: # USED !!!
|
|
pad_size = sum(sample_lens) - packed_query_states.shape[0]
|
|
packed_query_states_ = pad_sequence(packed_query_states_.permute(1, 0, 2), pad_size)
|
|
packed_key_states_ = pad_sequence(packed_key_states_.permute(1, 0, 2), pad_size)
|
|
packed_value_states = pad_sequence(packed_value_states.permute(1, 0, 2), pad_size)
|
|
packed_attn_output = flex_attention(
|
|
packed_query_states_.unsqueeze(0), # 1, num_head, L, head_dim
|
|
packed_key_states_.unsqueeze(0),
|
|
packed_value_states.unsqueeze(0),
|
|
enable_gqa=True,
|
|
block_mask=attention_mask,
|
|
)
|
|
end_index = packed_attn_output.shape[2] - pad_size
|
|
packed_attn_output = packed_attn_output[0, :, :end_index, :]
|
|
|
|
packed_attn_output = packed_attn_output.transpose(0, 1).reshape(-1, self.num_heads * self.head_dim)
|
|
packed_attn_output_ = packed_attn_output.new_zeros(packed_attn_output.shape)
|
|
packed_attn_output_[packed_und_token_indexes] = self.o_proj(packed_attn_output[packed_und_token_indexes])
|
|
packed_attn_output_[packed_gen_token_indexes] = self.o_proj_moe_gen(packed_attn_output[packed_gen_token_indexes])
|
|
|
|
return packed_attn_output_
|
|
|
|
def forward_inference(
|
|
self,
|
|
packed_query_sequence: torch.Tensor,
|
|
query_lens: torch.Tensor,
|
|
packed_query_position_embeddings: torch.Tensor,
|
|
packed_query_indexes: torch.Tensor,
|
|
past_key_values: Optional[NaiveCache] = None,
|
|
key_values_lens: Optional[torch.Tensor] = None,
|
|
packed_key_value_indexes: Optional[torch.Tensor] = None,
|
|
update_past_key_values=True,
|
|
is_causal=True,
|
|
mode="und",
|
|
packed_vae_token_indexes=None,
|
|
packed_text_indexes=None,
|
|
**kwargs
|
|
):
|
|
if mode == "und":
|
|
packed_query_states = self.q_proj(packed_query_sequence).view(-1, self.num_heads, self.head_dim)
|
|
packed_key_states = self.k_proj(packed_query_sequence).view(-1, self.num_key_value_heads, self.head_dim)
|
|
packed_value_states = self.v_proj(packed_query_sequence).view(-1, self.num_key_value_heads, self.head_dim)
|
|
packed_query_states = self.q_norm(packed_query_states)
|
|
packed_key_states = self.k_norm(packed_key_states)
|
|
elif mode == "gen":
|
|
packed_query_sequence = packed_query_sequence.to(torch.bfloat16)
|
|
packed_query_states = packed_query_sequence.new_zeros((packed_query_sequence.shape[0], self.num_heads * self.head_dim))
|
|
packed_key_states = packed_query_sequence.new_zeros((packed_query_sequence.shape[0], self.num_key_value_heads * self.head_dim))
|
|
packed_value_states = packed_query_sequence.new_zeros((packed_query_sequence.shape[0], self.num_key_value_heads * self.head_dim))
|
|
|
|
packed_text_query_sequence = packed_query_sequence[packed_text_indexes]
|
|
packed_vae_query_sequence = packed_query_sequence[packed_vae_token_indexes]
|
|
|
|
packed_query_states[packed_text_indexes] = self.q_proj(packed_text_query_sequence)
|
|
packed_query_states[packed_vae_token_indexes] = self.q_proj_moe_gen(packed_vae_query_sequence)
|
|
|
|
packed_key_states[packed_text_indexes] = self.k_proj(packed_text_query_sequence)
|
|
packed_key_states[packed_vae_token_indexes] = self.k_proj_moe_gen(packed_vae_query_sequence)
|
|
|
|
packed_value_states[packed_text_indexes] = self.v_proj(packed_text_query_sequence)
|
|
packed_value_states[packed_vae_token_indexes] = self.v_proj_moe_gen(packed_vae_query_sequence)
|
|
|
|
packed_query_states = packed_query_states.view(-1, self.num_heads, self.head_dim)
|
|
packed_key_states = packed_key_states.view(-1, self.num_key_value_heads, self.head_dim)
|
|
packed_value_states = packed_value_states.view(-1, self.num_key_value_heads, self.head_dim)
|
|
|
|
packed_query_states = packed_query_states.to(torch.float32)
|
|
packed_query_states[packed_text_indexes] = self.q_norm(packed_query_states[packed_text_indexes])
|
|
packed_query_states[packed_vae_token_indexes] = self.q_norm_moe_gen(packed_query_states[packed_vae_token_indexes])
|
|
|
|
packed_key_states = packed_key_states.to(torch.float32)
|
|
packed_key_states[packed_text_indexes] = self.k_norm(packed_key_states[packed_text_indexes])
|
|
packed_key_states[packed_vae_token_indexes] = self.k_norm_moe_gen(packed_key_states[packed_vae_token_indexes])
|
|
|
|
packed_cos, packed_sin = packed_query_position_embeddings
|
|
if kwargs.get("apply_qwen_2_5_vl_pos_emb"):
|
|
packed_query_states = rearrange(packed_query_states, "l (b h) d -> b h l d", h=self.num_heads)
|
|
packed_key_states = rearrange(packed_key_states, "l (b h) d -> b h l d", h=self.num_key_value_heads)
|
|
packed_query_states, packed_key_states = apply_multimodal_rotary_pos_emb(
|
|
packed_query_states, packed_key_states, packed_cos, packed_sin, self.config.rope_scaling["mrope_section"]
|
|
)
|
|
packed_query_states = rearrange(packed_query_states, "b h l d -> l (b h) d")
|
|
packed_key_states = rearrange(packed_key_states, "b h l d -> l (b h) d")
|
|
|
|
else:
|
|
packed_query_states, packed_key_states = apply_rotary_pos_emb(packed_query_states, packed_key_states, packed_cos, packed_sin, unsqueeze_dim=1)
|
|
|
|
packed_query_states = packed_query_states.to(torch.bfloat16)
|
|
packed_key_states = packed_key_states.to(torch.bfloat16)
|
|
packed_value_states = packed_value_states.to(torch.bfloat16)
|
|
|
|
if past_key_values is not None and past_key_values.key_cache[self.layer_idx] is not None:
|
|
past_key_states = past_key_values.key_cache[self.layer_idx]
|
|
past_value_states = past_key_values.value_cache[self.layer_idx]
|
|
|
|
seqlens = sum(query_lens) + sum(key_values_lens)
|
|
merged_key_states = past_key_states.new_zeros(size=[seqlens, self.num_key_value_heads, self.head_dim])
|
|
merged_value_states = past_key_states.new_zeros(size=[seqlens, self.num_key_value_heads, self.head_dim])
|
|
merged_key_states[packed_query_indexes] = packed_key_states
|
|
merged_key_states[packed_key_value_indexes] = past_key_states
|
|
merged_value_states[packed_query_indexes] = packed_value_states
|
|
merged_value_states[packed_key_value_indexes] = past_value_states
|
|
key_values_lens = key_values_lens + query_lens
|
|
else:
|
|
merged_key_states = packed_key_states
|
|
merged_value_states = packed_value_states
|
|
key_values_lens = query_lens
|
|
|
|
cu_seqlens_q = torch.nn.functional.pad(torch.cumsum(query_lens, dim=0), (1, 0))
|
|
cu_seqlens_k = torch.nn.functional.pad(torch.cumsum(key_values_lens, dim=0), (1, 0))
|
|
|
|
packed_attn_output = flash_attn_varlen_func(
|
|
q=packed_query_states,
|
|
k=merged_key_states,
|
|
v=merged_value_states,
|
|
cu_seqlens_q=cu_seqlens_q.to(torch.int32),
|
|
cu_seqlens_k=cu_seqlens_k.to(torch.int32),
|
|
max_seqlen_q=max(query_lens).item(),
|
|
max_seqlen_k=max(key_values_lens).item(),
|
|
causal=is_causal,
|
|
)
|
|
packed_attn_output = packed_attn_output.reshape(-1, self.hidden_size)
|
|
if mode == "und":
|
|
packed_attn_output = self.o_proj(packed_attn_output)
|
|
elif mode == "gen":
|
|
packed_attn_output[packed_text_indexes] = self.o_proj(packed_attn_output[packed_text_indexes])
|
|
packed_attn_output[packed_vae_token_indexes] = self.o_proj_moe_gen(packed_attn_output[packed_vae_token_indexes])
|
|
|
|
if update_past_key_values:
|
|
past_key_values.key_cache[self.layer_idx] = merged_key_states
|
|
past_key_values.value_cache[self.layer_idx] = merged_value_states
|
|
|
|
return packed_attn_output, past_key_values
|
|
|
|
|
|
class Qwen2DecoderLayer(nn.Module):
|
|
def __init__(self, config, layer_idx: Optional[int] = None):
|
|
super().__init__()
|
|
self.hidden_size = config.hidden_size
|
|
|
|
self.self_attn = PackedAttention(config, layer_idx)
|
|
|
|
self.mlp = Qwen2MLP(config)
|
|
self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
|
|
def forward(self, *args, **kwargs):
|
|
if self.training or kwargs.get("mode_forward") == "validation":
|
|
return self.forward_train(*args, **kwargs)
|
|
else:
|
|
return self.forward_inference(*args, **kwargs)
|
|
|
|
def forward_train(
|
|
self,
|
|
packed_sequence: torch.Tensor,
|
|
sample_lens: List[int],
|
|
attention_mask,
|
|
packed_position_embeddings: Tuple[torch.Tensor, torch.Tensor],
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
|
|
residual = packed_sequence
|
|
packed_sequence = self.input_layernorm(packed_sequence)
|
|
|
|
# Self Attention
|
|
packed_sequence = self.self_attn(
|
|
packed_sequence=packed_sequence,
|
|
sample_lens=sample_lens,
|
|
attention_mask=attention_mask,
|
|
packed_position_embeddings=packed_position_embeddings,
|
|
**kwargs,
|
|
)
|
|
packed_sequence = residual + packed_sequence
|
|
|
|
# Fully Connected
|
|
residual = packed_sequence
|
|
packed_sequence = self.post_attention_layernorm(packed_sequence)
|
|
packed_sequence = self.mlp(packed_sequence)
|
|
packed_sequence = residual + packed_sequence
|
|
|
|
return packed_sequence
|
|
|
|
def forward_inference(
|
|
self,
|
|
packed_query_sequence: torch.Tensor,
|
|
query_lens: torch.Tensor,
|
|
packed_query_position_embeddings: torch.Tensor,
|
|
packed_query_indexes: torch.Tensor,
|
|
past_key_values: Optional[NaiveCache] = None,
|
|
key_values_lens: Optional[torch.Tensor] = None,
|
|
packed_key_value_indexes: Optional[torch.Tensor] = None,
|
|
update_past_key_values=True,
|
|
is_causal=True,
|
|
**kwargs
|
|
) -> BaseNavitOutputWithPast:
|
|
|
|
residual = packed_query_sequence
|
|
packed_query_sequence = self.input_layernorm(packed_query_sequence)
|
|
|
|
# Self Attention
|
|
packed_query_sequence, past_key_values = self.self_attn(
|
|
packed_query_sequence=packed_query_sequence,
|
|
query_lens=query_lens,
|
|
packed_query_position_embeddings=packed_query_position_embeddings,
|
|
packed_query_indexes=packed_query_indexes,
|
|
past_key_values=past_key_values,
|
|
key_values_lens=key_values_lens,
|
|
packed_key_value_indexes=packed_key_value_indexes,
|
|
update_past_key_values=update_past_key_values,
|
|
is_causal=is_causal,
|
|
**kwargs
|
|
)
|
|
packed_query_sequence = residual + packed_query_sequence
|
|
|
|
# Fully Connected
|
|
residual = packed_query_sequence
|
|
packed_query_sequence = self.post_attention_layernorm(packed_query_sequence)
|
|
packed_query_sequence = self.mlp(packed_query_sequence)
|
|
packed_query_sequence = residual + packed_query_sequence
|
|
|
|
return packed_query_sequence, past_key_values
|
|
|
|
|
|
class Qwen2MoTDecoderLayer(nn.Module):
|
|
def __init__(
|
|
self,
|
|
config,
|
|
layer_idx: Optional[int] = None,
|
|
attn_module: Optional[PackedAttentionMoT] = PackedAttentionMoT,
|
|
):
|
|
super().__init__()
|
|
self.hidden_size = config.hidden_size
|
|
self.freeze_und = config.freeze_und
|
|
|
|
self.self_attn: PackedAttentionMoT = attn_module(config, layer_idx)
|
|
|
|
self.mlp = Qwen2MLP(config)
|
|
self.mlp_moe_gen = Qwen2MLP(config)
|
|
self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
self.input_layernorm_moe_gen = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
self.post_attention_layernorm_moe_gen = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
|
|
def forward(self, *args, **kwargs):
|
|
if self.training or kwargs.get("mode_forward") == "validation":
|
|
return self.forward_train(*args, **kwargs)
|
|
else:
|
|
return self.forward_inference(*args, **kwargs)
|
|
|
|
def forward_train(
|
|
self,
|
|
packed_sequence: torch.Tensor,
|
|
sample_lens: List[int],
|
|
attention_mask,
|
|
packed_position_embeddings: Tuple[torch.Tensor, torch.Tensor],
|
|
packed_und_token_indexes: torch.LongTensor,
|
|
packed_gen_token_indexes: torch.LongTensor,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
|
|
residual = packed_sequence
|
|
packed_sequence_ = packed_sequence.new_zeros(packed_sequence.shape)
|
|
packed_sequence_[packed_und_token_indexes] = self.input_layernorm(packed_sequence[packed_und_token_indexes])
|
|
packed_sequence_[packed_gen_token_indexes] = self.input_layernorm_moe_gen(packed_sequence[packed_gen_token_indexes])
|
|
|
|
# Self Attention
|
|
if attention_mask is not None:
|
|
attention_mask = attention_mask.to(device=packed_sequence_.device)
|
|
|
|
packed_sequence_ = self.self_attn(
|
|
packed_sequence=packed_sequence_,
|
|
sample_lens=sample_lens,
|
|
attention_mask=attention_mask,
|
|
packed_position_embeddings=packed_position_embeddings,
|
|
packed_und_token_indexes=packed_und_token_indexes,
|
|
packed_gen_token_indexes=packed_gen_token_indexes,
|
|
**kwargs,
|
|
)
|
|
if self.freeze_und:
|
|
packed_sequence_[packed_und_token_indexes] = packed_sequence_[packed_und_token_indexes].detach()
|
|
packed_sequence = residual + packed_sequence_
|
|
|
|
# Fully Connected
|
|
residual = packed_sequence
|
|
packed_sequence_ = packed_sequence.new_zeros(packed_sequence.shape)
|
|
packed_sequence_[packed_und_token_indexes] = self.mlp(self.post_attention_layernorm(packed_sequence[packed_und_token_indexes]))
|
|
if self.freeze_und:
|
|
packed_sequence_[packed_und_token_indexes] = packed_sequence_[packed_und_token_indexes].detach()
|
|
|
|
packed_sequence_[packed_gen_token_indexes] = self.mlp_moe_gen(self.post_attention_layernorm_moe_gen(packed_sequence[packed_gen_token_indexes]))
|
|
packed_sequence = residual + packed_sequence_
|
|
|
|
return packed_sequence
|
|
|
|
def forward_inference(
|
|
self,
|
|
packed_query_sequence: torch.Tensor,
|
|
query_lens: torch.Tensor,
|
|
packed_query_position_embeddings: torch.Tensor,
|
|
packed_query_indexes: torch.Tensor,
|
|
past_key_values: Optional[NaiveCache] = None,
|
|
key_values_lens: Optional[torch.Tensor] = None,
|
|
packed_key_value_indexes: Optional[torch.Tensor] = None,
|
|
update_past_key_values=True,
|
|
is_causal=True,
|
|
mode="und",
|
|
packed_vae_token_indexes=None,
|
|
packed_text_indexes=None,
|
|
**kwargs
|
|
) -> BaseNavitOutputWithPast:
|
|
|
|
residual = packed_query_sequence
|
|
if mode == "und":
|
|
packed_query_sequence = self.input_layernorm(packed_query_sequence)
|
|
elif mode == "gen":
|
|
packed_query_sequence_ = torch.zeros_like(packed_query_sequence)
|
|
packed_query_sequence_[packed_text_indexes] = self.input_layernorm(packed_query_sequence[packed_text_indexes])
|
|
packed_query_sequence_[packed_vae_token_indexes] = self.input_layernorm_moe_gen(packed_query_sequence[packed_vae_token_indexes])
|
|
packed_query_sequence = packed_query_sequence_
|
|
|
|
# Self Attention
|
|
packed_query_sequence, past_key_values = self.self_attn(
|
|
packed_query_sequence=packed_query_sequence,
|
|
query_lens=query_lens,
|
|
packed_query_position_embeddings=packed_query_position_embeddings,
|
|
packed_query_indexes=packed_query_indexes,
|
|
past_key_values=past_key_values,
|
|
key_values_lens=key_values_lens,
|
|
packed_key_value_indexes=packed_key_value_indexes,
|
|
update_past_key_values=update_past_key_values,
|
|
is_causal=is_causal,
|
|
mode=mode,
|
|
packed_vae_token_indexes=packed_vae_token_indexes,
|
|
packed_text_indexes=packed_text_indexes,
|
|
**kwargs,
|
|
)
|
|
packed_query_sequence = residual + packed_query_sequence
|
|
|
|
# Fully Connected
|
|
residual = packed_query_sequence
|
|
if mode == "und":
|
|
packed_query_sequence = self.post_attention_layernorm(packed_query_sequence)
|
|
packed_query_sequence = self.mlp(packed_query_sequence)
|
|
elif mode == "gen":
|
|
packed_text_query_sequence = packed_query_sequence[packed_text_indexes]
|
|
packed_vae_query_sequence = packed_query_sequence[packed_vae_token_indexes]
|
|
packed_text_query_sequence = self.post_attention_layernorm(packed_text_query_sequence).to(torch.bfloat16)
|
|
packed_vae_query_sequence = self.post_attention_layernorm_moe_gen(packed_vae_query_sequence).to(torch.bfloat16)
|
|
|
|
packed_query_sequence_ = torch.zeros_like(packed_query_sequence).to(torch.bfloat16)
|
|
packed_query_sequence_[packed_text_indexes] = self.mlp(packed_text_query_sequence)
|
|
packed_query_sequence_[packed_vae_token_indexes] = self.mlp_moe_gen(packed_vae_query_sequence)
|
|
packed_query_sequence = packed_query_sequence_
|
|
|
|
packed_query_sequence = residual + packed_query_sequence
|
|
return packed_query_sequence, past_key_values
|
|
|
|
|
|
class Qwen2MoEDecoderLayer(nn.Module):
|
|
def __init__(self, config, layer_idx: Optional[int] = None):
|
|
super().__init__()
|
|
self.hidden_size = config.hidden_size
|
|
|
|
self.self_attn = PackedAttention(config, layer_idx)
|
|
|
|
self.mlp = Qwen2MLP(config)
|
|
self.mlp_moe_gen = Qwen2MLP(config)
|
|
self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
|
|
def forward(self, *args, **kwargs):
|
|
if self.training or kwargs.get("mode_forward") == "validation":
|
|
return self.forward_train(*args, **kwargs)
|
|
else:
|
|
return self.forward_inference(*args, **kwargs)
|
|
|
|
def forward_train(
|
|
self,
|
|
packed_sequence: torch.Tensor,
|
|
sample_lens: List[int],
|
|
attention_mask,
|
|
packed_position_embeddings: Tuple[torch.Tensor, torch.Tensor],
|
|
packed_und_token_indexes: torch.LongTensor,
|
|
packed_gen_token_indexes: torch.LongTensor,
|
|
mode=None,
|
|
) -> torch.Tensor:
|
|
|
|
residual = packed_sequence
|
|
packed_sequence = self.input_layernorm(packed_sequence)
|
|
|
|
# Self Attention
|
|
packed_sequence = self.self_attn(
|
|
packed_sequence=packed_sequence,
|
|
sample_lens=sample_lens,
|
|
attention_mask=attention_mask,
|
|
packed_position_embeddings=packed_position_embeddings,
|
|
)
|
|
packed_sequence = residual + packed_sequence
|
|
|
|
# Fully Connected
|
|
residual = packed_sequence
|
|
packed_sequence = self.post_attention_layernorm(packed_sequence)
|
|
|
|
packed_sequence_new = packed_sequence.new_zeros(packed_sequence.shape)
|
|
packed_sequence_und = self.mlp(packed_sequence[packed_und_token_indexes])
|
|
packed_sequence_gen = self.mlp_moe_gen(packed_sequence[packed_gen_token_indexes])
|
|
packed_sequence_new[packed_und_token_indexes] = packed_sequence_und
|
|
packed_sequence_new[packed_gen_token_indexes] = packed_sequence_gen
|
|
|
|
packed_sequence = residual + packed_sequence_new
|
|
|
|
return packed_sequence
|
|
|
|
def forward_inference(
|
|
self,
|
|
packed_query_sequence: torch.Tensor,
|
|
query_lens: torch.Tensor,
|
|
packed_query_position_embeddings: torch.Tensor,
|
|
packed_query_indexes: torch.Tensor,
|
|
past_key_values: Optional[NaiveCache] = None,
|
|
key_values_lens: Optional[torch.Tensor] = None,
|
|
packed_key_value_indexes: Optional[torch.Tensor] = None,
|
|
update_past_key_values=True,
|
|
is_causal=True,
|
|
mode="und",
|
|
packed_vae_token_indexes=None,
|
|
packed_text_indexes=None,
|
|
) -> BaseNavitOutputWithPast:
|
|
|
|
residual = packed_query_sequence
|
|
packed_query_sequence = self.input_layernorm(packed_query_sequence)
|
|
|
|
# Self Attention
|
|
packed_query_sequence, past_key_values = self.self_attn(
|
|
packed_query_sequence=packed_query_sequence,
|
|
query_lens=query_lens,
|
|
packed_query_position_embeddings=packed_query_position_embeddings,
|
|
packed_query_indexes=packed_query_indexes,
|
|
past_key_values=past_key_values,
|
|
key_values_lens=key_values_lens,
|
|
packed_key_value_indexes=packed_key_value_indexes,
|
|
update_past_key_values=update_past_key_values,
|
|
is_causal=is_causal,
|
|
)
|
|
packed_query_sequence = residual + packed_query_sequence
|
|
|
|
# Fully Connected
|
|
residual = packed_query_sequence
|
|
packed_query_sequence = self.post_attention_layernorm(packed_query_sequence)
|
|
if mode == "und":
|
|
packed_query_sequence = self.mlp(packed_query_sequence)
|
|
elif mode == "gen":
|
|
packed_query_sequence_ = torch.zeros_like(packed_query_sequence).to(torch.bfloat16)
|
|
packed_query_sequence_[packed_text_indexes] = self.mlp(packed_query_sequence[packed_text_indexes])
|
|
packed_query_sequence_[packed_vae_token_indexes] = self.mlp_moe_gen(packed_query_sequence[packed_vae_token_indexes])
|
|
packed_query_sequence = packed_query_sequence_
|
|
packed_query_sequence = residual + packed_query_sequence
|
|
|
|
return packed_query_sequence, past_key_values
|
|
|
|
|
|
Decoder_layer_dict = {
|
|
"Qwen2DecoderLayer": Qwen2DecoderLayer,
|
|
"Qwen2MoEDecoderLayer": Qwen2MoEDecoderLayer,
|
|
"Qwen2MoTDecoderLayer": partial(Qwen2MoTDecoderLayer, attn_module=PackedAttentionMoT),
|
|
}
|
|
|
|
|
|
class Qwen2Model(Qwen2PreTrainedModel):
|
|
def __init__(self, config: Qwen2Config):
|
|
super().__init__(config)
|
|
self.padding_idx = config.pad_token_id
|
|
self.vocab_size = config.vocab_size
|
|
self.use_moe = "Mo" in config.layer_module
|
|
|
|
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
|
layer_module = Decoder_layer_dict[config.layer_module]
|
|
self.layers = nn.ModuleList([layer_module(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) # here is very slow
|
|
|
|
self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
if self.use_moe:
|
|
self.norm_moe_gen = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
|
|
self.apply_qwen_2_5_vl_pos_emb = config.apply_qwen_2_5_vl_pos_emb
|
|
if self.apply_qwen_2_5_vl_pos_emb:
|
|
self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config)
|
|
else:
|
|
self.rotary_emb = Qwen2RotaryEmbedding(config=config)
|
|
|
|
# Initialize weights and apply final processing
|
|
# self.post_init() # NOTE too slow, not used in inference
|
|
|
|
def forward(self, *args, **kwargs):
|
|
if self.training or kwargs.get("mode_forward") == "validation":
|
|
return self.forward_train(*args, **kwargs)
|
|
else:
|
|
return self.forward_inference(*args, **kwargs)
|
|
|
|
def forward_train(
|
|
self,
|
|
packed_sequence: torch.Tensor,
|
|
sample_lens: List[int],
|
|
attention_mask,
|
|
packed_position_ids: torch.Tensor,
|
|
packed_und_token_indexes: Optional[torch.LongTensor] = None,
|
|
packed_gen_token_indexes: Optional[torch.LongTensor] = None,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
if self.config.freeze_und:
|
|
packed_sequence[packed_und_token_indexes] = packed_sequence[packed_und_token_indexes].detach()
|
|
|
|
# create position embeddings to be shared across the decoder layers
|
|
if self.apply_qwen_2_5_vl_pos_emb:
|
|
packed_position_embeddings = self.rotary_emb(packed_sequence.unsqueeze(0), packed_position_ids)
|
|
kwargs.update({"apply_qwen_2_5_vl_pos_emb": self.apply_qwen_2_5_vl_pos_emb})
|
|
else:
|
|
cos, sin = self.rotary_emb(packed_sequence, packed_position_ids.unsqueeze(0))
|
|
cos = cos.squeeze(0)
|
|
sin = sin.squeeze(0)
|
|
packed_position_embeddings = (cos, sin)
|
|
kwargs.update({"apply_qwen_2_5_vl_pos_emb": self.apply_qwen_2_5_vl_pos_emb})
|
|
|
|
extra_inputs = {}
|
|
if self.use_moe:
|
|
assert packed_und_token_indexes is not None
|
|
if packed_gen_token_indexes is None:
|
|
packed_gen_token_indexes = packed_und_token_indexes.new_ones(size=[0])
|
|
extra_inputs.update(
|
|
packed_und_token_indexes=packed_und_token_indexes,
|
|
packed_gen_token_indexes=packed_gen_token_indexes,
|
|
)
|
|
|
|
for decoder_layer in self.layers:
|
|
attention_mask_ = attention_mask
|
|
packed_sequence = decoder_layer(
|
|
packed_sequence=packed_sequence,
|
|
sample_lens=sample_lens,
|
|
attention_mask=attention_mask_,
|
|
packed_position_embeddings=packed_position_embeddings,
|
|
**extra_inputs,
|
|
**kwargs,
|
|
)
|
|
|
|
if self.use_moe:
|
|
packed_sequence_ = torch.zeros_like(packed_sequence)
|
|
packed_sequence_[packed_und_token_indexes] = self.norm(packed_sequence[packed_und_token_indexes]).to(dtype=packed_sequence.dtype)
|
|
if self.config.freeze_und:
|
|
packed_sequence_[packed_und_token_indexes] = packed_sequence_[packed_und_token_indexes].detach()
|
|
packed_sequence_[packed_gen_token_indexes] = self.norm_moe_gen(packed_sequence[packed_gen_token_indexes]).to(dtype=packed_sequence.dtype)
|
|
return packed_sequence_
|
|
else:
|
|
return self.norm(packed_sequence)
|
|
|
|
def forward_inference(
|
|
self,
|
|
packed_query_sequence: torch.Tensor,
|
|
query_lens: torch.Tensor,
|
|
packed_query_position_ids: torch.Tensor,
|
|
packed_query_indexes: torch.Tensor,
|
|
past_key_values: Optional[NaiveCache] = None,
|
|
key_values_lens: Optional[torch.Tensor] = None,
|
|
packed_key_value_indexes: Optional[torch.Tensor] = None,
|
|
update_past_key_values=True,
|
|
is_causal=True,
|
|
mode="und",
|
|
packed_vae_token_indexes=None,
|
|
packed_text_indexes=None,
|
|
**kwargs,
|
|
) -> BaseNavitOutputWithPast:
|
|
|
|
if self.apply_qwen_2_5_vl_pos_emb:
|
|
packed_query_position_embeddings = self.rotary_emb(packed_query_sequence.unsqueeze(0), packed_query_position_ids)
|
|
kwargs.update({"apply_qwen_2_5_vl_pos_emb": self.apply_qwen_2_5_vl_pos_emb})
|
|
else:
|
|
# create position embeddings to be shared across the decoder layers
|
|
cos, sin = self.rotary_emb(packed_query_sequence, packed_query_position_ids.unsqueeze(0))
|
|
cos = cos.squeeze(0)
|
|
sin = sin.squeeze(0)
|
|
packed_query_position_embeddings = (cos, sin)
|
|
kwargs.update({"apply_qwen_2_5_vl_pos_emb": self.apply_qwen_2_5_vl_pos_emb})
|
|
|
|
extra_inputs = {}
|
|
if self.use_moe:
|
|
extra_inputs.update(mode=mode)
|
|
if mode == "gen":
|
|
assert packed_vae_token_indexes is not None
|
|
assert packed_text_indexes is not None
|
|
extra_inputs.update(
|
|
packed_vae_token_indexes=packed_vae_token_indexes,
|
|
packed_text_indexes=packed_text_indexes,
|
|
)
|
|
|
|
for decoder_layer in self.layers:
|
|
packed_query_sequence, past_key_values = decoder_layer(
|
|
packed_query_sequence=packed_query_sequence,
|
|
query_lens=query_lens,
|
|
packed_query_position_embeddings=packed_query_position_embeddings,
|
|
packed_query_indexes=packed_query_indexes,
|
|
past_key_values=past_key_values,
|
|
key_values_lens=key_values_lens,
|
|
packed_key_value_indexes=packed_key_value_indexes,
|
|
update_past_key_values=update_past_key_values,
|
|
is_causal=is_causal,
|
|
**extra_inputs,
|
|
**kwargs,
|
|
)
|
|
|
|
if self.use_moe:
|
|
if mode == "und":
|
|
packed_query_sequence = self.norm(packed_query_sequence)
|
|
elif mode == "gen":
|
|
packed_query_sequence_ = torch.zeros_like(packed_query_sequence)
|
|
packed_query_sequence_[packed_text_indexes] = self.norm(packed_query_sequence[packed_text_indexes])
|
|
packed_query_sequence_[packed_vae_token_indexes] = self.norm_moe_gen(packed_query_sequence[packed_vae_token_indexes])
|
|
packed_query_sequence = packed_query_sequence_
|
|
else:
|
|
packed_query_sequence = self.norm(packed_query_sequence)
|
|
|
|
return BaseNavitOutputWithPast(
|
|
packed_query_sequence=packed_query_sequence,
|
|
past_key_values=past_key_values,
|
|
)
|
|
|
|
|
|
class Qwen2ForCausalLM(Qwen2PreTrainedModel):
|
|
_tied_weights_keys = ["lm_head.weight"]
|
|
|
|
def __init__(self, config: Qwen2Config):
|
|
super().__init__(config)
|
|
self.model: Qwen2Model = Qwen2Model(config)
|
|
self.vocab_size = config.vocab_size
|
|
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
|
|
|
# Initialize weights and apply final processing
|
|
# self.post_init() # NOTE too slow, not used in inference
|
|
|
|
# Untie and clone to avoid save-time errors from tied-weight aliasing.
|
|
def untie_lm_head(self):
|
|
in_emb = self.get_input_embeddings()
|
|
out_emb = self.get_output_embeddings()
|
|
if out_emb.weight.data.data_ptr() == in_emb.weight.data.data_ptr():
|
|
with torch.no_grad():
|
|
out_emb.weight = torch.nn.Parameter(in_emb.weight.detach().clone())
|
|
# Prevent later automatic re-tying.
|
|
self.config.tie_word_embeddings = False
|
|
if hasattr(self, "_tied_weights_keys"):
|
|
self._tied_weights_keys = []
|
|
|
|
# When the vocab grows, copy newly added rows from input embeddings to lm_head.
|
|
def copy_new_token_rows_to_lm_head(self, num_new_tokens: int):
|
|
with torch.no_grad():
|
|
if num_new_tokens and num_new_tokens > 0:
|
|
in_emb = self.get_input_embeddings()
|
|
out_emb = self.get_output_embeddings()
|
|
with torch.no_grad():
|
|
out_emb.weight[-num_new_tokens:].copy_(in_emb.weight[-num_new_tokens:])
|
|
|
|
def init_moe(self):
|
|
for name, param in self.named_parameters():
|
|
if "moe_gen" in name:
|
|
try:
|
|
original_name = name.replace("_moe_gen", "")
|
|
param.data.copy_(self.state_dict()[original_name].data)
|
|
except KeyError:
|
|
print(f"Warning: {original_name} not found in state_dict, skipping copy.")
|
|
|
|
def freeze_llm_params(self):
|
|
self.eval()
|
|
for param in self.parameters():
|
|
param.requires_grad = False
|
|
|
|
def freeze_embed_tokens(self):
|
|
for name, param in self.model.embed_tokens.named_parameters():
|
|
# print(f'freeze_embed_tokens: {name}')
|
|
param.requires_grad = False
|
|
|
|
def freeze_lm_head(self):
|
|
for name, param in self.lm_head.named_parameters():
|
|
param.requires_grad = False
|
|
|
|
def freeze_und_params(self):
|
|
# NOTE: freeze the understanding-side parameters.
|
|
for name, param in self.named_parameters():
|
|
if "moe_gen" not in name:
|
|
# print(name)
|
|
param.requires_grad = False
|
|
|
|
def get_input_embeddings(self):
|
|
return self.model.embed_tokens
|
|
|
|
def set_input_embeddings(self, value):
|
|
self.model.embed_tokens = value
|
|
|
|
def get_output_embeddings(self):
|
|
return self.lm_head
|
|
|
|
def set_output_embeddings(self, new_embeddings):
|
|
self.lm_head = new_embeddings
|
|
|
|
def set_decoder(self, decoder):
|
|
self.model = decoder
|
|
|
|
def get_decoder(self):
|
|
return self.model
|
|
|
|
def forward(self, *args, **kwargs):
|
|
if self.training or kwargs.get("mode_forward") == "validation":
|
|
return self.forward_train(*args, **kwargs)
|
|
else:
|
|
return self.forward_inference(*args, **kwargs)
|
|
|
|
def forward_train(
|
|
self,
|
|
packed_sequence: torch.Tensor,
|
|
sample_lens: List[int],
|
|
attention_mask,
|
|
packed_position_ids: torch.Tensor,
|
|
packed_und_token_indexes: Optional[torch.LongTensor] = None,
|
|
packed_gen_token_indexes: Optional[torch.LongTensor] = None,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
outputs = self.model.forward(
|
|
packed_sequence=packed_sequence,
|
|
sample_lens=sample_lens,
|
|
packed_position_ids=packed_position_ids,
|
|
attention_mask=attention_mask,
|
|
packed_und_token_indexes=packed_und_token_indexes,
|
|
packed_gen_token_indexes=packed_gen_token_indexes,
|
|
**kwargs,
|
|
)
|
|
return outputs
|
|
|
|
def forward_inference(
|
|
self,
|
|
packed_query_sequence: torch.Tensor,
|
|
query_lens: torch.Tensor,
|
|
packed_query_position_ids: torch.Tensor,
|
|
packed_query_indexes: torch.Tensor,
|
|
past_key_values: Optional[NaiveCache] = None,
|
|
key_values_lens: Optional[torch.Tensor] = None,
|
|
packed_key_value_indexes: Optional[torch.Tensor] = None,
|
|
update_past_key_values=True,
|
|
is_causal=True,
|
|
mode="und",
|
|
packed_vae_token_indexes=None,
|
|
packed_text_indexes=None,
|
|
**kwargs,
|
|
) -> BaseNavitOutputWithPast:
|
|
|
|
outputs = self.model.forward(
|
|
packed_query_sequence=packed_query_sequence,
|
|
query_lens=query_lens,
|
|
packed_query_position_ids=packed_query_position_ids,
|
|
packed_query_indexes=packed_query_indexes,
|
|
past_key_values=past_key_values,
|
|
key_values_lens=key_values_lens,
|
|
packed_key_value_indexes=packed_key_value_indexes,
|
|
update_past_key_values=update_past_key_values,
|
|
is_causal=is_causal,
|
|
mode=mode,
|
|
packed_vae_token_indexes=packed_vae_token_indexes,
|
|
packed_text_indexes=packed_text_indexes,
|
|
**kwargs,
|
|
)
|
|
|
|
return outputs
|
|
|
|
# Compute RoPE indexes for Qwen-VL.
|
|
def get_rope_index(
|
|
self,
|
|
input_ids: Optional[torch.LongTensor] = None,
|
|
image_grid_thw: Optional[torch.LongTensor] = None,
|
|
video_grid_thw: Optional[torch.LongTensor] = None,
|
|
second_per_grid_ts: Optional[torch.Tensor] = None, # Time interval for each video grid.
|
|
attention_mask: Optional[torch.Tensor] = None, # Masks padding tokens; all-ones masks may be omitted.
|
|
image_token_id: int = None,
|
|
video_token_id: int = None,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Calculate the 3D rope index based on image and video's temporal, height and width in LLM.
|
|
|
|
Explanation:
|
|
Each embedding sequence contains vision embedding and text embedding or just contains text embedding.
|
|
|
|
For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.
|
|
Examples:
|
|
input_ids: [T T T T T], here T is for text.
|
|
temporal position_ids: [0, 1, 2, 3, 4]
|
|
height position_ids: [0, 1, 2, 3, 4]
|
|
width position_ids: [0, 1, 2, 3, 4]
|
|
|
|
For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part
|
|
and 1D rotary position embeddin for text part.
|
|
Examples:
|
|
Temporal (Time): 3 patches, representing different segments of the video in time.
|
|
Height: 2 patches, dividing each frame vertically.
|
|
Width: 2 patches, dividing each frame horizontally.
|
|
We also have some important parameters:
|
|
fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second.
|
|
tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity.
|
|
temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames.
|
|
interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs.
|
|
input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.
|
|
vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]
|
|
vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]
|
|
vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
|
|
text temporal position_ids: [101, 102, 103, 104, 105]
|
|
text height position_ids: [101, 102, 103, 104, 105]
|
|
text width position_ids: [101, 102, 103, 104, 105]
|
|
Here we calculate the text start position_ids as the max vision position_ids plus 1.
|
|
|
|
Args:
|
|
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
|
|
Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
|
|
it.
|
|
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
|
|
The temporal, height and width of feature shape of each image in LLM.
|
|
video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
|
|
The temporal, height and width of feature shape of each video in LLM.
|
|
second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*):
|
|
The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.
|
|
attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
|
|
Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
|
|
|
|
- 1 for tokens that are **not masked**,
|
|
- 0 for tokens that are **masked**.
|
|
|
|
Returns:
|
|
position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
|
|
mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
|
|
"""
|
|
spatial_merge_size = self.config.vision_config['spatial_merge_size'] # 2
|
|
image_token_id = self.config.image_token_id
|
|
video_token_id = self.config.video_token_id
|
|
vision_start_token_id = self.config.vision_start_token_id
|
|
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
|
|
if attention_mask is None:
|
|
attention_mask = torch.ones_like(total_input_ids)
|
|
position_ids = torch.ones(
|
|
3,
|
|
input_ids.shape[0],
|
|
input_ids.shape[1],
|
|
dtype=input_ids.dtype,
|
|
device=input_ids.device,
|
|
) # [3, 1, L]
|
|
image_index, video_index = 0, 0
|
|
attention_mask = attention_mask.to(total_input_ids.device)
|
|
for i, input_ids in enumerate(total_input_ids):
|
|
input_ids = input_ids[attention_mask[i] == 1]
|
|
image_nums, video_nums = 0, 0
|
|
vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
|
|
vision_tokens = input_ids[vision_start_indices + 1] # Determine whether the next filled token is image or video.
|
|
image_nums = (vision_tokens == image_token_id).sum()
|
|
video_nums = (vision_tokens == video_token_id).sum()
|
|
input_tokens = input_ids.tolist()
|
|
llm_pos_ids_list: list = []
|
|
st = 0
|
|
remain_images, remain_videos = image_nums, video_nums
|
|
for _ in range(image_nums + video_nums): # Iterate over image/video tokens to find their end positions.
|
|
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],
|
|
)
|
|
second_per_grid_t = 0
|
|
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],
|
|
)
|
|
if second_per_grid_ts is not None:
|
|
second_per_grid_t = second_per_grid_ts[video_index]
|
|
else:
|
|
second_per_grid_t = 1.0
|
|
video_index += 1
|
|
remain_videos -= 1
|
|
ed = ed_video
|
|
llm_grid_t, llm_grid_h, llm_grid_w = (
|
|
t.item(),
|
|
h.item() // spatial_merge_size,
|
|
w.item() // 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)
|
|
|
|
range_tensor = torch.arange(llm_grid_t).view(-1, 1)
|
|
expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w)
|
|
|
|
time_tensor = expanded_range * second_per_grid_t * self.config.vision_config['tokens_per_second']
|
|
|
|
time_tensor_long = time_tensor.long()
|
|
t_index = time_tensor_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)
|
|
position_ids[..., i, attention_mask[i] == 1] = 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:
|
|
if attention_mask is not None:
|
|
position_ids = attention_mask.long().cumsum(-1) - 1
|
|
position_ids.masked_fill_(attention_mask == 0, 1)
|
|
position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
|
|
max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
|
|
mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
|
|
else:
|
|
position_ids = (
|
|
torch.arange(input_ids.shape[1], device=input_ids.device)
|
|
.view(1, 1, -1)
|
|
.expand(3, input_ids.shape[0], -1)
|
|
)
|
|
mrope_position_deltas = torch.zeros(
|
|
[input_ids.shape[0], 1],
|
|
device=input_ids.device,
|
|
dtype=input_ids.dtype,
|
|
)
|
|
|
|
return position_ids, mrope_position_deltas
|