Update modeling_internlm.py

modeling_internlm.py

@@ -243,22 +243,10 @@ def rotate_half(x):
 
 # Copied from transformers.model.llama.modeling_llama.apply_rotary_pos_emb
 def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
-    if position_ids.size(1) == 1:
-        q_cos = cos[position_ids].unsqueeze(1).expand(q.shape)
-        q_sin = sin[position_ids].unsqueeze(1).expand(q.shape)
-        q_embed = (q * q_cos) + (rotate_half(q) * q_sin)
-
-        position_ids = position_ids.flatten() + 1
-        max_length = max(position_ids)
-        position_ids = torch.stack([torch.cat([torch.ones(max_length - w, dtype=torch.long), torch.arange(w)]) for w in position_ids])
-        k_cos = cos[position_ids].unsqueeze(1).expand(k.shape)
-        k_sin = sin[position_ids].unsqueeze(1).expand(k.shape)
-        k_embed = (k * k_cos) + (rotate_half(k) * k_sin)
-    else:
-        cos = cos[position_ids].unsqueeze(1)
-        sin = sin[position_ids].unsqueeze(1)
-        q_embed = (q * cos) + (rotate_half(q) * sin)
-        k_embed = (k * cos) + (rotate_half(k) * sin)
+    cos = cos[position_ids].unsqueeze(1)
+    sin = sin[position_ids].unsqueeze(1)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
     return q_embed, k_embed