fix static cache data type miss-match
Author: jiqing-feng
Pull request: yes

Hi @SunMarc . This PR fixed the data type mismatch when using low-precision static cache. The following code can reproduce the bug:
```python
import torch
from transformers import pipeline

model_id = "EleutherAI/gpt-j-6b"
model_kwargs = {"torch_dtype": torch.bfloat16}

pipe = pipeline("text-generation", model=model_id, model_kwargs=model_kwargs)

generation_config = pipe.model.generation_config
generation_config.cache_implementation="static"

print(pipe("I am happy because", generation_config=generation_config))
```

Output:
```
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Traceback (most recent call last):
  File "/home/jiqingfe/test_gptj.py", line 13, in <module>
    print(pipe("I am happy because", generation_config=generation_config))
  File "/home/jiqingfe/transformers/src/transformers/pipelines/text_generation.py", line 272, in __call__
    return super().__call__(text_inputs, **kwargs)
  File "/home/jiqingfe/transformers/src/transformers/pipelines/base.py", line 1301, in __call__
    return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)
  File "/home/jiqingfe/transformers/src/transformers/pipelines/base.py", line 1308, in run_single
    model_outputs = self.forward(model_inputs, **forward_params)
  File "/home/jiqingfe/transformers/src/transformers/pipelines/base.py", line 1208, in forward
    model_outputs = self._forward(model_inputs, **forward_params)
  File "/home/jiqingfe/transformers/src/transformers/pipelines/text_generation.py", line 370, in _forward
    generated_sequence = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
  File "/opt/conda/envs/idp/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/home/jiqingfe/transformers/src/transformers/generation/utils.py", line 2263, in generate
    result = self._beam_search(
  File "/home/jiqingfe/transformers/src/transformers/generation/utils.py", line 3472, in _beam_search
    outputs = self(**model_inputs, return_dict=True)
  File "/opt/conda/envs/idp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/idp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jiqingfe/transformers/src/transformers/models/gptj/modeling_gptj.py", line 1098, in forward
    transformer_outputs = self.transformer(
  File "/opt/conda/envs/idp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/idp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jiqingfe/transformers/src/transformers/models/gptj/modeling_gptj.py", line 838, in forward
    outputs = block(
  File "/opt/conda/envs/idp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/idp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jiqingfe/transformers/src/transformers/models/gptj/modeling_gptj.py", line 453, in forward
    attn_outputs = self.attn(
  File "/opt/conda/envs/idp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/idp/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jiqingfe/transformers/src/transformers/models/gptj/modeling_gptj.py", line 246, in forward
    key, value = layer_past.update(key, value, self.layer_idx, cache_kwargs)
  File "/home/jiqingfe/transformers/src/transformers/cache_utils.py", line 1220, in update
    k_out.index_copy_(2, cache_position, key_states)
RuntimeError: index_copy_(): self and source expected to have the same dtype, but got (self) BFloat16 and (source) Float
```


Comments:
BTW, I suppose transformers missed some static cache tests, do you have any instructions about where can I add this kind of test? Thanks!

> BTW, I suppose transformers missed some static cache tests, do you have any instructions about where can I add this kind of test? Thanks!

All the tests related to the cache are in test_utils.py file. inside `GenerationTesterMixin`, you will find the test we perform on all models and in `GenerationIntegrationTests`, these are integration tests. 

Hi @SunMarc . I left the [comment ](https://github.com/huggingface/transformers/pull/34799#discussion_r1849382212)to explain why llama model doesn't have this issue. BTW, I also added the low-precision static cache tests to avoid this kind of issue in the future, please review it. Thanks!

> ➕ on marc's comment. The safe way to do this is to cast key and query to the cache's dtype no? And do this in the `cache_utils` rather than at the modeling level!


Yes, I have applied your suggestions, thanks!



> Thanks for the explanation ! LGTM ! Did you run the static cache tests you added to see if there are other models that requires this fix ?

The CI already contains the tests that I changed, so currently no other models require it. Besides, I have changed it into cache_utils which should be applied for all language models with static cache.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_34799). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Hi @SunMarc , please review the new changes, thanks!

> Hi @SunMarc , please review the new changes, thanks!

All good from my side. Pinging @ArthurZucker 

I think it was testing for float32 initally and @jiqing-feng added coverage for float16 @ArthurZucker 

Sounds good then, merging!˜