Can not use prompt tuning inference
Author: aohenuo
Pull request: no
Labels: name: bug

### System Info

transformers version: 4.48.0
Platform: Linux-4.15.0-76-generic-x86_64-with-glibc2.27
Python version: 3.9.0
Huggingface_hub version: 0.23.4
Safetensors version: 0.4.3
Accelerate version: 0.34.2
Accelerate config: - compute_environment: LOCAL_MACHINE
- distributed_type: DEEPSPEED
- mixed_precision: fp16
- use_cpu: False
- debug: False
- num_processes: 4
- machine_rank: 0
- num_machines: 1
- rdzv_backend: static
- same_network: True
- main_training_function: main
- enable_cpu_affinity: False
- deepspeed_config: {'gradient_accumulation_steps': 16, 'gradient_clipping': 1.0, 'offload_optimizer_device': 'cpu', 'offload_param_device': 'cpu', 'zero3_init_flag': False, 'zero3_save_16bit_model': False, 'zero_stage': 3}
- downcast_bf16: no
- tpu_use_cluster: False
- tpu_use_sudo: False
- tpu_env: []
PyTorch version (GPU?): 2.3.1+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?:
Using GPU in script?:
GPU type: NVIDIA GeForce RTX 3090

### Who can help?

@SunMarc @ArthurZucker 

### Information

- [ ] The official example scripts
- [x] 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
```
import torch
import json
import re
import random
import argparse
from datetime import datetime
from datasets import load_dataset
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM
import os
from peft import PeftModel
import numpy as np


def parse_args():
    parser = argparse.ArgumentParser(description="Llama3.1 8B for gsm 8k")
    parser.add_argument("--num_shots", type=int, default=0,
                       help="use for incontext learning, 0 domain zero-shot")
    parser.add_argument("--num_samples", type=int, default=None,
                       help="test num samples")
    parser.add_argument("--output_dir", type=str, default="./results",)
    parser.add_argument("--temperature", type=float, default=0.0,)
    parser.add_argument("--max_new_tokens", type=int, default=512,)
    return parser.parse_args()

def main():
  
    args = parse_args()
    for k,v in sorted(vars(args).items()):
        print(k,'=',v)
    
    MODEL_PATH = "meta-llama/Llama-3.1-8B-Instruct"
    DATA_PATHS = {
        "test": "openai/gsm8k/test-00000-of-00001.parquet",
        "train": "openai/gsm8k/train-00000-of-00001.parquet"
    }

    os.environ["CUDA_VISIBLE_DEVICES"] = "7" 
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

    # load the model
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    tokenizer.pad_token = tokenizer.eos_token

    basemodel = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.bfloat16,
    device_map=DEVICE 
    )

    
    model = PeftModel.from_pretrained(basemodel, "./peft_gsm8k")

    def load_data(path):
        return load_dataset("parquet", data_files=path)["train"]
    
    train_examples = []
    if args.num_shots > 0:
        print(f"{args.num_shots}-shot...")
        train_data = load_data(DATA_PATHS["train"])
        train_examples = [{"question": ex["question"], "answer": ex["answer"]} 
                        for ex in train_data]
        print(f"the icl select set size is {len(train_examples)}")

    # load test data
    test_data = load_data(DATA_PATHS["test"])
    test_examples = [{"question": ex["question"], "answer": ex["answer"]} 
                   for ex in test_data]
    
    if args.num_samples:
        test_examples = random.sample(test_examples, args.num_samples)
    total_samples = len(test_examples)

    def extract_answer(text):
        try: 
            matches = re.findall(r"\\boxed{([\d,]+)}", text)
            if not matches:
                matches = re.findall(r"\b\d{1,3}(?:,\d{3})*\b", text)
            return matches[-1].replace(',', '') if matches else None
        except:
            return None
    def build_prompt(question):
        messages = []
        
        if args.num_shots > 0:
            examples = random.sample(train_examples, args.num_shots)
            for ex in examples:
                messages.append({"role": "user", "content": ex["question"]})
                messages.append({"role": "assistant", "content": ex["answer"]})
        
   
        if args.num_shots == 0:
            content = f"Please solve this math problem, and surround the final answer with\\boxed{{}}.question:{question}"
        else:
            content = question
            
        messages.append({"role": "user", "content": content})
        
        return tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")  
    result_filename = f"prompt_results_{args.num_shots}shot_{timestamp}.json"
    
    os.makedirs(args.output_dir, exist_ok=True) 
    output_path = os.path.join(args.output_dir, result_filename)  
    
    results = {
        "metadata": {
            "model": MODEL_PATH,
            "num_shots": args.num_shots,
            "num_samples": total_samples,
            "temperature": args.temperature,
            "max_new_tokens": args.max_new_tokens,
            "test_time": timestamp
        },
        "cases": []
    }

    # generate the answer of gsm8k 
    progress_bar = tqdm(test_examples, desc="progerss")
    correct = 0
    
    for example in progress_bar:
        case_result = {
            "question": example["question"],
            "generated": "",
            "pred_answer": None,
            "true_answer": None,
            "is_correct": False,
            "error": None
        }
        
        try:
            # build the prompt
            prompt = build_prompt(example["question"])
            
            # generate the answer
            inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
            outputs = model.generate(
                inputs.input_ids,
                max_new_tokens=args.max_new_tokens,
                temperature=args.temperature,
                do_sample=args.temperature > 0,
                pad_token_id=tokenizer.eos_token_id
            )
            
            # decode the answer
            generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
            pred_answer = extract_answer(generated_text)
            true_answer = extract_answer(example["answer"])
            
            case_result.update({
                "generated": generated_text,
                "pred_answer": pred_answer,
                "true_answer": true_answer,
                "is_correct": pred_answer == true_answer
            })
            
            if case_result["is_correct"]:
                correct += 1
                
        except Exception as e:
            case_result["error"] = str(e)
            print(f"\Error: {str(e)}")
        
        results["cases"].append(case_result)
        progress_bar.set_postfix(accuracy=f"{correct/(progress_bar.n+1)*100:.1f}%")
        
      
        if (progress_bar.n + 1) % 10 == 0:
            
            os.makedirs(args.output_dir, exist_ok=True)  
            with open(output_path, "w") as f:
                json.dump(results, f, indent=2, ensure_ascii=False)

    os.makedirs(args.output_dir, exist_ok=True)  
    with open(output_path, "w") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    
    accuracy = correct / total_samples * 100
    print(f"\n ACC:{accuracy:.2f}% ({correct}/{total_samples})")
    print(f"save dir:{output_path}")

if __name__ == "__main__":
    main()
```
### Expected behavior

When I run my inference code, the scripts get the following error:
```
/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/transformers/generation/configuration_utils.py:628: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.0` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`.
  warnings.warn(
/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/transformers/generation/configuration_utils.py:633: UserWarning: `do_sample` is set to `False`. However, `top_p` is set to `0.9` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_p`.
  warnings.warn(
The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.
progess   0%|                                                                                                                                                                                      | 0/1319 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "/data/syk/pro/prompt_test.py", line 204, in <module>
    main()
  File "/data/syk/pro/prompt_test.py", line 159, in main
    outputs = model.generate(
  File "/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/peft/peft_model.py", line 1192, in generate
    outputs = self.base_model.generate(**kwargs)
  File "/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/transformers/generation/utils.py", line 2255, in generate
    result = self._sample(
  File "/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/transformers/generation/utils.py", line 3247, in _sample
    model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
  File "/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/peft/peft_model.py", line 1221, in prepare_inputs_for_generation
    if model_kwargs["past_key_values"][0][0].shape[-2] >= model_kwargs["input_ids"].shape[1]:
  File "/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/transformers/cache_utils.py", line 390, in __getitem__
    raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}")
KeyError: 'Cache only has 0 layers, attempted to access layer with index 0'
```

This error seems to be related to the conflict between cache and learnable soft prompt, so how to solve this problem?

Comments:
@aohenuo This error occurs because the model is trying to access past_key_values, but the cache is empty. Since you are using a PEFT model, caching may not be handled correctly. One possible solution is to modify your generate() function by adding use_cache=False. Let me know if this works for you.

> [@aohenuo](https://github.com/aohenuo) This error occurs because the model is trying to access past_key_values, but the cache is empty. Since you are using a PEFT model, caching may not be handled correctly. One possible solution is to modify your generate() function by adding use_cache=False. Let me know if this works for you.

When I set user_cache=False

```
outputs = model.generate(
        inputs.input_ids,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        do_sample=args.temperature > 0,
        pad_token_id=tokenizer.eos_token_id,
        use_cache=False,   
            )
```

 here is an another error:  

```
Traceback (most recent call last):
  File "/data/syk/pro/prompt_test.py", line 205, in <module>
    main()
  File "/data/syk/pro/prompt_test.py", line 159, in main
    outputs = model.generate(
  File "/data/syk/miniconda3/envs/debias/lib/python3.9/site-packages/peft/peft_model.py", line 1192, in generate
    outputs = self.base_model.generate(**kwargs)
  File "/data/syk/miniconda3/envs/debias/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/data/syk/miniconda3/envs/debias/lib/python3.9/site-packages/transformers/generation/utils.py", line 2223, in generate
    result = self._sample(
  File "/data/syk/miniconda3/envs/debias/lib/python3.9/site-packages/transformers/generation/utils.py", line 3204, in _sample
    model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
  File "/data/syk/miniconda3/envs/debias/lib/python3.9/site-packages/peft/peft_model.py", line 1217, in prepare_inputs_for_generation
    if uses_cache and (model_kwargs["past_key_values"] is not None):
KeyError: 'past_key_values'
```


the past_key_values still exists

Hi @aohenuo! Try this modified code. 

```python
outputs = model.generate(
    inputs.input_ids,
    max_new_tokens=args.max_new_tokens,
    temperature=args.temperature,
    do_sample=args.temperature > 0,
    pad_token_id=tokenizer.eos_token_id,
    use_cache=True,  # Keep cache enabled
    past_key_values=None  # Explicitly set past_key_values
)
```

> Hi [@aohenuo](https://github.com/aohenuo)! Try this modified code.
> 
> outputs = model.generate(
>     inputs.input_ids,
>     max_new_tokens=args.max_new_tokens,
>     temperature=args.temperature,
>     do_sample=args.temperature > 0,
>     pad_token_id=tokenizer.eos_token_id,
>     use_cache=True,  # Keep cache enabled
>     past_key_values=None  # Explicitly set past_key_values
> )

Thank you for your reply🤗. I tried this, but it seems to be stuck in a deadlock between using cache and not using cache. The problem may come from the official prompt tuning code of Huggingface. When I used a similar method to load the Lora model, it worked

```
Traceback (most recent call last):
  File "/data/syk/pro/prompt_test.py", line 206, in <module>
    main()
  File "/data/syk/pro/prompt_test.py", line 159, in main
    outputs = model.generate(
  File "/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/peft/peft_model.py", line 1192, in generate
    outputs = self.base_model.generate(**kwargs)
  File "/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/transformers/generation/utils.py", line 2255, in generate
    result = self._sample(
  File "/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/transformers/generation/utils.py", line 3247, in _sample
    model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
  File "/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/peft/peft_model.py", line 1221, in prepare_inputs_for_generation
    if model_kwargs["past_key_values"][0][0].shape[-2] >= model_kwargs["input_ids"].shape[1]:
  File "/data/syk/miniconda3/envs/test/lib/python3.9/site-packages/transformers/cache_utils.py", line 390, in __getitem__
    raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}")
KeyError: 'Cache only has 0 layers, attempted to access layer with index 0'
```


I have also encountered a similar problem, is there anyone else? This comes from Huggingface's prompt learning code repository.

cc @gante @benjaminbossan @sayakpaul for peft/generation interactions

@aohenuo Since the code you provided uses custom checkpoints and does a lot of unrelated stuff, I tried to condense it to a minimal reproducer. However, for me there was no error. Could you please:

1. Try if my code runs successfully for you?
2. If it succeeds, do you have an idea what could be different between the scripts?
3. If it fails, could you please tell us your PEFT and torch versions?

```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel, PromptTuningConfig, get_peft_model

def main():
    MODEL_PATH = "meta-llama/Llama-3.1-8B-Instruct"
    DEVICE = "cuda"
    path = "/tmp/peft/prompt-tuning"

    # load the model
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    tokenizer.pad_token = tokenizer.eos_token

    basemodel = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.bfloat16,
        device_map=DEVICE
    )

    peft_config = PromptTuningConfig(num_virtual_tokens=10, task_type="CAUSAL_LM")
    model = get_peft_model(basemodel, peft_config)
    model.save_pretrained(path)
    del model, basemodel

    basemodel = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.bfloat16,
        device_map=DEVICE
    )
    model = PeftModel.from_pretrained(basemodel, path)

    prompt = "This is a test"
    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
    outputs = model.generate(
        inputs.input_ids,
        max_new_tokens=20,
        temperature=1.0,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id
    )

    # decode the answer
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(generated_text)

if __name__ == "__main__":
    main()
```


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.