make flash attention usage configurable from user code

config.json

@@ -143,6 +143,5 @@
   "torch_dtype": "bfloat16",
   "transformers_version": "4.38.1",
   "use_cache": true,
-  "vocab_size": 32064,
-  "_attn_implementation": "flash_attention_2"
+  "vocab_size": 32064
 }

modeling_phi3_v.py

@@ -40,7 +40,6 @@ from transformers.utils import (
     add_code_sample_docstrings,
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
-    is_flash_attn_2_available,
     is_flash_attn_greater_or_equal_2_10,
     logging,
     replace_return_docstrings,
@@ -49,11 +48,13 @@ from .configuration_phi3_v import Phi3VConfig
 from .image_embedding_phi3_v import Phi3ImageEmbedding
 
 
-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
 
     _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
+except ImportError:
+    pass
 
 logger = logging.get_logger(__name__)
 
@@ -1000,8 +1001,8 @@ PHI3V_INPUTS_DOCSTRING = r"""
             is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
             model's internal embedding lookup matrix.
         pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)):
-            The tensors corresponding to the input images. Pixel values can be obtained using [`AutoImageProcessor`]. 
-            See [`Phi3ImageProcessor.__call__`] for details. 
+            The tensors corresponding to the input images. Pixel values can be obtained using [`AutoImageProcessor`].
+            See [`Phi3ImageProcessor.__call__`] for details.
         image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`, *optional*):
             The sizes of the images in the batch, being (height, width) for each image.
         use_cache (`bool`, *optional*):
@@ -1046,7 +1047,7 @@ class Phi3VModel(Phi3VPreTrainedModel):
                 **config.embd_layer
             }
             self.vision_embed_tokens = Phi3ImageEmbedding(config, wte=self.embed_tokens, **embedding_config)
-            # # set wte the same for vision embedding 
+            # # set wte the same for vision embedding
             # self.vision_embed_tokens.wte.weight = self.embed_tokens.weight
 
         self.layers = nn.ModuleList(
@@ -1629,4 +1630,4 @@ class Phi3VForTokenClassification(Phi3VPreTrainedModel):
             logits=logits,
             hidden_states=model_outputs.hidden_states,
             attentions=model_outputs.attentions,
-        )
+        )