RecurrentGemma crashes during inference for inputs longer than sliding window width

[assafbk]

System Info

System Info:

• transformers version: 4.50.3
• Platform: Linux-6.8.0-52-generic-x86_64-with-glibc2.35
• Python version: 3.10.16
• Huggingface_hub version: 0.30.1
• Safetensors version: 0.5.3
• Accelerate version: 1.6.0
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (GPU?): 2.5.1+cu118 (True)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?: no
• Using GPU in script?: yes
• GPU type: NVIDIA RTX A6000

Who can help?

@ArthurZucker, @gante

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Code snippet for reproduction:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("google/recurrentgemma-9b-it")
model = AutoModelForCausalLM.from_pretrained("google/recurrentgemma-9b-it", device_map="cuda", torch_dtype=torch.float16)

input_text = "Write me a poem about Machine Learning." * 300    # This string is 2402 tokens long, which is larger than 2048, the sliding window attention width
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=20)
print(tokenizer.decode(outputs[0]))

Error message:

Traceback (most recent call last):
File "/data2/assaf/tmp/test_rg.py", line 13, in
outputs = model.generate(**input_ids, max_new_tokens=20)
File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/transformers/generation/utils.py", line 2326, in generate
result = self._sample(
File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/transformers/generation/utils.py", line 3289, in _sample
outputs = model_forward(**model_inputs, return_dict=True)
File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py", line 852, in forward
outputs = self.model(
File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py", line 717, in forward
causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position)
File "/data2/assaf/conda_envs/recurrent_gemma_tmp/lib/python3.10/site-packages/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py", line 764, in _update_causal_mask
padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0)
RuntimeError: The size of tensor a (2048) must match the size of tensor b (2402) at non-singleton dimension 3

Expected behavior

If the sequence is longer than the sliding window width (like it is now in the script) then the script crashes with the error message above.

If the sequence is shorter than the sliding window width (e.g. replace *300 by *200) then the script runs fine.

The bug was seen in transformers version v4.50.3
It does not reproduce on earlier transformers versions, such as v4.42.4

[Rocketknight1]

Hi @assafbk, can you use git bisect to identify the commit where this bug appeared? It'll help a lot in identifying the cause and assigning the right person to fix it!

[manueldeprada]

RecurrentGemma generation was functional until commit 295a90c (v4.46.0), which introduced two regressions:

1. Removed the model-specific prepare_inputs_for_generation, and in doing so, stopped forwarding past_key_values via **kwargs. This broke generation entirely.
2. Even if past_key_values were passed correctly, the same commit also introduced a bug for long prompts due to incorrect attention_mask handling during decoding.

Commit 3d6e55c (v4.48.3) later restored forwarding of **kwargs, which incidentally fixed generation. However, the long-context bug persists.

Root cause

RecurrentGemma assumes that the attention_mask received during decoding is already cropped to the model’s sliding attention window (attention_window_size). The generic prepare_inputs_for_generation does not currently enforce this, leading to shape mismatches in _update_causal_mask() when the prompt length exceeds the attention window.

Proposed fix

Add a one-line safeguard inside _update_causal_mask() to truncate attention_mask when it's longer than the causal mask, as its was done on the original prepare_inputs_for_generation.

Submitting a PR with the fix.