fix: make `flash_attn` optional (`trust_remote_code` breaks dynamic module check)

README.md

@@ -59,7 +59,7 @@ model = AutoModelForCausalLM.from_pretrained(
   "stabilityai/stablelm-3b-4e1t",
   trust_remote_code=True,
   torch_dtype="auto",
-+ use_flash_attention_2=True,
+  attn_implementation="flash_attention_2",
 )
 model.cuda()
 inputs = tokenizer("The weather is always wonderful", return_tensors="pt").to(model.device)

modeling_stablelm_epoch.py

@@ -33,14 +33,16 @@ from transformers.modeling_outputs import (
     CausalLMOutputWithPast,
 )
 from transformers.modeling_utils import PreTrainedModel
-from transformers.utils import logging, is_flash_attn_2_available, is_flash_attn_greater_or_equal_2_10
+from transformers.utils import logging, is_flash_attn_greater_or_equal_2_10
 
 from .configuration_stablelm_epoch import StableLMEpochConfig
 
-
-if is_flash_attn_2_available():
+try:
     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
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
+except:
+    flash_attn_func, flash_attn_varlen_func = None, None
+    index_first_axis, pad_input, unpad_input = None, None, None
 
 
 logger = logging.get_logger(__name__)