[FA] Fix FA tests for unique models

src/transformers/models/bark/modeling_bark.py

@@ -189,16 +189,6 @@ def __init__(self, *args, **kwargs):
         # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
         self._flash_attn_uses_top_left_mask = flash_attn_supports_top_left_mask()
 
-    def _split_heads(self, tensor, num_heads, attn_head_size):
-        """
-        Splits hidden_size dim into attn_head_size and num_heads
-        """
-        new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
-        tensor = tensor.view(new_shape)
-        # Flash attention requires the input to have the shape
-        # batch_size x seq_length x head_dim x hidden_dim - (batch, seq_length, head, head_features)
-        return tensor
-
     def _merge_heads(self, tensor, num_heads, attn_head_size):
         """
         Merges attn_head_size dim and num_attn_heads dim into hidden_size
@@ -232,9 +222,9 @@ def forward(
         target_dtype = get_target_dtype(query, self)  # if the query is in float32, this is the dtype to cast to for FA
 
         attn_output = _flash_attention_forward(
-            query,
-            key,
-            value,
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
             attention_mask,
             query_len,
             dropout=self.dropout if self.training else 0.0,

tests/models/bamba/test_modeling_bamba.py

@@ -67,7 +67,7 @@ def __init__(
         use_input_mask=True,
         use_labels=True,
         vocab_size=99,
-        hidden_size=32,
+        hidden_size=64,
         num_hidden_layers=2,
         num_attention_heads=4,
         num_key_value_heads=2,

tests/models/clipseg/test_modeling_clipseg.py

@@ -259,6 +259,7 @@ def prepare_config_and_inputs(self):
     def get_config(self):
         return CLIPSegTextConfig(
             vocab_size=self.vocab_size,
+            eos_token_id=self.vocab_size - 1,
             hidden_size=self.hidden_size,
             num_hidden_layers=self.num_hidden_layers,
             num_attention_heads=self.num_attention_heads,
@@ -399,6 +400,12 @@ def create_and_check_model_for_image_segmentation(self, config, input_ids, atten
     def prepare_config_and_inputs_for_common(self):
         config_and_inputs = self.prepare_config_and_inputs()
         config, input_ids, attention_mask, pixel_values = config_and_inputs
+
+        # Ensure eos pooling uses an unmasked token
+        eos_id = config.text_config.eos_token_id
+        input_ids[input_ids == eos_id] = 3
+        input_ids[:, -2] = eos_id
+
         inputs_dict = {
             "input_ids": input_ids,
             "attention_mask": attention_mask,

tests/models/efficientloftr/test_modeling_efficientloftr.py

@@ -49,15 +49,15 @@ def __init__(
         image_width=6,  # need to be a multiple of `stage_stride[0] * stage_stride[1]`
         image_height=4,  # need to be a multiple of `stage_stride[0] * stage_stride[1]`
         stage_num_blocks: list[int] = [1, 1],
-        out_features: list[int] = [16, 16],  # need to be >= 2 to make `config.fine_fusion_dims > 0`
+        out_features: list[int] = [16, 128],  # need to be >= 2 to make `config.fine_fusion_dims > 0`
         stage_stride: list[int] = [2, 1],
         q_aggregation_kernel_size: int = 1,
         kv_aggregation_kernel_size: int = 1,
         q_aggregation_stride: int = 1,
         kv_aggregation_stride: int = 1,
         num_attention_layers: int = 2,
         num_attention_heads: int = 8,
-        hidden_size: int = 16,
+        hidden_size: int = 128,
         coarse_matching_threshold: float = 0.0,
         fine_kernel_size: int = 2,
         coarse_matching_border_removal: int = 0,

tests/models/falcon/test_modeling_falcon.py

@@ -48,6 +48,8 @@ class FalconModelTester(CausalLMModelTester):
     def __init__(self, parent, new_decoder_architecture=True):
         super().__init__(parent)
         self.new_decoder_architecture = new_decoder_architecture
+        self.hidden_size = 64
+        self.num_attention_heads = 4
 
 
 @require_torch
@@ -67,6 +69,12 @@ def is_pipeline_test_to_skip(
     ):
         return True
 
+    @staticmethod
+    def _prepare_config_headdim(config, requested_dim):
+        # `head_dim` is a property only in falcon, we adjusted the initial values instead
+        assert config.head_dim >= requested_dim
+        return config
+
     @unittest.skip(reason="To support alibi, we are forced to create a mask in all cases")
     def test_sdpa_can_dispatch_on_flash(self):
         pass

tests/models/higgs_audio_v2_tokenizer/test_modeling_higgs_audio_v2_tokenizer.py

@@ -288,6 +288,9 @@ def test_initialization(self):
     @is_flaky()
     def test_flash_attn_2_inference_equivalence(self):
         for model_class in self.all_model_classes:
+            if not model_class._supports_flash_attn:
+                self.skipTest(f"{model_class.__name__} does not support Flash Attention 2")
+
             config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
             model = model_class(config)
 

tests/models/idefics2/test_modeling_idefics2.py

@@ -79,7 +79,7 @@ def __init__(
             "resampler_depth": 2,
             "resampler_n_heads": 2,
             "num_key_value_heads": 1,
-            "resampler_head_dim": 12,
+            "resampler_head_dim": 16,
             "attention_dropout": 0.0,
         },
         text_config={

tests/models/mimi/test_modeling_mimi.py

@@ -20,18 +20,14 @@
 import numpy as np
 import pytest
 from datasets import Audio, load_dataset
-from pytest import mark
 
 from tests.utils.test_audio_utils import normalize_waveform
 from transformers import AutoFeatureExtractor, MimiConfig, MimiDecoderConfig, MimiEncoderConfig, set_seed
 from transformers.audio_utils import load_audio
 from transformers.testing_utils import (
     Expectations,
-    is_flaky,
     is_torch_available,
-    require_flash_attn,
     require_torch,
-    require_torch_accelerator,
     slow,
     torch_device,
 )
@@ -297,12 +293,10 @@ def test_identity_shortcut(self):
         config.use_conv_shortcut = False
         self.model_tester.create_and_check_model_forward(config, inputs_dict)
 
-    @require_flash_attn
-    @require_torch_accelerator
-    @mark.flash_attn_test
-    @slow
-    @is_flaky()
-    def test_flash_attn_2_inference_equivalence(self):
+    def flash_attn_inference_equivalence(self, attn_implementation, padding_side="left"):
+        if padding_side == "right":
+            self.skipTest("The MimiModel does not support right padding")
+
         for model_class in self.all_model_classes:
             # Set seed for deterministic test - ensures reproducible model initialization and inputs
             set_seed(42)
@@ -312,7 +306,7 @@ def test_flash_attn_2_inference_equivalence(self):
             with tempfile.TemporaryDirectory() as tmpdirname:
                 model.save_pretrained(tmpdirname)
                 model_fa = model_class.from_pretrained(
-                    tmpdirname, dtype=torch.bfloat16, attn_implementation="flash_attention_2"
+                    tmpdirname, dtype=torch.bfloat16, attn_implementation=attn_implementation
                 )
                 model_fa.to(torch_device)
 
@@ -331,10 +325,6 @@ def test_flash_attn_2_inference_equivalence(self):
 
                 assert torch.allclose(logits_fa, logits, atol=4e-2, rtol=4e-2)
 
-    @unittest.skip(reason="The MimiModel does not support right padding")
-    def test_flash_attn_2_inference_equivalence_right_padding(self):
-        pass
-
     @unittest.skip(reason="The MimiModel does not have support dynamic compile yet")
     @pytest.mark.torch_compile_test
     def test_sdpa_can_compile_dynamic(self):

tests/models/nemotron3_5_transcribe/test_modeling_nemotron3_5_transcribe.py

@@ -17,6 +17,7 @@
 import tempfile
 import unittest
 from pathlib import Path
+from unittest.mock import patch
 
 from transformers import (
     AutoModelForMultimodalLM,
@@ -66,8 +67,8 @@ def __init__(self, parent, **kwargs):
         kwargs.setdefault("head_dim", 16)
         kwargs.setdefault("moe_intermediate_size", 16)
         kwargs.setdefault("moe_shared_expert_intermediate_size", 16)
-        kwargs.setdefault("mamba_num_heads", 4)
-        kwargs.setdefault("mamba_n_groups", 1)
+        kwargs.setdefault("mamba_num_heads", 8)
+        kwargs.setdefault("mamba_n_groups", 8)
         kwargs.setdefault("mamba_head_dim", 16)
         kwargs.setdefault("ssm_state_size", 8)
         kwargs.setdefault("mamba_chunk_size", 16)
@@ -81,13 +82,13 @@ def __init__(self, parent, **kwargs):
         kwargs.setdefault("audio_asr_num_hidden_layers", 2)
         kwargs.setdefault("audio_asr_num_attention_heads", 2)
 
-        kwargs.setdefault("audio_diarization_audio_hidden_size", 16)
+        kwargs.setdefault("audio_diarization_audio_hidden_size", 32)
         kwargs.setdefault("audio_diarization_audio_intermediate_size", 32)
         kwargs.setdefault("audio_diarization_audio_num_hidden_layers", 1)
         kwargs.setdefault("audio_diarization_audio_num_attention_heads", 2)
 
-        kwargs.setdefault("audio_diarization_head_hidden_size", 8)
-        kwargs.setdefault("audio_hidden_size", 16)
+        kwargs.setdefault("audio_diarization_head_hidden_size", 16)
+        kwargs.setdefault("audio_hidden_size", 32)
 
         kwargs.setdefault("chunk_size", 16)
         kwargs.setdefault("chunk_length", 500)
@@ -297,6 +298,26 @@ def test_audio_token_past_the_vocabulary(self):
         prompt_length = input_ids.shape[1]
         torch.testing.assert_close(sequences[:, prompt_length:], expected_sequences[:, prompt_length:], rtol=0, atol=0)
 
+    def test_sdpa_can_dispatch_on_flash(self):
+        """Audio encoder can produce fp32 (attention) outputs => avoid by autocasting + a mask is required so dummy replace it"""
+
+        def full_ones_mask(batch_size: int | None = None):
+            batch_size = batch_size if batch_size is not None else self.model_tester.batch_size
+            return torch.ones(
+                [batch_size, self.model_tester.feat_seq_length],
+                dtype=torch.bool,
+                device=torch_device,
+            )
+
+        with patch.object(self.model_tester, "create_audio_mask", new=full_ones_mask):
+            with torch.autocast(device_type=torch_device, dtype=torch.float16):
+                super().test_sdpa_can_dispatch_on_flash()
+
+    def _test_attention_implementation(self, attn_implementation):
+        """Audio encoder can produce fp32 (attention) outputs => avoid by autocasting"""
+        with torch.autocast(device_type=torch_device, dtype=torch.float16):
+            super()._test_attention_implementation(attn_implementation)
+
 
 @require_torch
 @require_torch_accelerator_memory(memory=70)

tests/models/squeezebert/test_modeling_squeezebert.py

@@ -47,7 +47,7 @@ def __init__(
         use_token_type_ids=False,
         use_labels=True,
         vocab_size=99,
-        hidden_size=32,
+        hidden_size=64,
         num_hidden_layers=2,
         num_attention_heads=4,
         intermediate_size=64,