Make flash attention configurable in user code

README.md

@@ -105,7 +105,7 @@ from transformers import AutoProcessor
 
 model_id = "microsoft/Phi-3-vision-128k-instruct" 
 
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", trust_remote_code=True, torch_dtype="auto")
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", trust_remote_code=True, torch_dtype="auto", _attn_implementation='flash_attention_2') # use _attn_implementation='eager' to disable flash attention
 
 processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) 
 
@@ -217,24 +217,6 @@ Note that by default, the Phi-3-Vision-128K model uses flash attention, which re
 * NVIDIA A6000
 * NVIDIA H100
 
-### Running on Windows or without flash attention
-To enable the model on these enviroment here are steps that you may consider to follow:
-
-Step 1: comment flash attention import code in modeling_phi3_v.py from line 52 to line 56.
-```python
-# if is_flash_attn_2_available():
-#     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)
-```
-
-Step 2: change _"_attn_implementation"_ from _"flash_attention_2"_ to _"eager"_ in config.json or disable flash attention when you create the model as below.
-
-```python
-model = AutoModelForCausalLM.from_pretrained('microsoft/Phi-3-vision-128k-instruct', device_map="cuda", trust_remote_code=True, torch_dtype="auto", _attn_implementation="eager")
-```
-
 ## License
 
 The model is licensed under the [MIT license](https://huggingface.co/microsoft/Phi-3-vision-128k-instruct/resolve/main/LICENSE).

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,
-        )
+        )