Repository: huggingface/transformers
Pull request: Gemma3 fixes

This PR fixes three things in `gemma3`:
- a multiple-device error where `torch.where` takes some of its coefficients from a tensor that is not on the right device and is a full_like, so we just replace it with the filling element
- an error in the `flash_attn_inference_equivalence` which is due to the model needing more parameters than are generated by defualt. To avoid this, we add a flag that specifies if we need to check the forward pass with training or not, and make this check default for both and left padding (cc. @vasqu )
- the test `flash_attn_from_config` was failing for the same reasons (`token_type_ids` is required as a model input when training) so I added a `.eval()` to avoid this. It does not seem the model needs to be in train mode for this test, but I can also add an option to the test to only call `.eval()` if a flag is passed

Changed files:
filename: src/transformers/models/gemma3/modeling_gemma3.py
patch: @@ -798,7 +798,7 @@ def create_causal_mask_mapping(
         is_previous_image = nn.functional.pad(is_image, (1, 0), value=0)[:, :-1]
         new_image_start = is_image & ~is_previous_image
         image_group_ids = torch.cumsum(new_image_start.int(), dim=1) - 1
-        image_group_ids = torch.where(is_image, image_group_ids, torch.full_like(token_type_ids, -1))
+        image_group_ids = torch.where(is_image, image_group_ids, -1)
         mask_kwargs["or_mask_function"] = token_type_ids_mask_function(
             token_type_ids.to(cache_position.device), image_group_ids
         )

filename: src/transformers/models/gemma3/modular_gemma3.py
patch: @@ -764,7 +764,7 @@ def create_causal_mask_mapping(
         is_previous_image = nn.functional.pad(is_image, (1, 0), value=0)[:, :-1]
         new_image_start = is_image & ~is_previous_image
         image_group_ids = torch.cumsum(new_image_start.int(), dim=1) - 1
-        image_group_ids = torch.where(is_image, image_group_ids, torch.full_like(token_type_ids, -1))
+        image_group_ids = torch.where(is_image, image_group_ids, -1)
         mask_kwargs["or_mask_function"] = token_type_ids_mask_function(
             token_type_ids.to(cache_position.device), image_group_ids
         )

filename: tests/models/gemma3/test_modeling_gemma3.py
patch: @@ -19,6 +19,7 @@
 
 import pytest
 from parameterized import parameterized
+from pytest import mark
 
 from transformers import (
     AutoModelForCausalLM,
@@ -33,9 +34,11 @@
     is_flash_attn_2_available,
     require_deterministic_for_xpu,
     require_flash_attn,
+    require_flash_attn_3,
     require_read_token,
     require_torch,
     require_torch_accelerator,
+    require_torch_gpu,
     require_torch_large_accelerator,
     slow,
     torch_device,
@@ -342,6 +345,20 @@ def test_automodelforcausallm(self):
             for_causal_lm = AutoModelForCausalLM.from_pretrained(tmp_dir)
             self.assertIsInstance(for_causal_lm, Gemma3ForConditionalGeneration)
 
+    @require_flash_attn
+    @require_torch_gpu
+    @mark.flash_attn_test
+    @slow
+    def test_flash_attn_2_from_config(self):
+        self.flash_attn_from_config(attn_implementation="flash_attention_2", test_fwd_in_train=False)
+
+    @require_flash_attn_3
+    @require_torch_gpu
+    @mark.flash_attn_3_test
+    @slow
+    def test_flash_attn_3_from_config(self):
+        self.flash_attn_from_config(attn_implementation="flash_attention_3", test_fwd_in_train=False)
+
 
 @slow
 @require_torch_accelerator

filename: tests/test_modeling_common.py
patch: @@ -2976,7 +2976,7 @@ def test_model_is_small(self):
 
     def flash_attn_inference_equivalence(
         self, attn_implementation: str, padding_side: str, atol: float = 4e-2, rtol: float = 4e-2
-    ):
+    ) -> None:
         r"""
         Tests the equivalence between the eager and flash attention implementations.
         This test is only for inference and runs with `dtype=torch.bfloat16`.
@@ -3114,9 +3114,6 @@ def flash_attn_inference_equivalence(
                 torch.testing.assert_close(logits_1_eager, logits_1_fa, atol=atol, rtol=rtol)
                 if padding_side == "left":
                     torch.testing.assert_close(logits_2_eager[1:], logits_2_fa[1:], atol=atol, rtol=rtol)
-                    # Check it can run in training mode
-                    model.train()
-                    _ = model(**second_inputs)
                 else:
                     torch.testing.assert_close(logits_2_eager[:-1], logits_2_fa[:-1], atol=atol, rtol=rtol)
 
@@ -3651,7 +3648,7 @@ def test_flash_attn_2_can_compile_with_attention_mask_None_without_graph_break(s
 
         assert not loss.isnan().any()
 
-    def flash_attn_from_config(self, attn_implementation: str):
+    def flash_attn_from_config(self, attn_implementation: str, test_fwd_in_train: bool = True):
         r"""
         Tests if the model can be loaded with `attn_implementation` from the config and if the
         weights are not randomly initialized.
@@ -3669,6 +3666,14 @@ def flash_attn_from_config(self, attn_implementation: str):
                 config, attn_implementation=attn_implementation, dtype=torch.bfloat16
             ).to(torch_device)
 
+            # By default, we perform the forward pass in train mode, because it's more sctrict than eval mode. If the
+            # forward pass is successful in train mode, it will also be successful in eval mode. But since some models
+            # (eg. gemma3) need different inputs in train mode we have the option to test the forward pass in eval mode.
+            if test_fwd_in_train:
+                fa_model = fa_model.train()
+            else:
+                fa_model = fa_model.eval()
+
             dummy_input = inputs_dict[fa_model.main_input_name]
             if dummy_input.dtype in [torch.float32, torch.float16]:
                 dummy_input = dummy_input.to(torch.bfloat16)