Gemma2 Encounters RuntimeErrors when Using past_key_values
Author: Eric2i
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.44.0
- Platform: Linux-5.4.0-189-generic-x86_64-with-glibc2.31
- Python version: 3.12.2
- Huggingface_hub version: 0.24.5
- Safetensors version: 0.4.4
- Accelerate version: 0.33.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.0+cu121 (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?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA RTX A4000

### Who can help?

@ArthurZucker 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I am trying KV-Cache with Gemma2 but encountered a RuntimeError as follows: 
```bash
RuntimeError: (*bias): last dimension must be contiguous
```
To reproduce the bug, please see the code below:
```python
import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        model_id = "google/gemma-2-2b"

        tokenizer = AutoTokenizer.from_pretrained(
            model_id,
            device_map="cuda",
            torch_dtype=torch.float16,
        )

        tokenizer.pad_token = tokenizer.eos_token

        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            device_map="cuda",
            torch_dtype=torch.float16,
        )

        prompt = "Paris is the"
        input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")

        past_key_values = None
        from transformers import DynamicCache

        kv_cache = DynamicCache()
        next_token_id = input_ids
        generated_tokens = []

        for _ in range(7):
            with torch.no_grad():
                outputs = model(next_token_id, use_cache=True, past_key_values=kv_cache)

            logits = outputs.logits[:, -1, :]
            kv_cache = outputs.past_key_values

            next_token_id = logits.argmax(dim=-1)
            generated_tokens.append(next_token_id)
            next_token_id = next_token_id.unsqueeze(0)
            # breakpoint()

        generated_tokens = tokenizer.batch_decode(generated_tokens)
        print(generated_tokens)
```
On the other hand, if I set 'model_id = "google/gemma-2b" (or meta-llama/Meta-Llama-3.1-8B)', the program ran smoothly without any errors.

### Expected behavior

The program is expected to generate something like: 
```bash
[' capital', ' of', ' France', ' and', ' the', ' most', ' populous']
```
which is the results outputted by gemma-2b with the similar code demonstrated above.

Comments:
Hey! Gemma2 is special and requires a specific cache type, Hybrid Cache (see more: https://huggingface.co/docs/transformers/v4.44.0/en/internal/generation_utils#transformers.HybridCache)

Also, if you are doing a custom cached generation loop with Gemma2, you have to pass correct `cache_position`. Not passing cache positions will raise an error if you install transformers from source. For more details, we have docs for cache and how to implement a custom generation loop che-object-work

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "google/gemma-2-2b"

tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    device_map="cuda",
    torch_dtype=torch.float16,
)

tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="cuda",
    torch_dtype=torch.float16,
)

prompt = "Paris is the"
input_ids = tokenizer(prompt, return_tensors="pt")

from transformers.cache_utils import DynamicCache, HybridCache

kv_cache = HybridCache(config=model.config, max_batch_size=1, max_cache_len=100, device="cuda", dtype=torch.float16)
next_token_id = input_ids.input_ids.to("cuda")
attention_mask = input_ids.attention_mask.to("cuda")
cache_position = torch.arange(attention_mask.shape[1], device="cuda")
generated_tokens = []

for _ in range(7):
    with torch.no_grad():
        outputs = model(next_token_id, attention_mask=attention_mask, cache_position=cache_position, use_cache=True, past_key_values=kv_cache)
    logits = outputs.logits[:, -1, :]
    kv_cache = outputs.past_key_values
    next_token_id = logits.argmax(dim=-1)
    generated_tokens.append(next_token_id)
    next_token_id = next_token_id.unsqueeze(0)
    attention_mask = torch.cat([attention_mask, torch.ones(1, 1, device="cuda")], dim=-1)
    cache_position = cache_position[-1:] + 1
    # breakpoint()

generated_tokens = tokenizer.batch_decode(generated_tokens)
print(generated_tokens)

```

Thanks, Raushan.