Update modeling_yuanlm2.py

modeling_yuanlm2.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+# Copyright 2022 YuanLab 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
@@ -32,11 +32,6 @@ 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 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
 import copy
 try:
     import grouped_gemm as gg
@@ -70,23 +65,6 @@ class RMSNorm(torch.nn.Module):
         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):
         super().__init__()
@@ -143,17 +121,11 @@ 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)
@@ -166,7 +138,6 @@ def _rotate_half(x, rotary_interleaved):
 def apply_rotary_pos_emb(t, freqs, position_ids, rotary_interleaved=False):
 
     rot_dim = freqs.shape[-1]
-    #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
@@ -180,26 +151,7 @@ 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)
 
 class LocalizedFiltering(torch.nn.Module):
     """
@@ -287,54 +239,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(
         self,
@@ -444,9 +348,7 @@ class FlashSelfAttention(torch.nn.Module):
             # turn off FA causal mask after first inference autoregressive iteration
             # 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]]
+            cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32, device=q.device)
             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)
@@ -473,7 +375,6 @@ class ParallelAttention_router(nn.Module):
         is_first_step = False
         before_hidden_states = None
         
-        #mixed_x_layer = torch.matmul(hidden_states, self.query_key_value)
         mixed_x_layer = self.query_key_value(hidden_states)
         (query_layer, key_layer, value_layer) = torch.split(mixed_x_layer, self.projection_size, -1)
         b, s, z = query_layer.shape
@@ -493,7 +394,6 @@ class YuanExpertMLP(nn.Module):
     def __init__(self, config):
         super(YuanExpertMLP, self).__init__()
         self.gated_linear_unit = config.moe_config['gated_linear_unit']
-        #self.ffn_hidden_size = config.moe_config['ffn_hidden_size']
         self.ffn_hidden_size = config.ffn_hidden_size
 
 
@@ -579,8 +479,7 @@ 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()
@@ -597,7 +496,7 @@ class YuanAttention(nn.Module):
         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')
+        hidden_states = hidden_states
         is_first_step = False
         if use_cache:
             if past_key_value is None:
@@ -621,7 +520,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 +533,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 +551,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)
         
@@ -764,11 +655,8 @@ class GroupedMLP(nn.Module):
         self.w2 = nn.ModuleList([nn.Linear(self.ffn_hidden_size, self.config.hidden_size, bias=False) for _ in range(num_experts)])
     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)
+        permuted_hidden_states = permuted_hidden_states
 
-        #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,18 +664,15 @@ 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
         fc2_output = torch.cat(fc2_outputs, dim=0)
-        return fc2_output#.to('cuda:1')
+        return fc2_output
 
 class YuanMoeLayer(nn.Module):
     def __init__(self, config:YuanConfig):
@@ -800,7 +685,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 +696,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)
@@ -865,9 +748,9 @@ class YuanDecoderLayer(nn.Module):
                 (see `past_key_values`).
             past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
         """
-        residual = hidden_states#.to('cuda:1')
+        residual = hidden_states
         torch.cuda.set_device(hidden_states.device)
-        hidden_states = self.input_layernorm(hidden_states) #.to('cuda:0')).to('cuda:1')
+        hidden_states = self.input_layernorm(hidden_states)
         # Self Attention
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
             hidden_states=hidden_states,
@@ -881,10 +764,10 @@ class YuanDecoderLayer(nn.Module):
         )
         hidden_states = residual + hidden_states.permute(1, 0, 2)
         # Fully Connected
-        residual = hidden_states#.to('cuda:1')
+        residual = hidden_states
         torch.cuda.set_device(hidden_states.device)
-        hidden_states = self.post_attention_layernorm(hidden_states) #.to('cuda:0')).to('cuda:1')
-        hidden_states = self.mlp(hidden_states)# .to('cuda:1')
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
         outputs = (hidden_states,)
 
@@ -1160,13 +1043,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')
+            pass
 
         # modify to reset position ids
         if past_key_values is not None:
@@ -1203,7 +1083,6 @@ class YuanModel(YuanPreTrainedModel):
                     attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
                 )
 
-        #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_pos_emb = self.rotary_pos_emb(self.max_position_embeddings)
@@ -1220,8 +1099,7 @@ 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()
+
         for idx, decoder_layer in enumerate(self.layers):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
@@ -1261,10 +1139,9 @@ 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)
+        hidden_states = hidden_states
+
         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,)
@@ -1284,7 +1161,6 @@ class YuanForCausalLM(YuanPreTrainedModel, GenerationMixin):
         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):