Repository: huggingface/transformers
Pull request: [`FlexAttn`] Fix models with unique characteristics

For context, flex attention cannot work with dimensions less than 16; hence, the config was manipulated to ensure the test works. Before, most models failed, including llama.

There are some models such as idefics 2+3, smolvlm which do not have the `_is_composite` flag and as I do not want to affect other tests - so, I added a new condition to skip the test. They may have passed before but it's not future-proof. For Zamba2, I overwrote the test since some other dims don't add up when changing `hidden_size`.

There are other options:
- Rewrite the test to handle subconfigs --> tried that but there are so many edge cases and weird configs that lead to some issues one way or another.
- Adjust the dimensions in all models and avoid the hidden dim manipulation in the first place. Not sure if this is good as it will strain the tests even more imo :eyes: 


Edit: #38434 took care of the composite models. This PR is left to fix some of the more unique models such as zamba2 and deepseek3.



Changed files:
filename: tests/models/deepseek_v3/test_modeling_deepseek_v3.py
patch: @@ -24,6 +24,7 @@
     require_read_token,
     require_torch,
     require_torch_accelerator,
+    require_torch_gpu,
     require_torch_large_accelerator,
     require_torch_sdpa,
     slow,
@@ -494,6 +495,35 @@ def test_eager_matches_sdpa_generate(self):
                     msg=f"\n{tokenizer.batch_decode(res_eager)} \nvs\n{tokenizer.batch_decode(res_sdpa)}",
                 )
 
+    @require_torch_gpu
+    def test_flex_attention_with_grads(self):
+        """
+        Overwriting as the namings/functionality on the attention part are different; for now it's more of a unique model.
+        Original issue is also due to dimensionalities, here specifically due to dims not being a multiple of 2.
+        """
+        for model_class in self.all_model_classes:
+            config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+            config._attn_implementation = "flex_attention"
+
+            # Disable dropout
+            config.attention_dropout = 0.0
+
+            # Deepseek 3 specific - manipulate nope and adjust calculated total head dim
+            config.qk_nope_head_dim = 16
+            config.qk_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
+
+            model = model_class(config).to(device=torch_device)
+            self.assertTrue(model.config._attn_implementation == "flex_attention")
+
+            # Elaborate workaround for encoder-decoder models as some do not specify their main input
+            dummy_inputs = {model.main_input_name: inputs_dict[model.main_input_name].to(torch_device)}
+            if config.is_encoder_decoder:
+                dummy_inputs["decoder_input_ids"] = inputs_dict["decoder_input_ids"].to(torch_device)
+                dummy_inputs["decoder_attention_mask"] = inputs_dict["decoder_attention_mask"].to(torch_device)
+
+            # If this does not raise an error, the test passes (see https://github.com/huggingface/transformers/pull/35605)
+            _ = model(**dummy_inputs)
+
 
 @require_torch_accelerator
 class DeepseekV3IntegrationTest(unittest.TestCase):

filename: tests/models/zamba2/test_modeling_zamba2.py
patch: @@ -578,6 +578,28 @@ def test_flash_attn_2_inference_equivalence_right_padding(self):
     def test_new_cache_format(self, num_beams, do_sample):
         pass
 
+    @require_torch_gpu
+    def test_flex_attention_with_grads(self):
+        """
+        Overwriting as the base hidden size is big enough for compile.
+        Manipulation of dims causes issues due to other constraints not being satisfied anymore.
+        """
+        for model_class in self.all_model_classes:
+            config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+            config._attn_implementation = "flex_attention"
+
+            model = model_class(config).to(device=torch_device)
+            self.assertTrue(model.config._attn_implementation == "flex_attention")
+
+            # Elaborate workaround for encoder-decoder models as some do not specify their main input
+            dummy_inputs = {model.main_input_name: inputs_dict[model.main_input_name].to(torch_device)}
+            if config.is_encoder_decoder:
+                dummy_inputs["decoder_input_ids"] = inputs_dict["decoder_input_ids"].to(torch_device)
+                dummy_inputs["decoder_attention_mask"] = inputs_dict["decoder_attention_mask"].to(torch_device)
+
+            # If this does not raise an error, the test passes (see https://github.com/huggingface/transformers/pull/35605)
+            _ = model(**dummy_inputs)
+
 
 @require_torch
 class Zamba2ModelIntegrationTest(unittest.TestCase):