Merge branch 'lmcafee/flash-attn-fix' into 'main'

Test NVIDIA#2: Memory, timing See merge request ADLR/megatron-lm!677
TrustLLMeu · Jul 19, 2023 · b030472 · b030472
2 parents 02c6229 + 1a03e5d
commit b030472
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
@@ -27,8 +27,10 @@
 try:
     from flash_attn.flash_attn_interface import flash_attn_unpadded_func
 except ImportError:
-    flash_attn_unpadded_func = None
-
+    try:
+        from flash_attn.flash_attn_interface import flash_attn_varlen_func as flash_attn_unpadded_func
+    except ImportError:
+        flash_attn_unpadded_func = None
 
 """ We use the following notation throughout this file:
      h: hidden size