Obtaining logits during model.generate()
Author: QiyaoWei
Pull request: no

### System Info

- `transformers` version: 4.35.0
- Platform: Linux-6.2.0-39-generic-x86_64-with-glibc2.35
- Python version: 3.10.9
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.25.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

_No response_

### 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

What I am trying to do is explained in the title. As far as I understand, there are two ways of doing this

(1) pass in the output_logits=True parameter into model.generate() directly
(2) forward pass the model on the generated textual output

MWE here, is there something that I am not understanding correctly?

```
import random
random.seed("1234")
from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed, pipeline
set_seed(1234)

model_name="lvwerra/gpt2-imdb"
model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# I want to have the logits when I do generation
encoding = tokenizer(["Hi there, how are you?"], return_tensors="pt").to("cuda")
generation_output = model.generate(**encoding, return_dict_in_generate=True, output_logits=True)
# To be clear, this method returns the logits of the generated tokens, which does not include the prompt, but that is a minor detail

# obtain textual and logits output from gen
sequences = generation_output.sequences
sanity_check_logits = generation_output.logits

# we already have our encoding, now we can look at the logits from our models
dummy_var = {"input_ids": sequences}
model_output = model(**dummy_var)

# assert sanity_check_logits == model_output.logits fails
```

### Expected behavior

I expect the forward pass logits to be equal to the generated logits

Comments:
cc @gante @zucchini-nlp 

Hi @QiyaoWei 👋 Thank you for opening this issue 

Correct, we are aware of the issue -- https://github.com/huggingface/transformers/issues/29551. Let's focus all discussion there :)

@gante is right, but I also wanted to add that the way check is done in the minimal script is not correct. In this given setup there should be no difference between logits of `generate` and model `forward`, since we do not have anything in generation config that could have changed of the model output logits in-place. 

The below script shows that the logits are nearly equal :)
 
```python

from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

set_seed(1234)

model_id = "lvwerra/gpt2-imdb"
model = AutoModelForCausalLM.from_pretrained(model_id).to("cuda:0")
tok = AutoTokenizer.from_pretrained(model_id)
encoding = tok(["Hi there, how are you?"], return_tensors="pt").to("cuda:0")
generation_output = model.generate(**encoding, return_dict_in_generate=True, output_logits=True)

sequences = generation_output.sequences
sanity_check_logits = generation_output.logits

model_output = model(input_ids=sequences)

new_token_length = sequences.shape[-1] - encoding["input_ids"].shape[-1]
for i in range(1, new_token_length+1):
    gen_logits = sanity_check_logits[-i] # curr token logits (token at -i index in sequences)
    model_logits = model_output.logits[:, -i-1, :] # -1 because model has logits for the next token also
    diff = torch.max(torch.abs(gen_logits - model_logits)).cpu().item() # we can see how much is the difference
    print(f"{torch.allclose(gen_logits, model_logits)} {diff:.07f}") # allclose because generate uses kv cache, which might lead to small precision errors

```

That makes sense. Many thanks!