`Gemma2Model` not returning cache
Author: naimenz
Pull request: no

### System Info

- `transformers` version: 4.42.4
- Platform: macOS-13.4-arm64-arm-64bit
- Python version: 3.10.13
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.3
- Accelerate version: 0.29.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.2 (False)
- Tensorflow version (GPU?): 2.16.1 (False)
- 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

### 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

`Gemma2Model` is not returning `past_key_values` even with `use_cache=True`. The code snippet below reproduces the issue and compares it to `GemmaModel`.

From a quick look at the code in `modeling_gemma2.py`, `Gemma2Model.forward` doesn't have any of the cache handling in `GemmaModel.forward`, and in fact never modifies the value of `past_key_values`.

```python
from transformers import AutoTokenizer, AutoConfig, AutoModel

model_names = ["google/gemma-1.1-7b-it", "google/gemma-2-9b-it"]
for model_name in model_names:
    print(f"=====\n{model_name}")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    config = AutoConfig.from_pretrained(model_name)
    config.intermediate_size=10
    config.hidden_size=32
    config.use_cache = True
    model = AutoModel.from_config(config)
    print(f"{type(model) = }")
    inps = tokenizer("Hello, my dog is cute", return_tensors="pt")
    out = model(**inps, use_cache=True)
    print(f"{type(out) = }")
    print(f"{out.past_key_values = }")

# Outputs:
# =====
# google/gemma-1.1-7b-it
# type(model) = <class 'transformers.models.gemma.modeling_gemma.GemmaModel'>
# type(out) = <class 'transformers.modeling_outputs.BaseModelOutputWithPast'>
# out.past_key_values = ((tensor([[[[-0.1430,  0.0217,  0.0094,  ...,  0.0696,  0.2129,  0.0378],
#           [ 0.1115,  0.0830,  0.0562,  ...,  0.0123,  0.0671, -0.1179],
#           [-0.0751,  0.0045,  0.1079,  ...,  0.1013,  0.1240, -0.0283],
#           ...,
# =====
# google/gemma-2-9b-it
# type(model) = <class 'transformers.models.gemma2.modeling_gemma2.Gemma2Model'>
# type(out) = <class 'transformers.modeling_outputs.BaseModelOutputWithPast'>
# out.past_key_values = None
```

### Expected behavior

I expect the `past_key_values` to be returned, as in `GemmaModel` and other `transformers` model families.

Comments:
Hey! This is related to the new cache format. You need to pass a `Cache` object to the model! 🤗 
The main is reason is that what you are doing is a forward pass that does not need the cache per say, if you need it back it's modified on the fly ! 

@ArthurZucker Ah I see, thanks! For my use case we do actually use the cache on forward passes like this: we want to get logits on a bunch of sequences with a common prefix, like in this [GCG implementation](https://github.com/centerforaisafety/HarmBench/blob/main/baselines/gcg/gcg.py). In particular we run a sequence through once, store the correspondence between the input tokens and the cache values, and then on subsequent runs look for a common prefix between the cached tokens and the new values and return as much of the cache as matches. 

Does this mean that the current API for caching is going to break for older models as well, and I should try to adapt my code to the `Cache` object? Or would it be possible to support the current API for `Gemma2` as well?

One note about trying to use the new `DynamicCache`: if I use the same snippet above but add a `DynamicCache`,
```python
    # out = model(**inps, use_cache=True)
    out = model(**inps, use_cache=True, past_key_values=DynamicCache())
```
 then I get this output:
```
=====
google/gemma-1.1-7b-it
type(model) = <class 'transformers.models.gemma.modeling_gemma.GemmaModel'>
type(out) = <class 'transformers.modeling_outputs.BaseModelOutputWithPast'>
out.past_key_values = DynamicCache()
=====
google/gemma-2-9b-it
type(model) = <class 'transformers.models.gemma2.modeling_gemma2.Gemma2Model'>
Traceback (most recent call last):
  File "/Users/ian/code/farai/robust-llm/robust_llm/SCRATCH/gemma2_repro_scratch.py", line 15, in <module>
    out = model(**inps, use_cache=True, past_key_values=DynamicCache())
  File "/Users/ian/code/farai/robust-llm/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/Users/ian/code/farai/robust-llm/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/ian/code/farai/robust-llm/venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 881, in forward
    causal_mask = self._update_causal_mask(
  File "/Users/ian/code/farai/robust-llm/venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 972, in _update_causal_mask
    causal_mask = torch.full(
TypeError: full(): argument 'size' failed to unpack the object at pos 2 with error "type must be tuple of ints,but got NoneType"
```

We are indeed going to break previous API but just the tuple part. If you init the cache outside the model it should work as expected. 


can you give me the full reproducer? 

To fix this issue, you need to update the `Gemma2Model`'s `forward` method to handle `past_key_values` correctly.

> **May be this helps:**

1. **Identify the Issue**: The `forward` method in `Gemma2Model` needs to manage `past_key_values`.
2. **Modify the `forward` Method**: Add logic to handle cache similar to the `GemmaModel`.

Example of how you might modify the `Gemma2Model` to include the necessary cache handling:

### Example

```python
# Example modification in modeling_gemma2.py

class Gemma2Model(PreTrainedModel):
    def forward(self, input_ids, past_key_values=None, use_cache=None, **kwargs):
        # Your existing forward logic
        # ...

        if use_cache:
            if past_key_values is None:
                past_key_values = self.init_cache()  # Initialize cache if not provided
            # Implement caching logic here
            # ...
            # Ensure past_key_values are properly updated and returned
            return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values)

        return BaseModelOutput(last_hidden_state=hidden_states)
    
    def init_cache(self):
        # Initialize cache logic
        # ...
        return past_key_values

# Updated script to test the fix
from transformers import AutoTokenizer, AutoConfig, AutoModel

model_names = ["google/gemma-1.1-7b-it", "google/gemma-2-9b-it"]
for model_name in model_names:
    print(f"=====\n{model_name}")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    config = AutoConfig.from_pretrained(model_name)
    config.intermediate_size = 10
    config.hidden_size = 32
    config.use_cache = True
    model = AutoModel.from_config(config)
    print(f"{type(model) = }")
    inps = tokenizer("Hello, my dog is cute", return_tensors="pt")
    out = model(**inps, use_cache=True)
    print(f"{type(out) = }")
    print(f"{out.past_key_values = }")
```

- **Initialize Cache**: Add a method to initialize `past_key_values`.
- **Modify Forward Method**: Update the `forward` method to manage `past_key_values` correctly and return them when `use_cache` is `True`.

By using these changes, the `Gemma2Model` should correctly handle and return `past_key_values` similar to the `GemmaModel`.

No I think the best way to do this is just to have a simple function that initialized the cache outside the model. Something like `cache = DynamicCache()` note that the cache for gemma is `HybridCache()`. Then you can just re-use it / create copies of it

> can you give me the full reproducer?

@ArthurZucker Yes sorry, it's just commenting out that line and adding the other:
```python
from transformers import AutoTokenizer, AutoConfig, AutoModel, DynamicCache

model_names = ["google/gemma-1.1-7b-it", "google/gemma-2-9b-it"]
for model_name in model_names:
    print(f"=====\n{model_name}")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    config = AutoConfig.from_pretrained(model_name)
    config.intermediate_size=10
    config.hidden_size=32
    config.use_cache = True
    model = AutoModel.from_config(config)
    print(f"{type(model) = }")
    inps = tokenizer("Hello, my dog is cute", return_tensors="pt")
    # out = model(**inps, use_cache=True)
    out = model(**inps, use_cache=True, past_key_values=DynamicCache())
    print(f"{type(out) = }")
    print(f"{out.past_key_values = }")
```
This gives the following output:
```
=====
google/gemma-1.1-7b-it
type(model) = <class 'transformers.models.gemma.modeling_gemma.GemmaModel'>
type(out) = <class 'transformers.modeling_outputs.BaseModelOutputWithPast'>
out.past_key_values = DynamicCache()
=====
google/gemma-2-9b-it
type(model) = <class 'transformers.models.gemma2.modeling_gemma2.Gemma2Model'>
Traceback (most recent call last):
  File "/root/robust-llm/SCRATCH/gemma2_repro.py", line 15, in <module>
    out = model(**inps, use_cache=True, past_key_values=DynamicCache())
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 881, in forward
    causal_mask = self._update_causal_mask(
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 972, in _update_causal_mask
    causal_mask = torch.full(
TypeError: full(): argument 'size' failed to unpack the object at pos 2 with error "type must be tuple of ints,but got NoneType"
```
If I try to use a `HybridCache` with `gemma-1.1` I get
```
TypeError: HybridCache.__init__() missing 3 required positional arguments: 'config', 'max_batch_size', and 'max_cache_len'
```
and I don't know what reasonable values are for these ahead of time; the old way I'd just get back `past_key_values` for whatever I put in, which was very useful. Plus I don't want to do any kind of sliding window, I just want a basic cache.

`HybridCache(model.config, inps["input_ids"].shape, max_length)` are good values. 
Got it for the Dynamic cache, it would be nice to have cc @zucchini-nlp if you have time to have look! 

I'm trying this out; a few minor fixes were needed:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.cache_utils import HybridCache

model_name = 'google/gemma-2-9b-it'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

inps = tokenizer("The capital of France is", return_tensors="pt")
cache = HybridCache(config=model.config, max_batch_size=inps["input_ids"].shape[0], max_cache_len=10)
out = model(**inps, use_cache=True, past_key_values=cache)
```

That works. Now let's try to use it in a mini-generation context:

```python
next_token = out.logits[0, -1].argmax()
print(tokenizer.decode(next_token.item()))
# Add on that token and ask for the next one.
out2 = model(next_token.expand(1, 1), use_cache=True, past_key_values=cache)
next_next_token = out2.logits[0, -1].argmax()
print(tokenizer.decode(next_next_token.item()))
```

Both `print` statements output `Paris`. I'd expect the second one to output `.` or something based on how I thought cache was supposed to work. Is there more different about the Gemma2 cache impl than meets the eye?

Trying to pass the precomputed cache to `generate` got me in even more trouble, so I decided I'd start back with a single generation step.

A clearer example:

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.cache_utils import HybridCache, DynamicCache

model_name = 'google/gemma-2-9b-it'
#model_name = 'google/gemma-1.1-2b-it'

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

inps = tokenizer("1 3 2 4 5", return_tensors="pt")
cache = HybridCache(config=model.config, max_batch_size=inps["input_ids"].shape[0], max_cache_len=10)
#cache = DynamicCache()
out = model(**inps, use_cache=True, past_key_values=cache)

next_token = out.logits[0, -1].argmax()
print(next_token, repr(tokenizer.decode(next_token.item())))

concatted = torch.cat([inps["input_ids"], next_token.unsqueeze(0).unsqueeze(0)], dim=-1)

out2_manual = model(concatted)
out2_opt = model(next_token.unsqueeze(0).unsqueeze(0), use_cache=True, past_key_values=cache)
assert torch.allclose(out2_manual.logits[0, -1], out2_opt.logits[0, -1])
```

The assertion fails on `transformers` 4.43.1. Uncomment the lines to use Gemma 1 instead, though, and the assertion succeeds.

@kcarnold seems like Gemma2 doesn't work correctly with iterative forward, since it breaks a few things compared to the prev Gemma. I will make a PR soon to fix it, in the meanwhile you can make the above code work by preparing `cache_position` outside and passing it to forward as follows

It will output max diff of `3.8147e-05` between two logits

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.cache_utils import HybridCache, DynamicCache

model_name = 'google/gemma-2-9b-it'
# model_name = 'google/gemma-1.1-2b-it'

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cuda")

inps = tokenizer("1 3 2 4 5", return_tensors="pt").to("cuda")
cache_position = torch.arange(inps.input_ids.shape[1], device="cuda", dtype=torch.int64)
cache = HybridCache(config=model.config, max_batch_size=inps["input_ids"].shape[0], max_cache_len=20)
# cache = DynamicCache()

with torch.no_grad():
    out = model(inps.input_ids, cache_position=cache_position, use_cache=True, past_key_values=cache)

next_token = out.logits[0, -1].argmax()
cache_position_next = cache_position[-1:] + 1
print(next_token, repr(tokenizer.decode(next_token.item())))

concatted = torch.cat([inps["input_ids"], next_token.unsqueeze(0).unsqueeze(0)], dim=-1)
cache_position_concatted = torch.cat([cache_position, cache_position_next], dim=-1)

with torch.no_grad():
    out2_manual = model(concatted, cache_position=cache_position_concatted)
    out2_opt = model(next_token.unsqueeze(0).unsqueeze(0), cache_position=cache_position_next, use_cache=True, past_key_values=cache)

    print((out2_manual.logits[0, -1] - out2_opt.logits[0, -1]).abs().max())
```



Thanks! It would be great to have a model-agnostic API for this. Is it intentional that `Cache` objects don't keep track of `cache_position` internally?

(My real goal is to be able to pass that cache off to a `generate`, but I'm hitting all kinds of issues with that so I figured I'd start here.)

Also note that I kinda stepped on this Issue; resolving my sub-issue won't actually resolve the OP's issue (which affects me too).

It is intentional: torch.compile does not support tensor updates in the graph. 
I don't think we tried recently with marking the adresse of the cache as static but AFAIK, it's incompatible 

@kcarnold when you ran `out2_opt = model(next_token.unsqueeze(0).unsqueeze(0), use_cache=True, past_key_values=cache)` the second time, you did not pass `cache_position=torch.tensor([input_ids.shape[-1])`