Change past_key_values shape of gpt_bigcode
Author: jiqing-feng
Pull request: yes

Hi @ArthurZucker @younesbelkada 

I found that the shape of `past_key_values` in `gpt_bigcode` is different from others, which will cause inconvenience or unexpected errors in `Optimum` and `Optimum-intel`. Therefore, I added some reshaping operations to make the shape insistent with other models. I have tested it on my device and didn't find any performance decay. Would you please help me review it? Thx!

BTW, could we add a restriction like the shape of `past_key_values` must be (batch_size, num_heads, sequence_length, head_dim)? It could avoid many issues if we have a fixed output and input shape of `past_key_values`

Comments:
WDYT @younesbelkada @ArthurZucker ?

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26402). All of your documentation changes will be reflected on that endpoint.

Sure. I see [bloom](https://github.com/huggingface/transformers/blob/main/src/transformers/models/bloom/modeling_bloom.py#L293-L299) and gpt-neox (see [_attn ](https://github.com/huggingface/transformers/blob/main/src/transformers/models/bloom/modeling_bloom.py#L293-L299)and [forward](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L164-L184)) also have this kind of operations. I think these operations only have slightly impact on performance, but it will keep input and output format insistent with others. Besides, I also remove some operations in `generation/util.py` since we use the standard format of `past_key_values`. Would like to hear opinions from BigCode team.

It's definitely there for performance reasons. The current format avoids a few unnecessary data copies and makes the kv cache concatenation faster. Not sure if it matters much though because the inference speed is almost always bottlenecked by either cpu or kv cache concatenation (if anything, the proposed changes would make both a bit worse).

In any case, I think the proposed changes should be treated as backward incompatible? So the only real option would be to add the kv cache format as an option?

> Thanks a lot for your review @jlamypoirier That makes sense, this might be backward in-compatible indeed @jiqing-feng regarding performance issues, would you be able to run some quick comparisons between transformers main branch and yours to determine how much is the overhead? But in any case it seems @jlamypoirier is right, the past key values shape gets changed so this might be backward incompatible, can you confirm that @jiqing-feng ?

Yes, Sure.

I tested on A100 and did not find any big difference. The script is as follows:

```python
import argparse
import time
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM


parser = argparse.ArgumentParser()
parser.add_argument("--model_id", default="bigcode/gpt_bigcode-santacoder", type=str, required=False)
parser.add_argument("--seq_len", default=32, type=int, required=False)
args = parser.parse_args()
model_id = args.model_id
seq_len = args.seq_len

name = model_id
name = "/home/bitsandbytes/hf_home/hub/models--bigcode--gpt_bigcode-santacoder"
model = AutoModelForCausalLM.from_pretrained(
  name,
  torch_dtype=torch.float16,
  use_cache=True,
  low_cpu_mem_usage=True,
  device_map="auto",
)

model = model.eval()
tokenizer = AutoTokenizer.from_pretrained(name)
input_sentence = "Example input"

generation_kwargs = dict(max_length=seq_len, min_length=seq_len, use_cache=True)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer, **generation_kwargs)
print(f"input sentence is {input_sentence}")
print(f"input tokens num is {len(tokenizer(input_sentence)['input_ids'])}")

for i in range(10):
    pre = time.time()
    out = generator(input_sentence, **generation_kwargs)
    print(f"Generate time costs {time.time()-pre} seconds")
    print(f"output = {out}")
    print(f"output token nums = {len(tokenizer(out[0]['generated_text'])['input_ids'])}")
```
You can just run `python test.py --seq_len 128` to double-check.

The performance results are shown in the following figure:
<img width="284" alt="image" src="https://github.com/huggingface/transformers/assets/107918818/e9e69ca4-9bd0-46f8-8f72-ec925a2140ae">


BTW, I am sorry that I can't understand backward incompatible. Does this mean backward broadcast in training?

Hi @younesbelkada @jlamypoirier . 

Would you please have a look at the results? Thx!

BTW, there are no backward incompatible issues since the pkv shape happened inside the codes. I also tested it before and after changing the pkv shape and there are no differences of outputs between them.

Hi @jiqing-feng 
Sorry for the delay replying on this matter

What I meant by a potential backward incompatibility is that the new past key values in this PR will have a different shape than the current model on main branch. If you can confirm that this is not the case then the PR is BC

I see this adds a ~5% overhead for seq_len = 512 and seems to grow linearly with the sequence length, perhaps with a longer sequence length & batch size the overhead will grow, therefore this PR introduces a regression which can be potentially important for long sequence length. However, given that this fixes some issues with external libs, we might consider merging this PR. I personally would prefer to avoid introducing additional overheads if possible, is there a way we could introduce a patch in optimum or optimum-intel ? Could you file an issue on optimum and tag me there to see what optimum maintainers think about it?

> Hi @jiqing-feng Sorry for the delay replying on this matter
> 
> What I meant by a potential backward incompatibility is that the new past key values in this PR will have a different shape than the current model on main branch. If you can confirm that this is not the case then the PR is BC
> 
> I see this adds a ~5% overhead for seq_len = 512 and seems to grow linearly with the sequence length, perhaps with a longer sequence length & batch size the overhead will grow, therefore this PR introduces a regression which can be potentially important for long sequence length. However, given that this fixes some issues with external libs, we might consider merging this PR. I personally would prefer to avoid introducing additional overheads if possible, is there a way we could introduce a patch in optimum or optimum-intel ? Could you file an issue on optimum and tag me there to see what optimum maintainers think about it?

Hi @younesbelkada . Thanks for your advice.

First, I confirm that there is no backward incompatible issue. Furthermore, we can see that it adds 5% overhead on both seq_len=32 and seq_len=512 so I think it is a fixed overhead and will not increase by the seq_len. Besides, most models have `reshape` or `transpose` operations for `past_key_values`. I think it is for consistency with other models even if it will have an impact on generation speed. For example, [gpt2](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py#L240). Maybe we can reshape the `attn_weight` instead of key-value in `_attn` function like gpt2.

For optimum, gpt_bigcode has a special onnx config, see [OnnxConfig](https://github.com/huggingface/optimum/blob/main/optimum/exporters/onnx/model_configs.py#L314) and [DummyPastKeyValuesGenerator](https://github.com/huggingface/optimum/blob/main/optimum/utils/input_generators.py#L840).

For optimum-intel, we also need a special process for gpt_bigcode in [BaseModelForCausalLM](https://github.com/huggingface/optimum-intel/blob/main/optimum/intel/generation/modeling.py#L281). 

There might be other external libs that rely on the shape of model inputs.

We can see external libs have special processes for gpt_bigcode, but we can avoid these kinds of changes if we can fix the shape of `past_key_values` for all models. It will help with code maintenance. WDYT?

hi @jiqing-feng thanks for providing more details!
Hmm I think here it is all about tradeoffs, IMO 5% slowdown is quite a lot since this means users that will update transformers will get a sudden slow down.
Regarding the fixes you shared, unless I am mistaken, this also means that those libraries will have to update again their patch for the next transformers version right? With that regard maybe the best is to keep things as they are, but if the bigcode team is happy to merge the PR I am not against it cc @ArthurZucker @jlamypoirier 

> hi @jiqing-feng thanks for providing more details! Hmm I think here it is all about tradeoffs, IMO 5% slowdown is quite a lot since this means users that will update transformers will get a sudden slow down. Regarding the fixes you shared, unless I am mistaken, this also means that those libraries will have to update again their patch for the next transformers version right? With that regard maybe the best is to keep things as they are, but if the bigcode team is happy to merge the PR I am not against it cc @ArthurZucker @jlamypoirier

Thanks for your advice. Yes, we have to update again if we change the shape of pkv.

The only benefit of keeping pkv shape the same is that it can improve the code maintenance of transformers and other libs like optimum and optimum-intel. It is true that it's all about trade-offs, we can keep things as they are if you think 5% performance is more important. Would like to hear your opinion, @ArthurZucker @jlamypoirier . Thx.

Hey @jiqing-feng I think our biggest concern when breaking backward compatibility is the incentive behind the change. In this case, consistency is not a big enough incentive, specifically given that we are also adding an overhead and other libraries out here probably had to update their codebase for this model. We'll be more careful with the cache when reviewing models @Rocketknight1 . Unless the community is super enthusiast, let's not merge this but feel free to keep it open!  

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.