update code

configuration_yuanvl.py

@@ -1,6 +1,6 @@
 # --------------------------------------------------------
 # InternVL
-# Copyright (c) 2024 OpenGVLab
+# Copyright (c) 2024 YuanLabAI
 # Licensed under The MIT License [see LICENSE for details]
 # --------------------------------------------------------
 

conversation.py

@@ -391,7 +391,7 @@ register_conv_template(
     Conversation(
         name='yuan-chat',
         system_template='<|im_start|>system\n{system_message}',
-        system_message='你是IEI-源多模态模型，英文名是YuanVL，是由浪潮信息开发的多模态大语言模型。',
+        system_message='你是Yuan3.0 Flash多模态大模型，由YuanLab.ai 团队开发的多模态大语言模型。',
         roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
         sep_style=SeparatorStyle.MPT,
         sep='<|im_end|>\n',

modeling_yuanlm2.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+# Copyright 2022 YuanLabAI and the HuggingFace Inc. team. All rights reserved.
 #
 # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
 # and OPT implementations in this library. It has been modified from its
@@ -26,17 +26,15 @@ import torch.utils.checkpoint
 from torch import einsum, nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers.activations import ACT2FN
-from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
-from .configuration_yuan import YuanConfig
+from configuration_yuan import YuanConfig
 from einops import rearrange
 # from flash_attn import flash_attn_varlen_func as flash_attn_unpadded_func
 #from apex.normalization import MixedFusedRMSNorm as RMSNorm
 #from flash_attn import flash_attn_func
-#from transformer_engine.pytorch import RMSNorm
-import pdb
+from transformer_engine.pytorch import RMSNorm
 import copy
 try:
     import grouped_gemm as gg
@@ -53,39 +51,6 @@ logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "YuanConfig"
 
-class RMSNorm(torch.nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
-        super().__init__()
-        self.weight = torch.nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states):
-        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-
-        # convert into half-precision if necessary
-        if self.weight.dtype in [torch.float16, torch.bfloat16]:
-            hidden_states = hidden_states.to(self.weight.dtype)
-
-        return self.weight * hidden_states
-
-
-"""
-class YuanRotaryEmbedding(nn.Module):
-    def __init__(self, dim, base=10000, dtype=torch.float32, device=None, scaling_factor=1.0, rope_type='default'):
-        super().__init__()
-        inv_freq = (1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))).to(dtype)#.to('cuda:1')
-        self.register_buffer('inv_freq', inv_freq)
-
-    def forward(self, max_seq_len, offset=0):
-        self.inv_freq = self.inv_freq.to(torch.float32)
-        seq = torch.arange(max_seq_len, device=self.inv_freq.device) + offset
-        freqs = einsum('i , j -> i j', seq.type_as(self.inv_freq), self.inv_freq)
-        # first part even vector components, second part odd vector components,
-        #  2 * dim in dimension size
-        emb = torch.cat((freqs, freqs), dim=-1)
-        # emb [seq_length, .., dim]
-        return emb[:, None, None, :]"""
 
 class YuanRotaryEmbedding(nn.Module):
     def __init__(self, dim, base=10000, dtype=torch.float32, rotary_interleaved=False, seq_len_interpolation_factor=None):
@@ -143,17 +108,10 @@ class YuanRotaryEmbedding(nn.Module):
             )
         # emb [seq_length, .., dim]
         emb = emb[:, None, None, :]
-        #emb = emb[:, None, :] 
         return emb
 
 
 def _rotate_half(x, rotary_interleaved):
-    """huggingface version
-    change sign so the last dimension becomes [-odd, +even]
-    
-    x1, x2 = torch.chunk(x, 2, dim=-1)
-    return torch.cat((-x2, x1), dim=-1)
-    """
     if not rotary_interleaved:
         x1, x2 = torch.chunk(x, 2, dim=-1)
         return torch.cat((-x2, x1), dim=-1)
@@ -180,24 +138,6 @@ def apply_rotary_pos_emb(t, freqs, position_ids, rotary_interleaved=False):
 
     t = (t * cos_) + (_rotate_half(t, rotary_interleaved) * sin_)
     return torch.cat((t, t_pass), dim=-1)
-    """huggingface version
-    input tensor t is of shape [seq_length, ..., dim]
-    rotary positional embeding tensor freqs is of shape [seq_length, ..., dim]
-    check https://kexue.fm/archives/8265 for detailed formulas
-    
-    dtype = t.dtype
-    rot_dim = freqs.shape[-1]
-    t_pass = t[..., rot_dim:]
-    if position_ids.shape[1] > 1:
-       freqs = freqs[position_ids]
-       freqs = freqs.view(t.shape[1],freqs.shape[1],freqs.shape[2],freqs.shape[4]).transpose(0,1)
-    # ideally t_pass is empty so rotary pos embedding is applied to all tensor t
-    t = t[..., :rot_dim]
-    # first part is cosine component
-    # second part is sine component, need to change signs with _rotate_half method
-    t = (t * freqs.cos()) + (_rotate_half(t) * freqs.sin())
-    t = t.to(dtype)
-    """
 
     return torch.cat((t, t_pass), dim=-1)
 
@@ -287,53 +227,6 @@ class LocalizedFiltering(torch.nn.Module):
             lf_output = self.output_layernorm(output2 + residual)
 
         return lf_output
-        '''#IEIyuan huggingface version
-        if before_hidden_states == None:
-            inputs = inputs.transpose(0,1)
-            seq_len, bsz, embed_dim = inputs.size()
-            if embed_dim != self.embed_dim:
-                raise ValueError(
-                    f"Unexpected embedding dimension received: input is {embed_dim}, model expects {self.embed_dim}"
-                )
-            residual = inputs
-            inputs = inputs.view(seq_len, 1, bsz, embed_dim).permute(2, 3, 0, 1)
-            inputs = torch.cat((torch.zeros(bsz, embed_dim, 1, 1, dtype=inputs.dtype, device=inputs.device), inputs), dim=2).contiguous()
-            output1 = self.conv1(inputs)
-
-            output1 = torch.cat((torch.zeros(bsz, embed_dim // 2, 1, 1, dtype=inputs.dtype, device=inputs.device), output1), dim=2).contiguous()
-            output2 = self.conv2(output1).permute(2, 3, 0, 1).contiguous()
-            output2 = output2.view(seq_len, bsz, embed_dim)
-            assert output2.shape == residual.shape
-            norm_input = (output2 + residual)#.to('cuda:0')
-            torch.cuda.set_device(norm_input.device)
-            lf_output = self.output_layernorm(norm_input)
-            lf_output = lf_output#.to('cuda:1')
-            lf_output = lf_output.transpose(0,1)
-            return lf_output
-        else:
-            inputs = inputs.transpose(0,1)
-            before_hidden_states = before_hidden_states.transpose(0,1)
-            seq_len, bsz, embed_dim = inputs.size()
-            if embed_dim != self.embed_dim:
-                raise ValueError(
-                    f"Unexpected embedding dimension received: input is {embed_dim}, model expects {self.embed_dim}"
-                )
-            residual = inputs
-            inputs = inputs.view(seq_len, 1, bsz, embed_dim).permute(2, 3, 0, 1)
-            before_hidden_states = before_hidden_states.view(2, 1, bsz, embed_dim).permute(2, 3, 0, 1)
-            inputs = torch.cat((before_hidden_states, inputs), dim=2).contiguous()
-            output1 = self.conv1(inputs)
-            output2 = self.conv2(output1).permute(2, 3, 0, 1).contiguous()
-            output2 = output2.view(seq_len, bsz, embed_dim)
-            assert output2.shape == residual.shape
-
-            norm_input = (output2 + residual)#.to('cuda:0')
-            torch.cuda.set_device(norm_input.device)
-            lf_output = self.output_layernorm(norm_input)
-            lf_output = lf_output#.to('cuda:1')
-            lf_output = lf_output.transpose(0,1)
-            return lf_output
-        '''
 
 
     def forward(
@@ -445,8 +338,6 @@ class FlashSelfAttention(torch.nn.Module):
             # only on first autoregressive step q,k,v have same seqlen
             is_causal = seqlen_q == seqlen_k
             cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32, device=q.device) 
-            #cu_seqlens_q = [cu_seqlens_q[0], cu_seqlens_q[-1]]
-            #cu_seqlens_k = [cu_seqlens_k[0], cu_seqlens_k[-1]]
             dropout_p = 0
 
         output = flash_attn_unpadded_func(q, k, v, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen_k, dropout_p, softmax_scale=self.softmax_scale, causal=is_causal)
@@ -579,8 +470,6 @@ class YuanAttention(nn.Module):
             self.lf_gate = LocalizedFiltering(self.hidden_size, self.lf_conv2d_group, self.lf_conv2d_num_pad)
             self.get_query_key = nn.Linear(self.hidden_size, 2 * self.attention_projection_size, bias=False)
         self.core_attention = FlashSelfAttention(causal=True, attention_dropout=config.attn_dropout, softmax_scale=self.softmax_scale)
-        #self.core_attention_flash = DotProductAttention(num_attention_heads=self.num_heads, 
-        #                                          kv_channels=self.head_dim)
 
     def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
         return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
@@ -596,6 +485,7 @@ class YuanAttention(nn.Module):
         output_attentions: bool = False,
         use_cache: bool = False,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        
         q_len, bsz, _ = hidden_states.size()
         hidden_states = hidden_states#.to('cuda:1')
         is_first_step = False
@@ -621,7 +511,6 @@ class YuanAttention(nn.Module):
         else:
             hidden_states = self.lf_gate(hidden_states, before_hidden_states)
             mixed_qk_layer = self.get_query_key(hidden_states)
-            #mixed_qk_layer = torch.matmul(hidden_states, qk_tensor)
             new_tensor_shape = mixed_qk_layer.size()[:-1] + (self.num_heads, 2 * self.head_dim)
             mixed_qk_layer = mixed_qk_layer.view(*new_tensor_shape)
             (query_states, key_states) = torch.split(mixed_qk_layer, self.head_dim, dim=-1)
@@ -635,7 +524,6 @@ class YuanAttention(nn.Module):
         if rotary_pos_emb is not None:
             if position_ids.shape[1] == 1:
                 q_seq_start = position_ids[0,-1]
-                #seq_start = past_key_value[0].shape[0]
                 q_seq_end = q_seq_start + 1
                 k_seq_end = q_seq_end
             else:
@@ -654,17 +542,11 @@ class YuanAttention(nn.Module):
             key_states = torch.cat([past_key_value[0], key_states], dim=0)
             value_states = torch.cat([past_key_value[1], value_states], dim=0)
         past_key_value = (key_states, value_states, inference_hidden_states_memory) if use_cache else None
-        #query_states = apply_rotary_pos_emb(query_states.permute(1, 0, 2, 3), q_pos_emb, position_ids)
-        #key_states = apply_rotary_pos_emb(key_states.permute(1, 0, 2, 3), k_pos_emb, position_ids)
         query_states = apply_rotary_pos_emb(query_states, q_pos_emb, position_ids)
         key_states = apply_rotary_pos_emb(key_states, k_pos_emb, position_ids_k)
 
         attn_weights = None
-        #query_states = query_states.transpose(0,1)
-        #key_states = key_states.transpose(0,1)
-        #value_states = value_states
         attn_output = self.core_attention(query_states, key_states, value_states)
-        #attn_output = self.core_attention(query_states, key_states, value_states, attention_mask)
         q_len, bsz, _, _ = attn_output.shape
         attn_output = attn_output.reshape(q_len, bsz, -1)
         
@@ -755,7 +637,6 @@ class GroupedMLP(nn.Module):
             return torch.nn.functional.silu(x[0]) * x[1]
 
         self.activation_func = glu
-        #self.ffn_hidden_size = config.moe_config['ffn_hidden_size']
         self.ffn_hidden_size = config.ffn_hidden_size
         fc1_output_size_per_partition = self.ffn_hidden_size * 2
         fc2_input_size = self.ffn_hidden_size
@@ -765,10 +646,6 @@ class GroupedMLP(nn.Module):
     def forward(self, permuted_hidden_states, tokens_per_expert):
         torch.cuda.set_device(permuted_hidden_states.device)
         permuted_hidden_states = permuted_hidden_states#.to('cuda:0')
-        #fc1_output = gg.ops.gmm(permuted_hidden_states, self.weight1, tokens_per_expert.cpu(), trans_b=False)
-
-        #intermediate_parallel = self.activation_func(fc1_output)
-        #fc2_output = gg.ops.gmm(intermediate_parallel, self.weight2, tokens_per_expert.cpu(), trans_b=False)
 
         fc2_outputs = []
         start_idx = 0
@@ -776,13 +653,10 @@ class GroupedMLP(nn.Module):
             if tokens_per_expert[i] == 0:
                 continue
             end_idx = start_idx + tokens_per_expert[i]
-            #fc1_output = torch.matmul(permuted_hidden_states[start_idx:end_idx], self.w1[i])
             # Use custom attributes for each expert's Linear layers
             
             fc1_output = self.w1[i](permuted_hidden_states[start_idx:end_idx])
-            #print("shape1:", self.w1[i].shape, "shape2:", permuted_hidden_states[start_idx:end_idx].shape)
             intermediate_parallel = self.activation_func(fc1_output)
-            #fc2_output = torch.matmul(intermediate_parallel, self.w2[i])
             fc2_output = self.w2[i](intermediate_parallel)
             fc2_outputs.append(fc2_output)
             start_idx = end_idx
@@ -800,7 +674,6 @@ class YuanMoeLayer(nn.Module):
         
         expert_indices_offset = (0)
 
-        #self.gate = ParallelAttention_router(config)
         self.router = ParallelAttention_router(config)
         self.token_dispatcher = MoEDroplessTokenDispatcher(self.num_experts, config=self.config)
         self.experts = GroupedMLP(self.num_experts, self.config)
@@ -812,7 +685,6 @@ class YuanMoeLayer(nn.Module):
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         batch_size, sequence_length, hidden_dim = hidden_states.shape
-        #logits = self.gate(hidden_states)
         logits = self.router(hidden_states)
         scores, indices = self.routing(logits)
         scores = scores.to(hidden_states.dtype)
@@ -868,6 +740,7 @@ class YuanDecoderLayer(nn.Module):
         residual = hidden_states#.to('cuda:1')
         torch.cuda.set_device(hidden_states.device)
         hidden_states = self.input_layernorm(hidden_states) #.to('cuda:0')).to('cuda:1')
+        
         # Self Attention
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
             hidden_states=hidden_states,
@@ -879,7 +752,9 @@ class YuanDecoderLayer(nn.Module):
             output_attentions=output_attentions,
             use_cache=use_cache,
         )
+
         hidden_states = residual + hidden_states.permute(1, 0, 2)
+
         # Fully Connected
         residual = hidden_states#.to('cuda:1')
         torch.cuda.set_device(hidden_states.device)
@@ -1159,14 +1034,10 @@ class YuanModel(YuanPreTrainedModel):
 
         seq_length_with_past = seq_length
         past_key_values_length = 0
+
         if past_key_values is not None:
-            #past_key_values_length = past_key_values[0][0].shape[2]
-            #modify
-            print('0000')
             past_key_values_length = past_key_values[0][0].shape[0]
             seq_length_with_past = seq_length_with_past + past_key_values_length
-        else:
-            print('1111')
 
         # modify to reset position ids
         if past_key_values is not None:
@@ -1187,6 +1058,7 @@ class YuanModel(YuanPreTrainedModel):
             )
             position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
         else:
+            #position_ids = position_ids.view(-1, seq_length).long()
             pass
         
         if inputs_embeds is None:
@@ -1206,6 +1078,11 @@ class YuanModel(YuanPreTrainedModel):
         #rotary_pos_emb = self.rotary_pos_emb(self.max_position_embeddings)
         # Rotary positional embeddings (embedding is None for PP intermediate devices)
         rotary_pos_emb = None
+        '''
+        rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
+                transformer_input=inputs_embeds
+            )
+        '''
         rotary_pos_emb = self.rotary_pos_emb(self.max_position_embeddings)
 
         hidden_states = inputs_embeds
@@ -1220,8 +1097,8 @@ class YuanModel(YuanPreTrainedModel):
         all_hidden_states = () if output_hidden_states else None
         all_self_attns = () if output_attentions else None
         next_decoder_cache = () if use_cache else None
-        #position_ids = position_ids.cpu()
-        #position_ids_k = position_ids_k.cpu()
+        position_ids = position_ids.cpu()
+        position_ids_k = position_ids_k.cpu()
         for idx, decoder_layer in enumerate(self.layers):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
@@ -1262,9 +1139,8 @@ class YuanModel(YuanPreTrainedModel):
             if output_attentions:
                 all_self_attns += (layer_outputs[1],)
         hidden_states = hidden_states#.to('cuda:0')
-        #torch.cuda.set_device(hidden_states.device)
+        torch.cuda.set_device(hidden_states.device)
         hidden_states = self.norm(hidden_states)
-        #print(hidden_states)
         # add hidden states from the last decoder layer
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
@@ -1279,12 +1155,11 @@ class YuanModel(YuanPreTrainedModel):
         )
 
 
-class YuanForCausalLM(YuanPreTrainedModel, GenerationMixin):
+class YuanForCausalLM(YuanPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.model = YuanModel(config)
-        #self.output = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.post_init()
 
     def get_input_embeddings(self):
@@ -1425,8 +1300,9 @@ class YuanForCausalLM(YuanPreTrainedModel, GenerationMixin):
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
         )
+
         hidden_states = outputs[0].transpose(0,1)
-        #print(hidden_states) 
+        
         logits = self.lm_head(hidden_states)
         
         loss = None

modeling_yuanvl_chat.py

@@ -1,6 +1,6 @@
 # --------------------------------------------------------
 # YuanVL
-# Copyright (c) 2024 OpenGVLab
+# Copyright (c) 2024 YuanLabAI
 # Licensed under The MIT License [see LICENSE for details]
 # --------------------------------------------------------
 
@@ -17,10 +17,9 @@ from transformers import (AutoModel, GenerationConfig, LlamaForCausalLM,
                           LlamaTokenizer)
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
-from transformers.generation import GenerationMixin
 from transformers.utils import ModelOutput, logging
 
-#from transformer_engine.pytorch import RMSNorm
+from transformer_engine.pytorch import RMSNorm
 from transformers.activations import ACT2FN
 
 from .configuration_yuanvl import YuanVLChatConfig
@@ -31,22 +30,6 @@ from .utils import flatten_bn, merge_multimodal_embeddings
 
 logger = logging.get_logger(__name__)
 
-class RMSNorm(torch.nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
-        super().__init__()
-        self.weight = torch.nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states):
-        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-
-        # convert into half-precision if necessary
-        if self.weight.dtype in [torch.float16, torch.bfloat16]:
-            hidden_states = hidden_states.to(self.weight.dtype)
-
-        return self.weight * hidden_states
-
 class InternVLImagePixelInputs(TypedDict):
     type: Literal["pixel_values"]
     data: Union[torch.Tensor, List[torch.Tensor]]
@@ -94,6 +77,9 @@ class YuanImageMLP(nn.Module):
         hidden_act: str,
     ) -> None:
         super().__init__()
+        #self.up_proj = ColumnParallelLinear(hidden_size, intermediate_size, bias=False,)
+        #self.gate_proj = ColumnParallelLinear(hidden_size, intermediate_size, bias=False,)
+        #self.down_proj = RowParallelLinear(intermediate_size, output_size, bias=False,)
         self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
         self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
         self.down_proj = nn.Linear(intermediate_size, output_size, bias=False)
@@ -108,13 +94,16 @@ class YuanImageMLP(nn.Module):
         return self.act_fn(y_1) * y_2
     
     def forward(self, x):
+        #import pdb
         x1 = self.up_proj(x)
         x2 = self.gate_proj(x)
         x3 = self.swiglu(x1, x2)
+        #x3 = self.act_fn(x1)
+        #x2 = self.gate_proj(x)
         x = self.down_proj(x3)
         return x
 
-class YuanVLChatModel(PreTrainedModel, GenerationMixin):
+class YuanVLChatModel(PreTrainedModel):
     config_class = YuanVLChatConfig
     main_input_name = 'pixel_values'
     base_model_prefix = 'language_model'
@@ -152,6 +141,10 @@ class YuanVLChatModel(PreTrainedModel, GenerationMixin):
                 raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.')
 
         self.pixel_unshuffle = torch.nn.PixelUnshuffle(downscale_factor=2)
+        #vit_hidden_size = config.vision_config.hidden_size
+        #llm_hidden_size = config.llm_config.hidden_size
+        #vit_mlp_ffn_hidden_size = config.vit_mlp_ffn_hidden_size
+        #layernorm_epsilon = config.llm_config.layernorm_epsilon
         layernorm_epsilon = config.llm_config.rms_norm_eps
 
         self.imagemlp_input_hiddensize = int(config.vision_config.hidden_size / self.downsample_ratio ** 2)
@@ -161,6 +154,18 @@ class YuanVLChatModel(PreTrainedModel, GenerationMixin):
                      output_size=config.llm_config.hidden_size, hidden_act="silu")
         self.imagemlp_layernorm = RMSNorm(config.llm_config.hidden_size, eps=layernorm_epsilon)
 
+        '''
+        # modify internvl vision
+        vit_hidden_size = config.vision_config.hidden_size
+        llm_hidden_size = config.llm_config.hidden_size
+        self.mlp1 = nn.Sequential(
+            nn.LayerNorm(vit_hidden_size * int(1/self.downsample_ratio) ** 2),
+            nn.Linear(vit_hidden_size * int(1/self.downsample_ratio) ** 2, llm_hidden_size),
+            nn.GELU(),
+            nn.Linear(llm_hidden_size, llm_hidden_size)
+        )
+        '''
+
         self.img_context_token_id = config.img_context_token_id
         self.conv_template = get_conv_template(self.template)
         self.system_message = self.conv_template.system_message
@@ -235,6 +240,7 @@ class YuanVLChatModel(PreTrainedModel, GenerationMixin):
         assert self.vision_model is not None
         # (total_patches, tokens_per_image, llm_config.hidden_size)
         image_embeds = self.extract_feature(image_input["data"])
+
         patches_per_image = image_input["patches_per_image"]
 
         # Only one image in the current batch
@@ -289,7 +295,7 @@ class YuanVLChatModel(PreTrainedModel, GenerationMixin):
                 input_ids, inputs_embeds, multimodal_embeddings,
                 self.img_context_token_id)
         return inputs_embeds
-        
+    
     def forward(
             self,
             input_ids: torch.LongTensor = None,
@@ -306,21 +312,22 @@ class YuanVLChatModel(PreTrainedModel, GenerationMixin):
             image_token_id: Optional[List[torch.Tensor]] = None,
             image_embeds: Optional[List[torch.Tensor]] = None,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
-        
+
+        import pdb
+        pdb.set_trace()
         if inputs_embeds is None:
             # (images, patches * token_per_image)
             vision_embeddings = self.get_multimodal_embeddings(pixel_values, image_token_id, image_embeds)
             # (tokens, hidden_size)
-            if input_ids is not None:
-                vision_embeddings = vision_embeddings.to(input_ids.device)
-            inputs_embeds = self.get_input_embeddings(input_ids, vision_embeddings)  #.permute(1, 0, 2)
+            inputs_embeds = self.get_input_embeddings(input_ids, vision_embeddings).permute(1, 0, 2)
             input_ids = None
         
         hidden_states = self.language_model.model(input_ids, attention_mask, position_ids, past_key_values, 
                                                   inputs_embeds, labels, use_cache, output_attentions, 
                                                   output_hidden_states, return_dict)
+
         return hidden_states
-    
+
     def pixel_shuffle(self, x, scale_factor=0.5):
         n, w, h, c = x.size()
         # N, W, H, C --> N, W, H * scale, C // scale
@@ -380,8 +387,6 @@ class YuanVLChatModel(PreTrainedModel, GenerationMixin):
                 vit_embeds = visual_features
             else:
                 vit_embeds = self.get_multimodal_embeddings(pixel_values)
-                if input_ids is not None:
-                    vit_embeds = vit_embeds.to(input_ids.device)
                 inputs_embeds = self.get_input_embeddings(input_ids, vit_embeds)
                 input_ids = None