From b7b0fbf915b9febf194fd2385b10d113d3531b31 Mon Sep 17 00:00:00 2001 From: Harshal Janjani Date: Mon, 9 Nov 2026 20:12:01 +0000 Subject: [PATCH] fix: Restore KV sharing for the Gemma 4 assistants and their targets Signed-off-by: Harshal Janjani --- .../models/gemma4/configuration_gemma4.py | 6 +-- .../configuration_gemma4_unified.py | 6 +-- .../test_modeling_gemma4_assistant.py | 51 +++++++++++++++++++ .../test_modeling_gemma4_unified_assistant.py | 47 +++++++++++++++++ 4 files changed, 104 insertions(+), 6 deletions(-) diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py index d2b9c8375e3b..51bf510569ca 100644 --- a/src/transformers/models/gemma4/configuration_gemma4.py +++ b/src/transformers/models/gemma4/configuration_gemma4.py @@ -226,9 +226,9 @@ def __post_init__(self, **kwargs): @property def kv_sharing_modes(self) -> list[str | None]: - # In this case there are no shared layers - if self.num_kv_shared_layers == 0: - return [None] * self.num_hidden_layers + # In this case all layers are shared + if self.num_kv_shared_layers == self.num_hidden_layers: + return ["reuse"] * self.num_hidden_layers # After this index, all layers are consumers first_shared_layer = self.num_hidden_layers - self.num_kv_shared_layers diff --git a/src/transformers/models/gemma4_unified/configuration_gemma4_unified.py b/src/transformers/models/gemma4_unified/configuration_gemma4_unified.py index 9fe3ef7e4bb7..ad5806b4b4aa 100644 --- a/src/transformers/models/gemma4_unified/configuration_gemma4_unified.py +++ b/src/transformers/models/gemma4_unified/configuration_gemma4_unified.py @@ -193,9 +193,9 @@ def __post_init__(self, **kwargs): @property def kv_sharing_modes(self) -> list[str | None]: - # In this case there are no shared layers - if self.num_kv_shared_layers == 0: - return [None] * self.num_hidden_layers + # In this case all layers are shared + if self.num_kv_shared_layers == self.num_hidden_layers: + return ["reuse"] * self.num_hidden_layers # After this index, all layers are consumers first_shared_layer = self.num_hidden_layers - self.num_kv_shared_layers diff --git a/tests/models/gemma4_assistant/test_modeling_gemma4_assistant.py b/tests/models/gemma4_assistant/test_modeling_gemma4_assistant.py index 870f5b84ec4f..0576ecff7212 100644 --- a/tests/models/gemma4_assistant/test_modeling_gemma4_assistant.py +++ b/tests/models/gemma4_assistant/test_modeling_gemma4_assistant.py @@ -14,6 +14,8 @@ import unittest +from parameterized import parameterized + from transformers import is_torch_available from transformers.testing_utils import ( Expectations, @@ -23,17 +25,66 @@ torch_device, ) +from ...test_modeling_common import ids_tensor from ...test_processing_common import url_to_local_path if is_torch_available(): from transformers import ( AutoModelForCausalLM, + Gemma4AssistantConfig, + Gemma4AssistantForCausalLM, + Gemma4ForCausalLM, Gemma4ForConditionalGeneration, Gemma4Processor, + Gemma4TextConfig, ) +@require_torch +class Gemma4AssistantTest(unittest.TestCase): + @parameterized.expand([0, 2]) + def test_assisted_decoding_matches_greedy_search(self, num_kv_shared_layers): + text_config_kwargs = { + "vocab_size": 99, + "hidden_size": 32, + "intermediate_size": 64, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 16, + "global_head_dim": 16, + "sliding_window": 8, + } + model_config = Gemma4TextConfig( + num_hidden_layers=4, + layer_types=["sliding_attention", "full_attention", "sliding_attention", "full_attention"], + num_kv_shared_layers=num_kv_shared_layers, + hidden_size_per_layer_input=16, + vocab_size_per_layer_input=99, + **text_config_kwargs, + ) + assistant_config = Gemma4AssistantConfig( + text_config=Gemma4TextConfig( + num_hidden_layers=2, + layer_types=["sliding_attention", "full_attention"], + hidden_size_per_layer_input=0, + vocab_size_per_layer_input=0, + **text_config_kwargs, + ), + backbone_hidden_size=32, + ) + model = Gemma4ForCausalLM(model_config).to(torch_device).eval() + assistant_model = Gemma4AssistantForCausalLM(assistant_config).to(torch_device).eval() + assistant_model.generation_config.num_assistant_tokens = 2 + + input_ids = ids_tensor([1, 7], vocab_size=99) + output_greedy = model.generate(input_ids, max_new_tokens=5, do_sample=False, eos_token_id=-1) + output_assisted = model.generate( + input_ids, assistant_model=assistant_model, max_new_tokens=5, do_sample=False, eos_token_id=-1 + ) + self.assertListEqual(output_greedy.tolist(), output_assisted.tolist()) + + @slow @require_torch @unittest.skip(reason="Update after release") # TODO @vasqu diff --git a/tests/models/gemma4_unified_assistant/test_modeling_gemma4_unified_assistant.py b/tests/models/gemma4_unified_assistant/test_modeling_gemma4_unified_assistant.py index b078e0d331bd..298c9ad3ec9f 100644 --- a/tests/models/gemma4_unified_assistant/test_modeling_gemma4_unified_assistant.py +++ b/tests/models/gemma4_unified_assistant/test_modeling_gemma4_unified_assistant.py @@ -14,6 +14,8 @@ import unittest +from parameterized import parameterized + from transformers import is_torch_available from transformers.testing_utils import ( Expectations, @@ -23,6 +25,7 @@ torch_device, ) +from ...test_modeling_common import ids_tensor from ...test_processing_common import url_to_local_path @@ -31,9 +34,53 @@ AutoModelForCausalLM, Gemma4ForConditionalGeneration, Gemma4Processor, + Gemma4UnifiedAssistantConfig, + Gemma4UnifiedAssistantForCausalLM, + Gemma4UnifiedForCausalLM, + Gemma4UnifiedTextConfig, ) +@require_torch +class Gemma4UnifiedAssistantTest(unittest.TestCase): + @parameterized.expand([0, 2]) + def test_assisted_decoding_matches_greedy_search(self, num_kv_shared_layers): + text_config_kwargs = { + "vocab_size": 99, + "hidden_size": 32, + "intermediate_size": 64, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 16, + "global_head_dim": 16, + "sliding_window": 8, + } + model_config = Gemma4UnifiedTextConfig( + num_hidden_layers=4, + layer_types=["sliding_attention", "full_attention", "sliding_attention", "full_attention"], + num_kv_shared_layers=num_kv_shared_layers, + **text_config_kwargs, + ) + assistant_config = Gemma4UnifiedAssistantConfig( + text_config=Gemma4UnifiedTextConfig( + num_hidden_layers=2, + layer_types=["sliding_attention", "full_attention"], + **text_config_kwargs, + ), + backbone_hidden_size=32, + ) + model = Gemma4UnifiedForCausalLM(model_config).to(torch_device).eval() + assistant_model = Gemma4UnifiedAssistantForCausalLM(assistant_config).to(torch_device).eval() + assistant_model.generation_config.num_assistant_tokens = 2 + + input_ids = ids_tensor([1, 7], vocab_size=99) + output_greedy = model.generate(input_ids, max_new_tokens=5, do_sample=False, eos_token_id=-1) + output_assisted = model.generate( + input_ids, assistant_model=assistant_model, max_new_tokens=5, do_sample=False, eos_token_id=-1 + ) + self.assertListEqual(output_greedy.tolist(), output_assisted.tolist()) + + @slow @require_torch @unittest.skip(reason="Update after release") # TODO(vasqu)