Attn implementation for composite models

src/transformers/modeling_utils.py

@@ -45,6 +45,7 @@
 from .dynamic_module_utils import custom_object_save
 from .generation import GenerationConfig, GenerationMixin
 from .integrations import PeftAdapterMixin, deepspeed_config, is_deepspeed_zero3_enabled
+from .models.auto.modeling_auto import MODEL_MAPPING
 from .pytorch_utils import (  # noqa: F401
     Conv1D,
     apply_chunking_to_forward,
@@ -1512,6 +1513,40 @@ def _autoset_attn_implementation(
             # If a config is passed with a preset attn_implementation, we skip the automatic dispatch and use the user-provided config, with hard checks that the requested attention implementation is available.
             requested_attn_implementation = config._attn_implementation_internal
 
+        # MultiModal LLM related hack: since they consist of two or might be more sub-models
+        # we have to check and dispatch SDPA to each sub-model, in case any of them support it.
+        # If one sub-model supports SDPA while other doesn't, an error will be raised following the
+        # typical SDPA-dispatch path.
+        # Same goes for the `_supports_cache_class` because we cannot know what flag LM has
+        # before knowing which class is that LM
+        if hasattr(config, "text_config") and hasattr(config, "vision_config"):
+            text_model_cls = MODEL_MAPPING.get(type(config.text_config), None)
+            if text_model_cls is not None:
+                config.text_config._attn_implementation = config._attn_implementation
+                cls._supports_sdpa = text_model_cls._supports_sdpa
+                cls._autoset_attn_implementation(
+                    config.text_config,
+                    use_flash_attention_2=use_flash_attention_2,
+                    torch_dtype=torch_dtype,
+                    device_map=device_map,
+                    check_device_map=check_device_map,
+                )
+
+            vision_model_cls = MODEL_MAPPING.get(type(config.vision_config), None)
+            if vision_model_cls is not None:
+                config.vision_config._attn_implementation = config._attn_implementation
+                cls._supports_sdpa = vision_model_cls._supports_sdpa
+                cls._autoset_attn_implementation(
+                    config.vision_config,
+                    use_flash_attention_2=use_flash_attention_2,
+                    torch_dtype=torch_dtype,
+                    device_map=device_map,
+                    check_device_map=check_device_map,
+                )
+
+            if vision_model_cls is not None and text_model_cls is not None:
+                cls._supports_sdpa = vision_model_cls._supports_sdpa or text_model_cls._supports_sdpa
+
         if use_flash_attention_2:
             logger.warning_once(
                 'The model was loaded with use_flash_attention_2=True, which is deprecated and may be removed in a future release. Please use `attn_implementation="flash_attention_2"` instead.'

src/transformers/models/blip_2/modeling_blip_2.py

@@ -1225,19 +1225,21 @@ class Blip2Model(Blip2PreTrainedModel):
     def __init__(self, config: Blip2Config):
         super().__init__(config)
 
-        self.vision_model = Blip2VisionModel(config.vision_config)
+        self.vision_model = Blip2VisionModel._from_config(
+            config.vision_config, attn_implementation=config.vision_config._attn_implementation
+        )
 
         self.query_tokens = nn.Parameter(torch.zeros(1, config.num_query_tokens, config.qformer_config.hidden_size))
         self.qformer = Blip2QFormerModel(config.qformer_config)
 
         self.language_projection = nn.Linear(config.qformer_config.hidden_size, config.text_config.hidden_size)
         if config.use_decoder_only_language_model:
             language_model = AutoModelForCausalLM.from_config(
-                config.text_config, attn_implementation=config._attn_implementation
+                config.text_config, attn_implementation=config.text_config._attn_implementation
             )
         else:
             language_model = AutoModelForSeq2SeqLM.from_config(
-                config.text_config, attn_implementation=config._attn_implementation
+                config.text_config, attn_implementation=config.text_config._attn_implementation
             )
 
         # Update _tied_weights_keys using the base model used.
@@ -1590,7 +1592,9 @@ class Blip2ForConditionalGeneration(Blip2PreTrainedModel):
     def __init__(self, config: Blip2Config):
         super().__init__(config)
 
-        self.vision_model = Blip2VisionModel(config.vision_config)
+        self.vision_model = Blip2VisionModel._from_config(
+            config.vision_config, attn_implementation=config.vision_config._attn_implementation
+        )
 
         self.query_tokens = nn.Parameter(torch.zeros(1, config.num_query_tokens, config.qformer_config.hidden_size))
         self.qformer = Blip2QFormerModel(config.qformer_config)

src/transformers/models/idefics2/modeling_idefics2.py

@@ -921,7 +921,9 @@ def __init__(self, config, layer_idx: int):
 
         self.input_latents_norm = Idefics2RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
         self.input_context_norm = Idefics2RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
-        self.self_attn = IDEFICS2_PERCEIVER_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx)
+        self.self_attn = IDEFICS2_PERCEIVER_ATTENTION_CLASSES[config.perceiver_config._attn_implementation](
+            config, layer_idx=layer_idx
+        )
         self.post_attention_layernorm = Idefics2RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
         self.mlp = Idefics2MLP(
             hidden_size=config.text_config.hidden_size,
@@ -1131,7 +1133,18 @@ def _autoset_attn_implementation(
             check_device_map=check_device_map,
             **kwargs,
         )
-        config.vision_config._attn_implementation = config._attn_implementation
+        # autoset-attn calls recursively all sub-configs (text-config, vision-config)
+        # and sets attn implementation if the config can be mapped bu auto-model
+        # Idefics2 vision config can't be mapped automcatically so we set it manually here
+        # We cant set vision attn same as general attn, because the general one can be sdpa if at
+        # least one sub-module (in this case LLM) supports sdpa
+        if hasattr(config, "vision_config"):
+            config.vision_config._attn_implementation = (
+                config._attn_implementation if config._attn_implementation != "sdpa" else "eager"
+            )
+            config.perceiver_config._attn_implementation = (
+                config._attn_implementation if config._attn_implementation != "sdpa" else "eager"
+            )
         return config
 
 

src/transformers/models/instructblip/modeling_instructblip.py

@@ -1281,7 +1281,9 @@ class InstructBlipForConditionalGeneration(InstructBlipPreTrainedModel):
     def __init__(self, config: InstructBlipConfig):
         super().__init__(config)
 
-        self.vision_model = InstructBlipVisionModel(config.vision_config)
+        self.vision_model = InstructBlipVisionModel(
+            config.vision_config, attn_implementation=config.vision_config._attn_implementation
+        )
 
         self.query_tokens = nn.Parameter(torch.zeros(1, config.num_query_tokens, config.qformer_config.hidden_size))
         self.qformer = InstructBlipQFormerModel(config.qformer_config)
@@ -1290,11 +1292,11 @@ def __init__(self, config: InstructBlipConfig):
 
         if config.use_decoder_only_language_model:
             language_model = AutoModelForCausalLM.from_config(
-                config.text_config, attn_implementation=config._attn_implementation
+                config.text_config, attn_implementation=config.text_config._attn_implementation
             )
         else:
             language_model = AutoModelForSeq2SeqLM.from_config(
-                config.text_config, attn_implementation=config._attn_implementation
+                config.text_config, attn_implementation=config.text_config._attn_implementation
             )
 
         if language_model._no_split_modules is not None:

src/transformers/models/instructblipvideo/modeling_instructblipvideo.py

@@ -1290,7 +1290,9 @@ class InstructBlipVideoForConditionalGeneration(InstructBlipVideoPreTrainedModel
     def __init__(self, config: InstructBlipVideoConfig):
         super().__init__(config)
 
-        self.vision_model = InstructBlipVideoVisionModel(config.vision_config)
+        self.vision_model = InstructBlipVideoVisionModel._from_config(
+            config.vision_config, attn_implementation=config.vision_config._attn_implementation
+        )
 
         self.query_tokens = nn.Parameter(torch.zeros(1, config.num_query_tokens, config.qformer_config.hidden_size))
         self.qformer = InstructBlipVideoQFormerModel(config.qformer_config)
@@ -1299,11 +1301,11 @@ def __init__(self, config: InstructBlipVideoConfig):
 
         if config.use_decoder_only_language_model:
             language_model = AutoModelForCausalLM.from_config(
-                config.text_config, attn_implementation=config._attn_implementation
+                config.text_config, attn_implementation=config.text_config._attn_implementation
             )
         else:
             language_model = AutoModelForSeq2SeqLM.from_config(
-                config.text_config, attn_implementation=config._attn_implementation
+                config.text_config, attn_implementation=config.text_config._attn_implementation
             )
 
         if language_model._no_split_modules is not None:

src/transformers/models/llava/modeling_llava.py

@@ -149,14 +149,6 @@ def _init_weights(self, module):
             if module.padding_idx is not None:
                 module.weight.data[module.padding_idx].zero_()
 
-    @property
-    def _supports_sdpa(self):
-        """
-        Retrieve language_model's attribute to check whether the model supports
-        SDPA or not.
-        """
-        return self.language_model._supports_sdpa
-
 
 LLAVA_INPUTS_DOCSTRING = r"""
     Args:
@@ -236,12 +228,14 @@ def _supports_sdpa(self):
 class LlavaForConditionalGeneration(LlavaPreTrainedModel):
     def __init__(self, config: LlavaConfig):
         super().__init__(config)
-        self.vision_tower = AutoModel.from_config(config.vision_config)
+        self.vision_tower = AutoModel.from_config(
+            config.vision_config, attn_implementation=config.text_config._attn_implementation
+        )
 
         self.multi_modal_projector = LlavaMultiModalProjector(config)
         self.vocab_size = config.text_config.vocab_size
         self.language_model = AutoModelForCausalLM.from_config(
-            config.text_config, attn_implementation=config._attn_implementation
+            config.text_config, attn_implementation=config.vision_config._attn_implementation
         )
         self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
         self.post_init()

src/transformers/models/llava_next/modeling_llava_next.py

@@ -255,14 +255,6 @@ def _init_weights(self, module):
             if module.padding_idx is not None:
                 module.weight.data[module.padding_idx].zero_()
 
-    @property
-    def _supports_sdpa(self):
-        """
-        Retrieve language_model's attribute to check whether the model supports
-        SDPA or not.
-        """
-        return self.language_model._supports_sdpa
-
 
 LLAVA_NEXT_INPUTS_DOCSTRING = r"""
     Args:
@@ -345,15 +337,17 @@ def _supports_sdpa(self):
 class LlavaNextForConditionalGeneration(LlavaNextPreTrainedModel):
     def __init__(self, config: LlavaNextConfig):
         super().__init__(config)
-        self.vision_tower = AutoModel.from_config(config.vision_config)
+        self.vision_tower = AutoModel.from_config(
+            config.vision_config, attn_implementation=config.vision_config._attn_implementation
+        )
 
         self.multi_modal_projector = LlavaNextMultiModalProjector(config)
         embed_std = 1 / math.sqrt(config.text_config.hidden_size)
         self.image_newline = nn.Parameter(torch.randn(config.text_config.hidden_size, dtype=self.dtype) * embed_std)
 
         self.vocab_size = config.text_config.vocab_size
         self.language_model = AutoModelForCausalLM.from_config(
-            config.text_config, attn_implementation=config._attn_implementation
+            config.text_config, attn_implementation=config.text_config._attn_implementation
         )
         self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
         self._padding_side = "left"  # set it to left by default, user can use setter to change padding_sides

src/transformers/models/llava_next_video/modeling_llava_next_video.py

@@ -295,14 +295,6 @@ def _init_weights(self, module):
             if module.padding_idx is not None:
                 module.weight.data[module.padding_idx].zero_()
 
-    @property
-    def _supports_sdpa(self):
-        """
-        Retrieve language_model's attribute to check whether the model supports
-        SDPA or not.
-        """
-        return self.language_model._supports_sdpa
-
 
 LLAVA_NEXT_VIDEO_INPUTS_DOCSTRING = r"""
     Args:
@@ -388,15 +380,17 @@ def __init__(
         config: LlavaNextVideoConfig,
     ):
         super().__init__(config)
-        self.vision_tower = AutoModel.from_config(config.vision_config)
-
+        self.vision_tower = AutoModel.from_config(
+            config.vision_config, attn_implementation=config.vision_config._attn_implementation
+        )
         self.multi_modal_projector = LlavaNextVideoMultiModalProjector(config)
+
         embed_std = 1 / math.sqrt(config.text_config.hidden_size)
         self.image_newline = nn.Parameter(torch.randn(config.text_config.hidden_size, dtype=self.dtype) * embed_std)
 
         self.vocab_size = config.text_config.vocab_size
         self.language_model = AutoModelForCausalLM.from_config(
-            config.text_config, attn_implementation=config._attn_implementation
+            config.text_config, attn_implementation=config.text_config._attn_implementation
         )
         self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
         self._padding_side = "left"  # set it to left by default, user can use setter to change padding_sides

src/transformers/models/paligemma/modeling_paligemma.py

@@ -126,7 +126,6 @@ class PaliGemmaPreTrainedModel(PreTrainedModel):
     _no_split_modules = ["PaliGemmaMultiModalProjector"]
     _skip_keys_device_placement = "past_key_values"
     _supports_flash_attn_2 = False
-    _supports_sdpa = True
 
     def _init_weights(self, module):
         # important: this ported version of PaliGemmaisn't meant for training from scratch - only
@@ -149,14 +148,6 @@ def _init_weights(self, module):
             if module.padding_idx is not None:
                 module.weight.data[module.padding_idx].zero_()
 
-    @property
-    def _supports_sdpa(self):
-        """
-        Retrieve language_model's attribute to check whether the model supports
-        SDPA or not.
-        """
-        return self.language_model._supports_sdpa
-
 
 PALIGEMMA_INPUTS_DOCSTRING = r"""
     Args:
@@ -231,13 +222,15 @@ def _supports_sdpa(self):
 class PaliGemmaForConditionalGeneration(PaliGemmaPreTrainedModel):
     def __init__(self, config: PaliGemmaConfig):
         super().__init__(config)
-        self.vision_tower = AutoModel.from_config(config=config.vision_config)
+        self.vision_tower = AutoModel.from_config(
+            config=config.vision_config, attn_implementation=config.vision_config._attn_implementation
+        )
         self.multi_modal_projector = PaliGemmaMultiModalProjector(config)
         self.vocab_size = config.text_config.vocab_size
         self._attn_implementation = config._attn_implementation
 
         language_model = AutoModelForCausalLM.from_config(
-            config=config.text_config, attn_implementation=self._attn_implementation
+            config=config.text_config, attn_implementation=config.text_config._attn_implementation
         )
 
         if language_model._tied_weights_keys is not None:

src/transformers/models/video_llava/modeling_video_llava.py

@@ -237,13 +237,17 @@ def _supports_sdpa(self):
 class VideoLlavaForConditionalGeneration(VideoLlavaPreTrainedModel):
     def __init__(self, config: VideoLlavaConfig):
         super().__init__(config)
-        self.video_tower = AutoModel.from_config(config.vision_config)
-        self.image_tower = AutoModel.from_config(config.vision_config)
+        self.video_tower = AutoModel.from_config(
+            config.vision_config, attn_implementation=config.vision_config._attn_implementation
+        )
+        self.image_tower = AutoModel.from_config(
+            config.vision_config, attn_implementation=config.vision_config._attn_implementation
+        )
 
         self.multi_modal_projector = VideoLlavaMultiModalProjector(config)
         self.vocab_size = config.text_config.vocab_size
         self.language_model = AutoModelForCausalLM.from_config(
-            config.text_config, attn_implementation=config._attn_implementation
+            config.text_config, attn_implementation=config.text_config._attn_implementation
         )
         self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
         self.post_init()

src/transformers/models/vipllava/modeling_vipllava.py

@@ -158,14 +158,6 @@ def _init_weights(self, module):
             if module.padding_idx is not None:
                 module.weight.data[module.padding_idx].zero_()
 
-    @property
-    def _supports_sdpa(self):
-        """
-        Retrieve language_model's attribute to check whether the model supports
-        SDPA or not.
-        """
-        return self.language_model._supports_sdpa
-
 
 VIPLLAVA_INPUTS_DOCSTRING = r"""
     Args:
@@ -241,12 +233,14 @@ def _supports_sdpa(self):
 class VipLlavaForConditionalGeneration(VipLlavaPreTrainedModel):
     def __init__(self, config: VipLlavaConfig):
         super().__init__(config)
-        self.vision_tower = AutoModel.from_config(config.vision_config)
+        self.vision_tower = AutoModel.from_config(
+            config.vision_config, attn_implementation=config.text_config._attn_implementation
+        )
 
         self.multi_modal_projector = VipLlavaMultiModalProjector(config)
         self.vocab_size = config.text_config.vocab_size
         self.language_model = AutoModelForCausalLM.from_config(
-            config.text_config, attn_implementation=config._attn_implementation
+            config.text_config, attn_implementation=config.vision_config._attn_implementation
         )
         self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
         self.post_init()