fix static cache data type miss-match

[jiqing-feng]

Hi @SunMarc . This PR fixed the data type mismatch when using low-precision static cache. The following code can reproduce the bug:

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

[jiqing-feng]

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

[SunMarc]

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.

[SunMarc]

Thanks for the bug fix ! Left a comment

[SunMarc]

could you explain why this is needed for this particular model and why this doesn't happen for llama for example ? Many models have approximately the same modeling code.

[jiqing-feng]

For llama, we can see sin and cos come from position_embeddings (bf16 tensor) which comes from here. You can see llama's rotary embedding converts the data type. But for gptj, the position embeddings come from here, it set the data type to float32, so data type miss-match happens when the input data type is bf16 or fp16.

[jiqing-feng]

Hi @SunMarc . I left the comment 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!

[SunMarc]

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 ?

[ArthurZucker]

➕ 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!

[jiqing-feng]

➕ 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.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[jiqing-feng]

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

[SunMarc]

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

All good from my side. Pinging @ArthurZucker

[ArthurZucker]

Not completely sure we want to test for float32 as it's quite heavy

[SunMarc]

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

[ArthurZucker]

Sounds good then, merging!˜