Use input attention mask instead of casual mask in attention

modelling_RW.py

@@ -271,13 +271,14 @@ class Attention(nn.Module):
         else:
             present = None
 
+        attention_mask_float = (attention_mask * 1.0).masked_fill(attention_mask, -1e9).to(query_layer.dtype)
         if alibi is None:
             query_layer_ = query_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)
             key_layer_ = key_layer.reshape(batch_size, self.num_kv, -1, self.head_dim)
             value_layer_ = value_layer.reshape(batch_size, self.num_kv, -1, self.head_dim)
 
             attn_output = F.scaled_dot_product_attention(
-                query_layer_, key_layer_, value_layer_, None, 0.0, is_causal=True
+                query_layer_, key_layer_, value_layer_, attention_mask_float, 0.0, is_causal=False
             )
 
             x = attn_output.view(batch_size, self.num_heads, q_length, self.head_dim)
@@ -290,7 +291,6 @@ class Attention(nn.Module):
             assert not output_attentions  # not supported.
             return outputs
         else:
-            attention_mask_float = (attention_mask * 1.0).masked_fill(attention_mask, -1e9).to(torch.bfloat16)
             matmul_result = query_layer @ key_layer.transpose(-1, -2)
 
             # change view to [batch_size, num_heads, q_length, kv_length]