past_key_values is not working as expected for falcon-7b
Author: orgadhadas
Pull request: no

### System Info

Hello,

I've been trying to use past_key_values to speed up text generation, but it doesn't seem to work: instead of generating coherent text like is done when I'm not using past_key_values, it seems to generate the same token over and over again. I've been trying to search the web for usage guidelines and it seemed to me like I'm doing everything correctly, but maybe I'm missing something.

Thank you!

### Who can help?

@ArthurZucker @younesbelkada - I think you're the relevant people for this.

### 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

```
from transformers import AutoTokenizer, AutoModelForCausalLM
import transformers
import torch
import pandas as pd
import pickle
from transformers import pipeline
device = "cuda" if torch.cuda.is_available() else "cpu"
model = "tiiuae/falcon-7b-instruct"

tokenizer = AutoTokenizer.from_pretrained(model, padding_side="left")
model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-7b-instruct", trust_remote_code=True).to(device)


# WITH past_key_values
def get_answer_from_model1(model, tokenizer, q, max_new_tokens=50):
    predicted_token_id = None
    prompt = q
    generated_text = ""
    n_new_tokens = 0
    past_key_values=None
    
    while (predicted_token_id != tokenizer.eos_token_id) and (n_new_tokens < max_new_tokens):
        if predicted_token_id is not None:
            model_input = tokenizer(predicted_token, return_tensors='pt').to(device)
        else:
            model_input = tokenizer(prompt, return_tensors='pt').to(device)
        with torch.no_grad():
            model_output = model(model_input['input_ids'], past_key_values=past_key_values)
        past_key_values = model_output['past_key_values']
        logits = model_output['logits']
        predicted_token_id = logits.argmax(-1)[0][-1]
        predicted_token = tokenizer.decode(predicted_token_id)
        
        if predicted_token_id != tokenizer.eos_token_id:
            prompt += predicted_token
            generated_text += predicted_token
            n_new_tokens += 1
        
    return generated_text

# WITHOUT past_key_values
def get_answer_from_model2(model, tokenizer, q, max_new_tokens=50):
    predicted_token_id = None
    prompt = q
    generated_text = ""
    n_new_tokens = 0
    past_key_values=None
    
    while (predicted_token_id != tokenizer.eos_token_id) and (n_new_tokens < max_new_tokens):
        model_input = tokenizer(prompt, return_tensors='pt').to(device)
        with torch.no_grad():
            model_output = model(model_input['input_ids'], past_key_values=past_key_values)
        logits = model_output['logits']
        predicted_token_id = logits.argmax(-1)[0][-1]
        predicted_token = tokenizer.decode(predicted_token_id)
        
        if predicted_token_id != tokenizer.eos_token_id:
            prompt += predicted_token
            generated_text += predicted_token
            n_new_tokens += 1
        
    return generated_text

q="hello"
answer1 = get_answer_from_model1(model, tokenizer, q)
print(answer1)
answer2 = get_answer_from_model2(model, tokenizer, q)
print(answer2)
```

### Expected behavior

answer1 and answer2 should be the same

Comments:
Hi @orgadhadas 
Thanks for the issue, I think the canonical way to use past key values is to set `use_cache=True` when calling `model.generate`. I think the remote code supports that argument as you can see here: https://huggingface.co/tiiuae/falcon-7b-instruct/blob/main/modelling_RW.py#L699 Can you share with us why you want to define a custom past key value mechanism?

I need to use the model.forward method, to get access to the hidden states computed during inference (for research). I hope this answers the question.

You have the `output_hidden_states` argument which should output the hidden states no? 

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.