fix rotary_seq_len bug

InternLM · Jul 12, 2024 · a104a42 · a104a42
1 parent 6d83c33
commit a104a42
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/xtuner/model/modules/dispatch/mistral.py b/xtuner/model/modules/dispatch/mistral.py
@@ -317,7 +317,7 @@ def mistral_varlen_attn_forward(
         value_states = value_states.transpose(1, 2)
         # Because the input can be padded, the absolute sequence length
         # depends on the max position id.
-        rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item() + 1)
+        rotary_seq_len = max(kv_seq_len, position_ids.max().item() + 1)
         cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
         query_states, key_states = apply_rotary_pos_emb(
             query_states, key_states, cos, sin, position_ids)

diff --git a/xtuner/model/modules/dispatch/phi3.py b/xtuner/model/modules/dispatch/phi3.py
@@ -339,7 +339,7 @@ def phi3_varlen_attn_forward(
                                                        self.layer_idx)
 
     assert position_ids is not None
-    rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1
+    rotary_seq_len = max(kv_seq_len, position_ids.max().item() + 1)
     cos, sin = self.rotary_emb(
         value_states, position_ids, seq_len=rotary_seq_len)
 

diff --git a/xtuner/model/modules/dispatch/qwen2.py b/xtuner/model/modules/dispatch/qwen2.py
@@ -233,7 +233,7 @@ def qwen2_varlen_attn_forward(
                                                        self.layer_idx)
 
     assert position_ids is not None
-    rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item() + 1)
+    rotary_seq_len = max(kv_seq_len, position_ids.max().item() + 1)
     cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
 
     query_states, key_states = apply_rotary_pos_emb(query_states, key_states,