Inconsistent Output with and without Prompt Caching in Llama-3.1-8B-Instruct.
Author: giulio98
Pull request: no
Labels: name: bug

name: Generation

name: Cache

### System Info

- `transformers` version: 4.45.1
- Platform: Linux-6.6.35-amd64-x86_64-with-glibc2.35
- Python version: 3.11.6
- Huggingface_hub version: 0.25.2
- Safetensors version: 0.4.3
- Accelerate version: 1.0.1
- Accelerate config: 	- compute_environment: LOCAL_MACHINE
	- distributed_type: MULTI_GPU
	- mixed_precision: no
	- use_cpu: False
	- debug: False
	- num_processes: 2
	- machine_rank: 0
	- num_machines: 1
	- gpu_ids: 0,1
	- rdzv_backend: static
	- same_network: True
	- main_training_function: main
	- enable_cpu_affinity: False
	- downcast_bf16: no
	- tpu_use_cluster: False
	- tpu_use_sudo: False
	- tpu_env: []
- PyTorch version (GPU?): 2.4.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: NO
- Using GPU in script?: YES
- GPU type: NVIDIA H100 80GB HBM3

### Who can help?

@gante @ArthurZucker @itaza

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. Generate Responses with Cache, following [Re-use Cache to continue generation](https://huggingface.co/docs/transformers/en/kv_cache#re-use-cache-to-continue-generation)
```python
import copy
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, StaticCache

model_id = "meta-llama/Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Init StaticCache with big enough max-length
prompt_cache = StaticCache(config=model.config, max_batch_size=1, max_cache_len=1024, device="cuda", dtype=torch.bfloat16)

INITIAL_PROMPT = "You are a helpful assistant. "
inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to("cuda")

with torch.no_grad():
    prompt_cache = model(**inputs_initial_prompt, past_key_values=prompt_cache).past_key_values
prompts = ["Help me to write a blogpost about travelling.", "What is the capital of France?"]
responses = []
for prompt in prompts:
    new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to("cuda")
    past_key_values = copy.deepcopy(prompt_cache)
    outputs = model.generate(**new_inputs, past_key_values=past_key_values, max_new_tokens=20)
    response = tokenizer.batch_decode(outputs)[0]
    responses.append(response)

print(responses)
```
2. Observed Output
```python
['
You are a helpful assistant. Help me to write a blogpost about travelling.  I have some ideas, but I’ts not clear how to structure the post.  I',
 '
You are a helpful assistant. What is the capital of France? Paris.  is the capital of the United States? Washington D.C.  is the capital of']
```
3. Generate response without cache
```python
responses = []
for prompt in prompts:
    new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(**new_inputs, max_new_tokens=20, use_cache=False)
    response = tokenizer.batch_decode(outputs)[0]
    responses.append(response)

print(responses)
```
4. Observed Output
```python
['
You are a helpful assistant. Help me to write a blogpost about travelling. Here’s what I need to write about:\nTitle: “The Magic of Exploring New Places:',
 '
You are a helpful assistant. What is the capital of France? Paris.\nWhat is the capital of Australia? Canberra.\nWhat is the capital of Brazil? Brasília']
```

### Expected behavior

The output without cache should be exactly the same as the one that uses the cache.

Comments:
Hi @giulio98 👋 

> The output without cache should be exactly the same as the one that uses the cache.

This statement is not true, see [this comment](https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535) :) However, the results should be similar, which doesn't seem to be the case -- I'm going to have a look.

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.