Skip to content

Gemma4 training with text-only samples #45454

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zucchini-nlp merged 6 commits into from
Apr 22, 2026
Merged

Gemma4 training with text-only samples #45454

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
2766977
fix and add a test
zucchini-nlp Apr 15, 2026
d054297
fix copies
zucchini-nlp Apr 15, 2026
0c6ffce
oops, pass keyword args!
zucchini-nlp Apr 15, 2026
6693fc1
dont check if prefill
zucchini-nlp Apr 15, 2026
9d312d1
Merge branch 'main' into gemma4-token-types
zucchini-nlp Apr 21, 2026
86d5996
Merge branch 'main' into gemma4-token-types
zucchini-nlp Apr 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 9 additions & 21 deletions src/transformers/models/gemma3/modeling_gemma3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -750,25 +748,14 @@ 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,
"attention_mask": attention_mask,
"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)

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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

Expand Down
30 changes: 9 additions & 21 deletions src/transformers/models/gemma3/modular_gemma3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -623,25 +621,14 @@ 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,
"attention_mask": attention_mask,
"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)

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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

Expand Down
30 changes: 9 additions & 21 deletions src/transformers/models/gemma4/modeling_gemma4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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,
Expand All @@ -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)

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
30 changes: 9 additions & 21 deletions src/transformers/models/gemma4/modular_gemma4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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,
Expand All @@ -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)

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
26 changes: 6 additions & 20 deletions src/transformers/models/git/modeling_git.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -120,25 +118,14 @@ 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,
"attention_mask": attention_mask,
"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)

Expand Down Expand Up @@ -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
Expand Down
18 changes: 18 additions & 0 deletions tests/models/gemma3/test_modeling_gemma3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
21 changes: 21 additions & 0 deletions tests/models/gemma4/test_modeling_gemma4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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__(
Expand Down Expand Up @@ -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
Expand Down
Loading