fix alibi with flash attention

lucidrains · May 15, 2023 · 4c3ec34 · 4c3ec34
1 parent bfaab9a
commit 4c3ec34
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/MEGABYTE_pytorch/attend.py b/MEGABYTE_pytorch/attend.py
@@ -101,6 +101,7 @@ def flash_attn(self, q, k, v, mask = None, attn_bias = None):
             if exists(mask):
                 attn_bias = attn_bias.masked_fill(~mask, mask_value)
 
+            mask = attn_bias
             causal = False
 
         # pytorch 2.0 flash attn: q, k, v, mask, dropout, causal, softmax_scale

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'MEGABYTE-pytorch',
   packages = find_packages(),
-  version = '0.0.4',
+  version = '0.0.5',
   license='MIT',
   description = 'MEGABYTE - Pytorch',
   long_description_content_type = 'text/markdown',