Fixes flash-attn import with a try/except statement

configuration_mixformer_sequential.py

@@ -30,6 +30,9 @@ class MixFormerSequentialConfig(PretrainedConfig):
         n_head_kv: Optional[int] = None,
         rotary_dim: Optional[int] = 32,
         activation_function: Optional[str] = "gelu_new",
+        flash_rotary: bool = False,
+        fused_dense: bool = False,
+        attn_pdrop: Optional[float] = 0.0,
         embd_pdrop: Optional[float] = 0.0,
         resid_pdrop: Optional[float] = 0.0,
         layer_norm_epsilon: Optional[float] = 1e-5,
@@ -47,6 +50,9 @@ class MixFormerSequentialConfig(PretrainedConfig):
         self.n_head_kv = n_head_kv
         self.rotary_dim = min(rotary_dim, n_embd // n_head)
         self.activation_function = activation_function
+        self.flash_rotary = flash_rotary
+        self.fused_dense = fused_dense
+        self.attn_pdrop = attn_pdrop
         self.embd_pdrop = embd_pdrop
         self.resid_pdrop = resid_pdrop
         self.layer_norm_epsilon = layer_norm_epsilon

modeling_mixformer_sequential.py

@@ -32,7 +32,6 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 from __future__ import annotations
-import importlib
 
 import math
 from typing import Any, Dict, Optional, Tuple, Union
@@ -49,14 +48,10 @@ from transformers.modeling_outputs import CausalLMOutputWithPast
 from .configuration_mixformer_sequential import MixFormerSequentialConfig
 
 
-def _is_flash_attn_available() -> bool:
-    return importlib.util.find_spec("flash_attn") is not None
-
-
-if _is_flash_attn_available():
+try:
     from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
     from flash_attn.ops.fused_dense import FusedDense
-else:
+except:
     FlashRotaryEmbedding = None
     FusedDense = None
 
@@ -549,9 +544,6 @@ class MHA(nn.Module):
         bias: bool = True,
         causal: bool = True,
         softmax_scale: Optional[float] = None,
-        dropout: float = 0.0,
-        flash_rotary: bool = True,
-        fused_dense: bool = True,
         layer_idx: Optional[int] = None,
         return_residual: bool = False,
         checkpointing: bool = False,
@@ -565,7 +557,7 @@ class MHA(nn.Module):
             if rotary_emb_scale_base is not None and rotary_emb_scale_base > 0.0:
                 rotary_kwargs["scale_base"] = rotary_emb_scale_base
             
-            rotary_cls = FlashRotaryEmbedding if flash_rotary else RotaryEmbedding
+            rotary_cls = FlashRotaryEmbedding if config.flash_rotary else RotaryEmbedding
             if rotary_cls is None:
                 rotary_cls = RotaryEmbedding
             self.rotary_emb = rotary_cls(self.rotary_emb_dim, **rotary_kwargs)
@@ -575,7 +567,7 @@ class MHA(nn.Module):
         op_size = self.head_dim * (self.n_head + 2 * self.n_head_kv)
         hidden_size = config.n_embd
 
-        linear_cls = FusedDense if fused_dense else nn.Linear
+        linear_cls = FusedDense if config.fused_dense else nn.Linear
         if linear_cls is None:
             linear_cls = nn.Linear
 
@@ -583,8 +575,8 @@ class MHA(nn.Module):
         self.out_proj = linear_cls(hidden_size, hidden_size, bias=bias, device=device, dtype=dtype)
 
         # Attention
-        self.inner_attn = SelfAttention(causal=causal, softmax_scale=softmax_scale, attention_dropout=dropout)
-        self.inner_cross_attn = CrossAttention(causal=causal, softmax_scale=softmax_scale, attention_dropout=dropout)
+        self.inner_attn = SelfAttention(causal=causal, softmax_scale=softmax_scale, attention_dropout=config.attn_pdrop)
+        self.inner_cross_attn = CrossAttention(causal=causal, softmax_scale=softmax_scale, attention_dropout=config.attn_pdrop)
 
         self.layer_idx = layer_idx
         self.return_residual = return_residual