Update modeling_Llamoe.py

modeling_Llamoe.py

@@ -562,13 +562,10 @@ class LlamoeSdpaAttention(LlamoeAttention):
         bsz, q_len, _ = hidden_states.size()
 
 
-        print("hidden_states:",hidden_states.shape)
         query_states = self.q_proj(hidden_states)
         key_states = self.k_proj(hidden_states)
         value_states = self.v_proj(hidden_states)
-        print("query_states:",query_states.shape)
-        print("key_states:",key_states.shape)
-        print("value_states:",value_states.shape)
+  
         query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
         key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
         value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
@@ -585,15 +582,12 @@ class LlamoeSdpaAttention(LlamoeAttention):
 
         key_states = repeat_kv(key_states, self.num_key_value_groups)
         value_states = repeat_kv(value_states, self.num_key_value_groups)
-        print("after_rb_key_states:",key_states)
-        print("after_rb_value_states:",value_states)
+
 
         causal_mask = attention_mask
-        print("causal_mask:",causal_mask)
         if attention_mask is not None and cache_position is not None:
             causal_mask = causal_mask[:, :, cache_position, : key_states.shape[-2]]
 
-        print("after_causal_masks:",causal_mask)
         # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
         # Reference: https://github.com/pytorch/pytorch/issues/112577.
         if query_states.device.type == "cuda" and causal_mask is not None:
@@ -605,7 +599,7 @@ class LlamoeSdpaAttention(LlamoeAttention):
             query_states,
             key_states,
             value_states,
-            attn_mask=causal_mask,
+            attn_mask=None,
             dropout_p=self.attention_dropout if self.training else 0.0,
         )