diff --git a/src/transformers/models/gemma3/modeling_gemma3.py b/src/transformers/models/gemma3/modeling_gemma3.py index a197a02ed590..a78117c29e59 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -739,8 +739,6 @@ def create_causal_mask_mapping( past_key_values: Cache | None, position_ids: torch.Tensor | None, token_type_ids: torch.Tensor | None = None, - pixel_values: torch.FloatTensor | None = None, - is_training: bool = False, is_first_iteration: bool | None = None, **kwargs, ) -> dict: @@ -750,9 +748,6 @@ def create_causal_mask_mapping( Uses `pixel_values` as an optional input to disambiguate edge cases. """ - if is_training and token_type_ids is None: - raise ValueError("`token_type_ids` is required as a model input when training") - mask_kwargs = { "config": config.get_text_config(), "inputs_embeds": inputs_embeds, @@ -760,15 +755,7 @@ def create_causal_mask_mapping( "past_key_values": past_key_values, "position_ids": position_ids, } - # NOTE: this `may_have_image_input` logic is not flawless, it fails when we're using a cache eagerly initialized - # (e.g. compiled prefill) AND `pixel_values` are not provided (i.e. the image data is provided through other - # means). Determining prefill in that case requires checking data values, which is not compile-compatible. - is_first_iteration = ( - is_first_iteration - if is_first_iteration is not None - else (past_key_values is None or not past_key_values.is_initialized or pixel_values is not None) - ) - if token_type_ids is not None and is_first_iteration: + if token_type_ids is not None: # We need to pass an additional mask function to account for token type ids, and it needs to be an `or` (to # undo the causal masking) @@ -915,13 +902,11 @@ def forward( if not isinstance(causal_mask_mapping := attention_mask, dict): causal_mask_mapping = create_causal_mask_mapping( self.config, - inputs_embeds, - attention_mask, - past_key_values, - position_ids, - token_type_ids, - pixel_values, - is_training=self.training, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + token_type_ids=token_type_ids, ) outputs = self.language_model( @@ -1110,6 +1095,9 @@ def prepare_inputs_for_generation( # iteration with a question and cached system prompt (continue generate from cache). NOTE: use_cache=False needs pixel_values always if is_first_iteration or not use_cache: model_inputs["pixel_values"] = pixel_values + else: + # Don't pass to not apply bidirectional mask on top + model_inputs["token_type_ids"] = None return model_inputs diff --git a/src/transformers/models/gemma3/modular_gemma3.py b/src/transformers/models/gemma3/modular_gemma3.py index 1e96f5acceb9..24b3f7043e7c 100644 --- a/src/transformers/models/gemma3/modular_gemma3.py +++ b/src/transformers/models/gemma3/modular_gemma3.py @@ -612,8 +612,6 @@ def create_causal_mask_mapping( past_key_values: Cache | None, position_ids: torch.Tensor | None, token_type_ids: torch.Tensor | None = None, - pixel_values: torch.FloatTensor | None = None, - is_training: bool = False, is_first_iteration: bool | None = None, **kwargs, ) -> dict: @@ -623,9 +621,6 @@ def create_causal_mask_mapping( Uses `pixel_values` as an optional input to disambiguate edge cases. """ - if is_training and token_type_ids is None: - raise ValueError("`token_type_ids` is required as a model input when training") - mask_kwargs = { "config": config.get_text_config(), "inputs_embeds": inputs_embeds, @@ -633,15 +628,7 @@ def create_causal_mask_mapping( "past_key_values": past_key_values, "position_ids": position_ids, } - # NOTE: this `may_have_image_input` logic is not flawless, it fails when we're using a cache eagerly initialized - # (e.g. compiled prefill) AND `pixel_values` are not provided (i.e. the image data is provided through other - # means). Determining prefill in that case requires checking data values, which is not compile-compatible. - is_first_iteration = ( - is_first_iteration - if is_first_iteration is not None - else (past_key_values is None or not past_key_values.is_initialized or pixel_values is not None) - ) - if token_type_ids is not None and is_first_iteration: + if token_type_ids is not None: # We need to pass an additional mask function to account for token type ids, and it needs to be an `or` (to # undo the causal masking) @@ -718,13 +705,11 @@ def forward( if not isinstance(causal_mask_mapping := attention_mask, dict): causal_mask_mapping = create_causal_mask_mapping( self.config, - inputs_embeds, - attention_mask, - past_key_values, - position_ids, - token_type_ids, - pixel_values, - is_training=self.training, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + token_type_ids=token_type_ids, ) outputs = self.language_model( @@ -897,6 +882,9 @@ def prepare_inputs_for_generation( # iteration with a question and cached system prompt (continue generate from cache). NOTE: use_cache=False needs pixel_values always if is_first_iteration or not use_cache: model_inputs["pixel_values"] = pixel_values + else: + # Don't pass to not apply bidirectional mask on top + model_inputs["token_type_ids"] = None return model_inputs diff --git a/src/transformers/models/gemma4/modeling_gemma4.py b/src/transformers/models/gemma4/modeling_gemma4.py index 5b147f95ba36..487359a190dc 100644 --- a/src/transformers/models/gemma4/modeling_gemma4.py +++ b/src/transformers/models/gemma4/modeling_gemma4.py @@ -2086,8 +2086,6 @@ def create_causal_mask_mapping( past_key_values: Cache | None, position_ids: torch.Tensor | None, mm_token_type_ids: torch.Tensor | None = None, - pixel_values: torch.FloatTensor | None = None, - is_training: bool = False, is_first_iteration: bool | None = None, **kwargs, ) -> dict: @@ -2097,9 +2095,6 @@ def create_causal_mask_mapping( Uses `pixel_values` as an optional input to disambiguate edge cases. """ - if is_training and mm_token_type_ids is None: - raise ValueError("`mm_token_type_ids` is required as a model input when training") - mask_kwargs = { "config": config.get_text_config(), "inputs_embeds": inputs_embeds, @@ -2109,15 +2104,7 @@ def create_causal_mask_mapping( } sliding_mask_kwargs = mask_kwargs.copy() - # NOTE: this `may_have_image_input` logic is not flawless, it fails when we're using a cache eagerly initialized - # (e.g. compiled prefill) AND `pixel_values` are not provided (i.e. the image data is provided through other - # means). Determining prefill in that case requires checking data values, which is not compile-compatible. - is_first_iteration = ( - is_first_iteration - if is_first_iteration is not None - else (past_key_values is None or not past_key_values.is_initialized or pixel_values is not None) - ) - if mm_token_type_ids is not None and is_first_iteration: + if mm_token_type_ids is not None: # We need to pass an additional mask function to account for token type ids, and it needs to be an `or` (to # undo the causal masking) @@ -2368,13 +2355,11 @@ def forward( # Larger Gemma 4 models use Gemma 3's bidirectional attention mask for vision inputs causal_mask_mapping = create_causal_mask_mapping( self.config, - inputs_embeds, - attention_mask, - past_key_values, - position_ids, - mm_token_type_ids, - pixel_values, - is_training=self.training, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + mm_token_type_ids=mm_token_type_ids, ) else: # Smaller Gemma models use a conventional casual attention mask @@ -2608,6 +2593,9 @@ def prepare_inputs_for_generation( model_inputs["pixel_values_videos"] = pixel_values_videos model_inputs["input_features"] = input_features model_inputs["input_features_mask"] = input_features_mask + else: + # Don't pass to not apply bidirectional mask on top + model_inputs["mm_token_type_ids"] = None return model_inputs diff --git a/src/transformers/models/gemma4/modular_gemma4.py b/src/transformers/models/gemma4/modular_gemma4.py index 12412b319b5c..7f2a22c79b8b 100644 --- a/src/transformers/models/gemma4/modular_gemma4.py +++ b/src/transformers/models/gemma4/modular_gemma4.py @@ -1660,8 +1660,6 @@ def create_causal_mask_mapping( past_key_values: Cache | None, position_ids: torch.Tensor | None, mm_token_type_ids: torch.Tensor | None = None, - pixel_values: torch.FloatTensor | None = None, - is_training: bool = False, is_first_iteration: bool | None = None, **kwargs, ) -> dict: @@ -1671,9 +1669,6 @@ def create_causal_mask_mapping( Uses `pixel_values` as an optional input to disambiguate edge cases. """ - if is_training and mm_token_type_ids is None: - raise ValueError("`mm_token_type_ids` is required as a model input when training") - mask_kwargs = { "config": config.get_text_config(), "inputs_embeds": inputs_embeds, @@ -1683,15 +1678,7 @@ def create_causal_mask_mapping( } sliding_mask_kwargs = mask_kwargs.copy() - # NOTE: this `may_have_image_input` logic is not flawless, it fails when we're using a cache eagerly initialized - # (e.g. compiled prefill) AND `pixel_values` are not provided (i.e. the image data is provided through other - # means). Determining prefill in that case requires checking data values, which is not compile-compatible. - is_first_iteration = ( - is_first_iteration - if is_first_iteration is not None - else (past_key_values is None or not past_key_values.is_initialized or pixel_values is not None) - ) - if mm_token_type_ids is not None and is_first_iteration: + if mm_token_type_ids is not None: # We need to pass an additional mask function to account for token type ids, and it needs to be an `or` (to # undo the causal masking) @@ -1952,13 +1939,11 @@ def forward( # Larger Gemma 4 models use Gemma 3's bidirectional attention mask for vision inputs causal_mask_mapping = create_causal_mask_mapping( self.config, - inputs_embeds, - attention_mask, - past_key_values, - position_ids, - mm_token_type_ids, - pixel_values, - is_training=self.training, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + mm_token_type_ids=mm_token_type_ids, ) else: # Smaller Gemma models use a conventional casual attention mask @@ -2183,6 +2168,9 @@ def prepare_inputs_for_generation( model_inputs["pixel_values_videos"] = pixel_values_videos model_inputs["input_features"] = input_features model_inputs["input_features_mask"] = input_features_mask + else: + # Don't pass to not apply bidirectional mask on top + model_inputs["mm_token_type_ids"] = None return model_inputs diff --git a/src/transformers/models/git/modeling_git.py b/src/transformers/models/git/modeling_git.py index 242d3ca960f2..9be97d01c425 100644 --- a/src/transformers/models/git/modeling_git.py +++ b/src/transformers/models/git/modeling_git.py @@ -109,8 +109,6 @@ def create_causal_mask_mapping( past_key_values: Cache | None, position_ids: torch.Tensor | None, token_type_ids: torch.Tensor | None = None, - pixel_values: torch.FloatTensor | None = None, - is_training: bool = False, is_first_iteration: bool | None = None, **kwargs, ) -> dict: @@ -120,9 +118,6 @@ def create_causal_mask_mapping( Uses `pixel_values` as an optional input to disambiguate edge cases. """ - if is_training and token_type_ids is None: - raise ValueError("`token_type_ids` is required as a model input when training") - mask_kwargs = { "config": config.get_text_config(), "inputs_embeds": inputs_embeds, @@ -130,15 +125,7 @@ def create_causal_mask_mapping( "past_key_values": past_key_values, "position_ids": position_ids, } - # NOTE: this `may_have_image_input` logic is not flawless, it fails when we're using a cache eagerly initialized - # (e.g. compiled prefill) AND `pixel_values` are not provided (i.e. the image data is provided through other - # means). Determining prefill in that case requires checking data values, which is not compile-compatible. - is_first_iteration = ( - is_first_iteration - if is_first_iteration is not None - else (past_key_values is None or not past_key_values.is_initialized or pixel_values is not None) - ) - if token_type_ids is not None and is_first_iteration: + if token_type_ids is not None: # We need to pass an additional mask function to account for token type ids, and it needs to be an `or` (to # undo the causal masking) @@ -948,12 +935,11 @@ def forward( # Images attend each other bidirectionally while text remains causal causal_mask = create_causal_mask_mapping( self.config, - embedding_output, - attention_mask, - past_key_values, - None, - token_type_ids, - pixel_values, + inputs_embeds=embedding_output, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=None, + token_type_ids=token_type_ids, ) hidden_states = embedding_output diff --git a/tests/models/gemma3/test_modeling_gemma3.py b/tests/models/gemma3/test_modeling_gemma3.py index fe65a3f83bcf..913d6b9cf5ff 100644 --- a/tests/models/gemma3/test_modeling_gemma3.py +++ b/tests/models/gemma3/test_modeling_gemma3.py @@ -304,6 +304,24 @@ class Gemma3Vision2TextModelTest(VLMModelTest, unittest.TestCase): test_disk_offload_safetensors = False test_disk_offload_bin = False + def test_training(self): + # Overwrite to test training with text-only samples, should not raise errors + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + config.return_dict = True + + model = Gemma3ForConditionalGeneration(config) + model.to(torch_device) + model.train() + inputs = self._prepare_for_class(inputs_dict, Gemma3ForConditionalGeneration, return_labels=True) + loss = model(**inputs).loss + loss.backward() + + # pop out image-related inputs and try to run forward + inputs.pop("token_type_ids", None) + inputs.pop("pixel_values", None) + loss = model(**inputs).loss + loss.backward() + @unittest.skip("Gemma3 applies key/query norm which doesn't work with packing") def test_flash_attention_2_padding_matches_padding_free_with_position_ids(self): pass diff --git a/tests/models/gemma4/test_modeling_gemma4.py b/tests/models/gemma4/test_modeling_gemma4.py index 91694b5c1d45..e694fae48362 100644 --- a/tests/models/gemma4/test_modeling_gemma4.py +++ b/tests/models/gemma4/test_modeling_gemma4.py @@ -123,6 +123,9 @@ def test_sdpa_padding_matches_padding_free_with_position_ids(self): def test_tp_generation_quantized(self): pass + def test_model_training(self): + pass + class Gemma4Audio2TextModelTester: def __init__( @@ -386,6 +389,24 @@ def setUp(self): self.model_tester = Gemma4Vision2TextModelTester(self) self.config_tester = ConfigTester(self, config_class=Gemma4Config, hidden_size=37) + def test_training(self): + # Overwrite to test training with text-only samples, should not raise errors + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + config.return_dict = True + + model = Gemma4ForConditionalGeneration(config) + model.to(torch_device) + model.train() + inputs = self._prepare_for_class(inputs_dict, Gemma4ForConditionalGeneration, return_labels=True) + loss = model(**inputs).loss + loss.backward() + + # pop out image-related inputs and try to run forward + inputs.pop("mm_token_type_ids", None) + inputs.pop("pixel_values", None) + loss = model(**inputs).loss + loss.backward() + @unittest.skip("The tester has no audios in input dict") def test_get_audio_features_hidden_states(self): pass