Reccurent Gemma forward pass does not match with Google DeepMind original implementation
Author: JRD971000
Pull request: no

### System Info

- `transformers` version: 4.40.0
- Platform: Linux-5.15.0-102-generic-x86_64-with-glibc2.31
- Python version: 3.11.8
- Huggingface_hub version: 0.22.2
- Safetensors version: 0.4.3
- Accelerate version: 0.29.3
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: NO
        - mixed_precision: fp16
        - use_cpu: False
        - debug: False
        - num_processes: 1
        - machine_rank: 0
        - num_machines: 1
        - gpu_ids: 0
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - enable_cpu_affinity: False
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.2.1+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): 0.8.2 (cpu)
- Jax version: 0.4.25
- JaxLib version: 0.4.25
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@ArthurZucker and @younesbelkada

### 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

## Steps to reproduce:

1. clone and install the original repo: https://github.com/google-deepmind/recurrentgemma
2. run the following reproducer:

```
from space_gemma import jax as space_gemma
import space_gemma as torch_space_gemma
import itertools
import torch 
from space_gemma.torch.griffin import Griffin
from transformers import AutoModelForCausalLM, AutoTokenizer

pattern = (
        space_gemma.TemporalBlockType.RECURRENT,
        space_gemma.TemporalBlockType.RECURRENT,
        space_gemma.TemporalBlockType.ATTENTION,
    )
num_layers = 26
config = space_gemma.GriffinConfig(
    vocab_size=256_128,
    width=2560,
    mlp_expanded_width=3 * 2560,
    attention_window_size=1024,
    num_heads=10,
    block_types=tuple(
        itertools.islice(itertools.cycle(pattern), num_layers)),
)

model = Griffin(config)
PATH_TO_DEEPMIND_MODEL_WEIGHTS = ...
dm_model_weight = torch.load(PATH_TO_DEEPMIND_MODEL_WEIGHTS)
model.load_state_dict(dm_model_weight, strict=True)
model = model.cpu()
print("Copied weights")

hf_model = AutoModelForCausalLM.from_pretrained("google/recurrentgemma-2b", device_map="cpu")
input_texts = [
        'query: how much protein should a female eat',
        'query: summit define',
        "passage: As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
        "passage: Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments.",
    ]
# Tokenize the input texts
hf_tokenizer = AutoTokenizer.from_pretrained("google/recurrentgemma-2b")
batch_dict = hf_tokenizer(input_texts, max_length=128, pad_to_max_length=True, truncation=True, return_tensors='pt')

batch_dict_cpu = {k: v.cpu() for k, v in batch_dict.items()}

prompt = batch_dict_cpu['input_ids']
prompt_pos = torch.arange(prompt.shape[1]).unsqueeze(0).repeat(prompt.shape[0], 1)

model = model.eval()
hf_model = hf_model.eval()

dm_logits, cache = model.forward(prompt, prompt_pos, cache=None)
hf_logits = hf_model(**batch_dict_cpu)['logits']

print(dm_logits[:,:,:256000] - hf_logits)

```

### Expected behavior

Same forward path as the original repo

Comments:
Helllo, as of now and my latest tests, logits do not match exactly, but generations do. 
I invite you to remove the `nn.Conv1D` in favor of the non native on, which will produce the same logits. 
Also inviting you to compare generation, not logits, as Llama3 logits do not match either due to distribution shifts from the ROPE embedding

TLDR; while we aime to have 1e-4 logit equivalence, for RecurrentGemma, the model was very much slow in torch, and removing the custom convolution in favor of the native one improved generation speed. 
Inviting you to benchmark this as well 🤗 

This tests suit was run against the original codebase: https://github.com/huggingface/transformers/blob/deab6cff7e402ff47284e9a6aa02ce54e669ca28/tests/models/gemma/test_modeling_gemma.py#L504 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.