# 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