add support for flash attn 2

modeling_qwen.py

@@ -36,10 +36,6 @@ SUPPORT_CUDA = torch.cuda.is_available()
 SUPPORT_BF16 = SUPPORT_CUDA and torch.cuda.is_bf16_supported()
 SUPPORT_FP16 = SUPPORT_CUDA and torch.cuda.get_device_capability(0)[0] >= 7
 
-apply_rotary_emb_func = None
-rms_norm = None
-flash_attn_unpadded_func = None
-
 from .configuration_qwen import QWenConfig
 from .qwen_generation_utils import (
     HistoryType,
@@ -57,6 +53,45 @@ _CONFIG_FOR_DOC = "QWenConfig"
 
 QWen_PRETRAINED_MODEL_ARCHIVE_LIST = ["qwen-7b"]
 
+apply_rotary_emb_func = None
+rms_norm = None
+flash_attn_unpadded_func = None
+
+
+def _import_flash_attn():
+    global apply_rotary_emb_func, rms_norm, flash_attn_unpadded_func
+    try:
+        from flash_attn.layers.rotary import apply_rotary_emb_func as __apply_rotary_emb_func
+        apply_rotary_emb_func = __apply_rotary_emb_func
+    except ImportError:
+        logger.warn(
+            "Warning: import flash_attn rotary fail, please install FlashAttention rotary to get higher efficiency "
+            "https://github.com/Dao-AILab/flash-attention/tree/main/csrc/rotary"
+        )
+
+    try:
+        from flash_attn.ops.rms_norm import rms_norm as __rms_norm
+        rms_norm = __rms_norm
+    except ImportError:
+        logger.warn(
+            "Warning: import flash_attn rms_norm fail, please install FlashAttention layer_norm to get higher efficiency "
+            "https://github.com/Dao-AILab/flash-attention/tree/main/csrc/layer_norm"
+        )
+
+    try:
+        import flash_attn
+        if int(flash_attn.__version__.split(".")[0]) >= 2:
+            from flash_attn.flash_attn_interface import flash_attn_varlen_func as __flash_attn_unpadded_func
+        else:
+            from flash_attn.flash_attn_interface import flash_attn_unpadded_func as __flash_attn_unpadded_func
+        flash_attn_unpadded_func = __flash_attn_unpadded_func
+    except ImportError:
+        logger.warn(
+            "Warning: import flash_attn fail, please install FlashAttention to get higher efficiency "
+            "https://github.com/Dao-AILab/flash-attention"
+        )
+
+
 class FlashSelfAttention(torch.nn.Module):
     def __init__(
         self,
@@ -794,33 +829,7 @@ class QWenLMHeadModel(QWenPreTrainedModel):
             logger.warn("Flash attention will be disabled because it does NOT support fp32.")
 
         if config.use_flash_attn:
-            global apply_rotary_emb_func, rms_norm, flash_attn_unpadded_func
-            try:
-                from flash_attn.layers.rotary import apply_rotary_emb_func as __apply_rotary_emb_func
-                apply_rotary_emb_func = __apply_rotary_emb_func
-            except ImportError:
-                logger.warn(
-                    "Warning: import flash_attn rotary fail, please install FlashAttention rotary to get higher efficiency "
-                    "https://github.com/Dao-AILab/flash-attention/tree/main/csrc/rotary"
-                )
-
-            try:
-                from flash_attn.ops.rms_norm import rms_norm as __rms_norm
-                rms_norm = __rms_norm
-            except ImportError:
-                logger.warn(
-                    "Warning: import flash_attn rms_norm fail, please install FlashAttention layer_norm to get higher efficiency "
-                    "https://github.com/Dao-AILab/flash-attention/tree/main/csrc/layer_norm"
-                )
-
-            try:
-                from flash_attn.flash_attn_interface import flash_attn_unpadded_func as __flash_attn_unpadded_func
-                flash_attn_unpadded_func = __flash_attn_unpadded_func
-            except ImportError:
-                logger.warn(
-                    "Warning: import flash_attn fail, please install FlashAttention to get higher efficiency "
-                    "https://github.com/Dao-AILab/flash-attention"
-                )
+            _import_flash_attn()
 
         self.transformer = QWenModel(config)
         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)