fix rotary_emb for llama

Author: EnflameGCU

paddlenlp/transformers/llama/modeling.py

@@ -935,7 +935,7 @@ def forward(
 
             else:
                 if self.config.use_long_sequence_strategies:
-                    cos, sin = self.rotary_emb(seq_len=kv_seq_len)
+                    cos, sin, _ = self.rotary_emb(seq_len=kv_seq_len)
                     cos = cos[None, :, None, :]
                     sin = sin[None, :, None, :]
                     cos, sin = (

paddlenlp/transformers/llama/modeling_auto.py

@@ -448,7 +448,7 @@ def forward(
                 assert past_key_value is None, "fuse rotary not support cache kv for now"
                 batch_size, seq_length, num_heads, head_dim = query_states.shape
                 _, kv_seq_len, num_key_value_heads, _ = key_states.shape
-                cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+                cos, sin, _ = self.rotary_emb(value_states, seq_len=kv_seq_len)
 
                 paddle_version = float(paddle.__version__[:3])
                 if ((paddle_version != 0.0) and (paddle_version <= 2.6)) and (num_heads != num_key_value_heads):
@@ -481,7 +481,7 @@ def forward(
                         use_neox_rotary_style=False,
                     )
             else:
-                cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+                cos, sin, _ = self.rotary_emb(value_states, seq_len=kv_seq_len)
                 # hack here, because elementwise infer spmd not support broadcast now
                 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
 

paddlenlp/transformers/llama/modeling_auto_static.py

@@ -421,7 +421,7 @@ def forward(
         if self.config.rope:
             if self.use_fused_rope:
                 assert past_key_value is None, "fuse rotary not support cache kv for now"
-                cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+                cos, sin, _ = self.rotary_emb(value_states, seq_len=kv_seq_len)
                 query_states, key_states, _ = fused_rotary_position_embedding(
                     query_states,
                     key_states,
@@ -432,7 +432,7 @@ def forward(
                     use_neox_rotary_style=False,
                 )
             else:
-                cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+                cos, sin, _ = self.rotary_emb(value_states, seq_len=kv_seq_len)
                 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
 
         # [bs, seq_len, num_head, head_dim]