Gemma 2 crashes with CUDA error on long context input
Author: waffoo
Pull request: no

### System Info

- `transformers` version: 4.42.3
- Platform: Linux-5.4.0-171-generic-x86_64-with-glibc2.31
- Python version: 3.10.14
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.3
- Accelerate version: 0.32.1
- PyTorch version (GPU?): 2.3.1+cu121 (True)
- Using distributed or parallel set-up in script?: No
- Using GPU in script?: Yes
- GPU type: NVIDIA A100 80GB PCIe

### 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

```python
import torch
from transformers import pipeline


def main():
    pipe = pipeline('text-generation',
                    model='google/gemma-2-9b-it',
                    model_kwargs={'torch_dtype': torch.bfloat16},
                    device=0)

    messages = [{
        "role": "user",
        "content": "What is the capital of France?" * 600,
    }]  # About 4k tokens. With *500, this code doesn't produce an error.

    outputs = pipe(messages, max_length=8000, do_sample=True)
    assistant_response = outputs[0]['generated_text'][-1]['content']
    print('response:', assistant_response)


if __name__ == '__main__':
    main()

```

### Expected behavior

I expected it to finish without crushing. The code snippet successfully worked for other models like `meta-llama/Meta-Llama-3-8B-Instruct` or `google/gemma-1.1-2b-it`.
However, with gemma-2 series, it produced an error when the input context length exceeded around 4k (not 8k).

```

.....

../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [16971,0,0], thread: [122,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [16971,0,0], thread: [123,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [16971,0,0], thread: [124,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [16971,0,0], thread: [125,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [16971,0,0], thread: [126,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [16971,0,0], thread: [127,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
Traceback (most recent call last):
...

  File "/home/user/gemma_main.py", line 16, in main
    outputs = pipe(messages, max_length=8000, do_sample=True)
  File "/home/user/venv/lib/python3.10/site-packages/transformers/pipelines/text_generation.py", line 257, in __call__
    return super().__call__(Chat(text_inputs), **kwargs)
  File "/home/user/venv/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1254, in __call__
    return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)
  File "/home/user/venv/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1261, in run_single
    model_outputs = self.forward(model_inputs, **forward_params)
  File "/home/user/venv/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1161, in forward
    model_outputs = self._forward(model_inputs, **forward_params)
  File "/home/user/venv/lib/python3.10/site-packages/transformers/pipelines/text_generation.py", line 349, in _forward
    generated_sequence = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
  File "/home/user/venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/user/venv/lib/python3.10/site-packages/transformers/generation/utils.py", line 1914, in generate
    result = self._sample(
  File "/home/user/venv/lib/python3.10/site-packages/transformers/generation/utils.py", line 2651, in _sample
    outputs = self(
  File "/home/user/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 1068, in forward
    outputs = self.model(
  File "/home/user/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 908, in forward
    layer_outputs = decoder_layer(
  File "/home/user/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 665, in forward
    hidden_states = self.post_feedforward_layernorm(hidden_states)
  File "/home/user/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 83, in forward
    output = self._norm(x.float())
  File "/home/user/venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 80, in _norm
    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

I'm new to LLMs so I'm sorry if this is not a bug.

Comments:
Same Error. + 1

Hey! That is because the Gemma2 cache is by default a static cache. This means that you do not have enough RAM to support the max 8000 tokens that you are pre-allocating your cache with! 

@ArthurZucker Thank you for the quick reply.
It doesn't seem the model uses up all the RAM or VRAM. A machine I'm using has plenty of memory (>500GB), and the model only consumes 20GB out of 80GB VRAM when it has 4k tokens (`*500` case in the example).
I also tested the code on another machine with an A100 40GB and got the same error at the same token count regardless of the VRAM size.

Is there any workaround for this issue? Is there some software limit that I can change?

Yeah you are right, the issue most probably not that (you would have a malloc error instead). 
I am pretty sure this comes from the sliding window cache (as you mention that 4K token does not have an issue. I'll investigate! Thanks for reporting

@ArthurZucker I have a similar error tied to a cache only occuring with an input of around 5k tokens running this : 

```python

    attn_implementation = 'eager'
    max_new_tokens=30
    model_id ="google/gemma-2-9b-it"

    pipe = pipeline("text-generation",  
                    model=model_id,
                    model_kwargs={'attn_implementation':attn_implementation,
                                  'torch_dtype':torch.bfloat16,},
                    device_map="auto",
                    max_new_tokens=max_new_tokens,
                    return_full_text =True,
                    add_special_tokens=True,
                    )

    res = [out["generated_text"][-1]['content'] for out in tqdm(pipe(dataset))]
                               
```

```
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [16388,0,0], thread: [29,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [16388,0,0], thread: [30,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [16388,0,0], thread: [31,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
Traceback (most recent call last):
  File "/home/tcharlot/inference/inference.py", line 57, in <module>
    main()
  File "/home/tcharlot/inference/inference.py", line 49, in main
    res = gemma.inference(dataset=dataset, verbose=True, **parameters)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/inference/gemma.py", line 42, in inference
    res = [out["generated_text"][-1]['content'] for out in tqdm(pipe(dataset))]
                                                                ^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/pipelines/text_generation.py", line 257, in __call__
    return super().__call__(Chat(text_inputs), **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/pipelines/base.py", line 1254, in __call__
    return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/pipelines/base.py", line 1261, in run_single
    model_outputs = self.forward(model_inputs, **forward_params)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/pipelines/base.py", line 1161, in forward
    model_outputs = self._forward(model_inputs, **forward_params)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/pipelines/text_generation.py", line 349, in _forward
    generated_sequence = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/generation/utils.py", line 1914, in generate
    result = self._sample(
             ^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/generation/utils.py", line 2651, in _sample
    outputs = self(
              ^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 1068, in forward
    outputs = self.model(
              ^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 908, in forward
    layer_outputs = decoder_layer(
                    ^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 650, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
                                                          ^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 252, in forward
    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/cache_utils.py", line 1071, in update
    return update_fn(
           ^^^^^^^^^^
  File "/home/tcharlot/.conda/envs/py311/lib/python3.11/site-packages/transformers/cache_utils.py", line 1046, in _static_update
    k_out[:, :, cache_position] = key_states
    ~~~~~^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```
maybe it can help, also it's running on a RTX 8000

I think this may be because of the following issues: #31781 
sliding_window is a value in cache_kwargs, rather than self.sliding_window. I set the sliding_window = cache_kwargs['sliding_window'] in the update function and it works without error.


that sounds right 


I confirmed that the latest release, transformers==4.42.4, has resolved the issue. Thank you so much!

I still got same issue with transformers version 4.43.3. 

> I think this may be because of the following issues: #31781 sliding_window is a value in cache_kwargs, rather than self.sliding_window. I set the sliding_window = cache_kwargs['sliding_window'] in the update function and it works without error.

When I tried to fix by this solution, it still gives me this error:

RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Edit: This problem alies when I trained google/gemma-2-9b with Supervised Fine Tuning and inference trained model.