Switch import mechanism for flash_attn

modeling_phimoe.py

@@ -50,14 +50,15 @@ from transformers.utils.import_utils import is_torch_fx_available
 from .configuration_phimoe import PhiMoEConfig
 
 from einops import rearrange
-from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
 
-
-if is_flash_attn_2_available():
+try:
+    from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
     from flash_attn import flash_attn_func, flash_attn_varlen_func
     from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
 
     _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
+except ImportError:
+    pass
 
 # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
 # It means that the function will not be traced through and simply appear as a node in the graph.