feat: check flash-attn version if installed

modeling_yi.py

@@ -29,6 +29,11 @@ try:
 except Exception:
     is_flash_attn_available = False
 
+if is_flash_attn_available:
+    from flash_attn import __version__
+
+    assert __version__ >= "2.3.0", "please update your flash_attn version"
+
 logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "YiConfig"
@@ -539,7 +544,9 @@ class YiModel(YiPreTrainedModel):
     def _prepare_decoder_attention_mask(
         self, attention_mask, input_ids, inputs_embeds, past_key_values_length
     ):
-        input_shape = input_ids.shape if input_ids is not None else inputs_embeds.shape[:-1]
+        input_shape = (
+            input_ids.shape if input_ids is not None else inputs_embeds.shape[:-1]
+        )
         # create causal mask
         # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
         combined_attention_mask = None