v1

meteor/arch/modeling_internlm2.py

@@ -277,6 +277,9 @@ 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, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors."""
+    print(f'------------------------------cos: {cos.device}------------------------------')
+    print(f'------------------------------position_ids: {position_ids.device}------------------------------')
+    print(f'------------------------------unsqueeze_dim: {unsqueeze_dim.device}------------------------------')
     cos = cos[position_ids].unsqueeze(unsqueeze_dim)
     sin = sin[position_ids].unsqueeze(unsqueeze_dim)
     q_embed = (q * cos) + (rotate_half(q) * sin)