bug fix in fast attention

VarunGumma · Aug 9, 2024 · ad64f6d · ad64f6d
1 parent 390bb36
commit ad64f6d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/fairseq/modules/fast_grouped_query_attention.py b/fairseq/modules/fast_grouped_query_attention.py
@@ -152,6 +152,8 @@ def forward(
         tgt_len, bsz, embed_dim = query.size()
         src_len = tgt_len
 
+        dropout_p = self.dropout_p if self.training else 0
+
         assert list(query.size()) == [tgt_len, bsz, embed_dim]
         if key is not None:
             src_len, key_bsz, _ = key.size()
@@ -378,7 +380,7 @@ def forward(
             value=v,
             is_causal=False,
             attn_mask=combined_mask,
-            dropout_p=self.dropout_p,
+            dropout_p=dropout_p,
         )
 
         attn = rearrange(

diff --git a/fairseq/modules/fast_multihead_attention.py b/fairseq/modules/fast_multihead_attention.py
@@ -146,6 +146,8 @@ def forward(
         tgt_len, bsz, embed_dim = query.size()
         src_len = tgt_len
 
+        dropout_p = self.dropout_p if self.training else 0
+
         assert list(query.size()) == [tgt_len, bsz, embed_dim]
         if key is not None:
             src_len, key_bsz, _ = key.size()
@@ -337,7 +339,7 @@ def forward(
             value=v,
             is_causal=False,
             attn_mask=combined_mask,
-            dropout_p=self.dropout_p,
+            dropout_p=dropout_p
         )
 
         attn = rearrange(