fix rotary_emb for llama

PaddlePaddle · May 20, 2024 · 98a9cc2 · 98a9cc2
1 parent b36b6a0
commit 98a9cc2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/paddlenlp/transformers/llama/modeling_auto.py b/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)
 

diff --git a/paddlenlp/transformers/llama/modeling_auto_static.py b/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]