make flash-attn optional

modeling_yi.py

@@ -6,7 +6,6 @@ import torch.utils.checkpoint
 from einops import repeat
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
 from transformers.activations import ACT2FN
 from transformers.modeling_outputs import (
     BaseModelOutputWithPast,
@@ -18,17 +17,17 @@ from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
 from transformers.utils import (
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
-    is_flash_attn_available,
     logging,
     replace_return_docstrings,
 )
 
 from .configuration_yi import YiConfig
 
-
-if is_flash_attn_available():
+is_flash_attn_available = True
+try:
     from flash_attn import flash_attn_func
-
+except Exception:
+    is_flash_attn_available = False
 
 logger = logging.get_logger(__name__)
 
@@ -224,7 +223,6 @@ class YiAttention(nn.Module):
         use_cache: bool = False,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         bsz, q_len, _ = hidden_states.size()
-        flash_attn_available = is_flash_attn_available()
 
         query_states = self.q_proj(hidden_states).view(
             bsz, q_len, self.num_heads, self.head_dim
@@ -237,7 +235,7 @@ class YiAttention(nn.Module):
             bsz, q_len, self.num_key_value_heads, self.head_dim
         )
 
-        if not flash_attn_available:
+        if not is_flash_attn_available:
             if self.num_key_value_groups > 1:
                 key_states = repeat(
                     key_states, f"b n h d -> b n (h {self.num_key_value_groups}) d"
@@ -251,13 +249,13 @@ class YiAttention(nn.Module):
             key_states = key_states.transpose(1, 2)
             value_states = value_states.transpose(1, 2)
 
-        seq_dim = 1 if flash_attn_available else 2
+        seq_dim = 1 if is_flash_attn_available else 2
         kv_seq_len = key_states.shape[seq_dim]
         if past_key_value is not None:
             kv_seq_len += past_key_value[0].shape[seq_dim]
         cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
         query_states, key_states = apply_rotary_pos_emb(
-            query_states, key_states, cos, sin, position_ids, flash_attn_available
+            query_states, key_states, cos, sin, position_ids, is_flash_attn_available
         )
 
         if past_key_value is not None:
@@ -267,7 +265,7 @@ class YiAttention(nn.Module):
 
         past_key_value = (key_states, value_states) if use_cache else None
 
-        if flash_attn_available:
+        if is_flash_attn_available:
             attn_output = flash_attn_func(
                 query_states, key_states, value_states, dropout_p=0.0, causal=True
             )
@@ -308,7 +306,7 @@ class YiAttention(nn.Module):
                     f" {attn_output.size()}"
                 )
 
-        if not flash_attn_available:
+        if not is_flash_attn_available:
             attn_output = attn_output.transpose(1, 2)
 
         attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
@@ -631,7 +629,7 @@ class YiModel(YiPreTrainedModel):
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
 
-        if not is_flash_attn_available():
+        if not is_flash_attn_available:
             # embed positions
             if attention_mask is None:
                 attention_mask = torch.ones(