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
@@ -35,7 +35,6 @@ from einops import rearrange
 #from apex.normalization import MixedFusedRMSNorm as RMSNorm
 #from flash_attn import flash_attn_func
 from transformer_engine.pytorch import RMSNorm
-import pdb
 import copy
 try:
     import grouped_gemm as gg
@@ -52,22 +51,6 @@ logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "YuanConfig"
 
-"""
-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):
@@ -125,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)
@@ -162,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)
 
@@ -269,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(
@@ -427,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)
@@ -561,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()
@@ -582,7 +489,6 @@ class YuanAttention(nn.Module):
         q_len, bsz, _ = hidden_states.size()
         hidden_states = hidden_states#.to('cuda:1')
         is_first_step = False
-        import pdb
         if use_cache:
             if past_key_value is None:
                 before_hidden_states = None
@@ -605,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)
@@ -619,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:
@@ -633,26 +537,16 @@ class YuanAttention(nn.Module):
             else:
                 rotary_pos_emb = ((rotary_pos_emb,) * 2)
             q_pos_emb, k_pos_emb = rotary_pos_emb
-            #q_pos_emb = q_pos_emb[q_seq_start:q_seq_end]
-            #k_pos_emb = k_pos_emb[:k_seq_end]
-        #import pdb
-        #pdb.set_trace()
         if past_key_value is not None:
             # reuse k, v, self_attention
             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)
         
@@ -743,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
@@ -753,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
@@ -764,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
@@ -788,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)
@@ -800,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)
@@ -853,7 +737,6 @@ class YuanDecoderLayer(nn.Module):
                 (see `past_key_values`).
             past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
         """
-        import pdb
         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')
@@ -870,9 +753,6 @@ class YuanDecoderLayer(nn.Module):
             use_cache=use_cache,
         )
 
-        import pdb
-        #print(hidden_states)
-        #pdb.set_trace()
         hidden_states = residual + hidden_states.permute(1, 0, 2)
 
         # Fully Connected
@@ -1156,8 +1036,6 @@ class YuanModel(YuanPreTrainedModel):
         past_key_values_length = 0
 
         if past_key_values is not None:
-            #past_key_values_length = past_key_values[0][0].shape[2]
-            #modify
             past_key_values_length = past_key_values[0][0].shape[0]
             seq_length_with_past = seq_length_with_past + past_key_values_length
 
@@ -1170,13 +1048,8 @@ class YuanModel(YuanPreTrainedModel):
             position_ids_k = torch.cat((position_ids, position_ids_k), dim=1)
             position_ids = position_ids[:,-1]+past_key_values[0][0].shape[0]-position_ids.shape[1]+1
             position_ids = position_ids.unsqueeze(0)
-            #print(position_ids_k,position_ids)
-            #print(position_ids_k.shape,position_ids.shape)
         else:
             position_ids_k = position_ids
-            #print(position_ids)
-            #import pdb
-            #pdb.set_trace()
 
         if position_ids is None:
             device = input_ids.device if input_ids is not None else inputs_embeds.device
@@ -1285,20 +1158,8 @@ class YuanModel(YuanPreTrainedModel):
 class YuanForCausalLM(YuanPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
-        '''
-        self.eod_token = config.eod_token
-        self.sep_token = config.sep_token
-        self.use_loss_mask = config.use_loss_mask
-        self.model = YuanModel(config)
-
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-        '''
         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):

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]
 # --------------------------------------------------------
 
@@ -353,11 +353,6 @@ class YuanVLChatModel(PreTrainedModel):
         # vit_embeds: (imbs * num_images, h*w, vit_dim)
         vit_embeds = vit_embeds[:, 1:, :]
 
-        '''h = w = int(vit_embeds.shape[1]**0.5)
-        # vit_embeds: (imbs * num_images, vit_dim, h, w)
-        vit_embeds = vit_embeds.view(vit_embeds.shape[0], h, w, -1)
-        vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
-        '''
         pn, phw, pc = vit_embeds.shape
         ph = pw = int(phw**0.5)
         vit_embeds = vit_embeds.view(pn, ph, pw, pc).permute(0, 3, 1, 2)