Test NVIDIA#2: Memory, timing

TrustLLMeu · Jul 19, 2023 · 1a03e5d · 1a03e5d
1 parent 040eac9
commit 1a03e5d
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