diff --git a/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py
index da7a5d8ee0da..6d6b70b1bf33 100644
--- a/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py
+++ b/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py
@@ -54,7 +54,7 @@ class RecurrentGemmaConfig(PreTrainedConfig):
"""
model_type = "recurrent_gemma"
- attribute_map = {"sliding_window": "attention_window_size"}
+ attribute_map = {"sliding_window": "attention_window_size", "layer_types": "layers_block_type"}
num_hidden_layers: int = 26
vocab_size: int = 256000
@@ -97,7 +97,9 @@ def validate_architecture(self):
@property
def layers_block_type(self):
- return (self.block_types * 100)[: self.num_hidden_layers]
+ # Follow the `layer_types` conventions: the attention blocks always use local (sliding window) attention
+ block_types = (self.block_types * 100)[: self.num_hidden_layers]
+ return ["linear_attention" if x == "recurrent" else "sliding_attention" for x in block_types]
__all__ = ["RecurrentGemmaConfig"]
diff --git a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py
index 5b276a58a06f..5d7196b81f5f 100644
--- a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py
+++ b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py
@@ -465,7 +465,10 @@ def _setup_cache(self, batch, device, dtype):
self.conv1d_state = torch.zeros((batch, self.hidden_size, self.conv1d_width - 1), device=device, dtype=dtype)
-TEMPORAL_BLOCK_CLASSES = {"recurrent": RecurrentGemmaRecurrentBlock, "attention": RecurrentGemmaAttention}
+TEMPORAL_BLOCK_CLASSES = {
+ "linear_attention": RecurrentGemmaRecurrentBlock,
+ "sliding_attention": RecurrentGemmaAttention,
+}
class RecurrentGemmaMlp(nn.Module):
@@ -652,13 +655,17 @@ def forward(
position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device) + past_seen_tokens
position_ids = position_ids.unsqueeze(0)
- causal_mask = create_sliding_window_causal_mask(
- config=self.config,
- inputs_embeds=inputs_embeds,
- attention_mask=attention_mask,
- past_key_values=past_key_values,
- position_ids=position_ids,
- )
+ # It may already have been prepared by e.g. `generate`
+ if isinstance(attention_mask, dict):
+ causal_mask = attention_mask["sliding_attention"]
+ else:
+ causal_mask = create_sliding_window_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
hidden_states = hidden_states * self.normalizer.type(hidden_states.dtype)
diff --git a/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py b/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py
index 055faa135fb2..a77c7994400a 100644
--- a/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py
+++ b/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py
@@ -181,6 +181,13 @@ def test_assisted_decoding_matches_greedy_search(self):
def test_left_padding_compatibility(self):
pass
+ @unittest.skip(
+ "RecurrentGemma's recurrent block has no chunked-continuation support: a cached multi-token forward is "
+ "treated as a new prefill, so the split-vs-single comparison cannot match regardless of padding masking."
+ )
+ def test_recurrent_layers_mask_padding_on_continued_forward(self):
+ pass
+
@pytest.mark.generate
@unittest.skip(reason="Relies on `past_key_values` returned by the model. Not supported with recurrent gemma")
def test_assisted_decoding_sample(self):
@@ -238,7 +245,7 @@ class RecurrentGemmaIntegrationTest(MemoryCleanupMixin, unittest.TestCase):
model_id = "google/recurrentgemma-2b"
def test_2b_generate(self):
- EXPECTED_TEXTS = ['Hello I am doing a project on the topic of "The impact of the internet on the society" and I am looking for some information on the topic. I am looking for some information on the impact of the internet on the society. I am looking for some information on the impact of the internet on the society. I am looking for some', 'Hi today is a new app that allows you to make money by watching videos.\n\nThe app is very simple to use and you can earn money by watching videos.\n\nThe app is available for both Android and iOS devices and you can download it from the Google Play Store or the App Store.\n\nOnce you have downloaded the app'] # fmt: skip
+ EXPECTED_TEXTS = ['Hello I am doing a project on the topic of "The impact of the internet on the society" and I am looking for some information on the topic. I am looking for some information on the impact of the internet on the society. I am looking for some information on the impact of the internet on the society. I am looking for some', "Hi today the 19th of June 2019, I was in the office of the Ministry of Health and Social Protection of the Republic of Uzbekistan. I was there to discuss the issue of the implementation of the project “Strengthening the capacity of the Ministry of Health and Social Protection of the Republic of Uzbekistan to"] # fmt: skip
model = AutoModelForCausalLM.from_pretrained(
self.model_id,
).to(torch_device)
@@ -254,7 +261,7 @@ def test_2b_generate(self):
self.assertEqual(output_text, EXPECTED_TEXTS)
tokenizer.padding_side = "left"
- EXPECTED_TEXTS = ['Hello I am doing a project on the topic of "The impact of the internet on the society" and I am looking for some information on the topic. I am looking for some information on the impact of the internet on the society. I am looking for some information on the impact of the internet on the society. I am looking for some', 'Hi today I’m going to show you how to make a simple and easy to make a DIY DIY DIY DIY DIY DIY DIY DIY DIY DIY DIY DIY'] # fmt: skip
+ EXPECTED_TEXTS = ['Hello I am doing a project on the topic of "The impact of the internet on the society" and I am looking for some information on the topic. I am looking for some information on the impact of the internet on the society. I am looking for some information on the impact of the internet on the society. I am looking for some', "Hi today I'm going to show you how to make a simple and easy to make a 3D model of a house.\n\nYou will need:\n\n* A 3D printer\n* A 3D model of a house\n* A 3D printer filament\n* A 3D printer pen\n"] # fmt: skip
inputs = tokenizer(self.input_text, return_tensors="pt", padding=True).to(torch_device)
output = model.generate(**inputs, max_new_tokens=64, do_sample=False)
@@ -297,7 +304,8 @@ def test_model_2b_8bit(self):
EXPECTED_TEXTS = Expectations(
{
("xpu", None): ['Hello I am doing a project on the topic of "The impact of the internet on the society" and I am stuck', "Hi today I'm going to show you how to make a simple and easy to make a 3D"],
- (None, None): ['Hello I am doing a project on the topic of "The impact of social media on the society" and I am looking', "Hi today I'm going to show you how to make a simple and easy to make a 3D"],
+ ("rocm", (9, 4)): ['Hello I am doing a project on the 1960s and I am looking for a picture of a ', 'Hi today I am going to show you how to make a simple and easy to make a DIY'],
+ (None, None): ['Hello I am doing a project on the effects of the environment on the human body. I am looking for a few good', "Hi today I'm going to show you how to make a simple and easy to make a 3D"],
}
)
# fmt: on