Bug version 4.42.4: KeyError: 'Cache only has 0 layers, attempted to access layer with index 0'
Author: Abhishek-TAMU
Pull request: no
Labels: name: Generation

name: Cache

name: PEFT

Hi @gante @amyeroberts 👋

I am using transformers version 4.42.4 and wrote a code which does prompt tuning using SFTTrainer, and furthermore uses the tuned model to run inference. I can successfully train the model, but when I do inferencing of the model, I am also getting the same error in the function `model.generate`. **I am not getting same issue with transformers version 4.40 and 4.41**. Any input would be helpful. 

Details:
Python version: 3.12.4
transformers version: 4.42.4
Torch Version: 2.5.0.dev20240715
Accelerate version: 0.32.1
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?: No (As the model and dataset is very small)
Using distributed or parallel set-up in script?: No
Using Flash-attention: No

Attaching the code and full error traceback.

```
import datasets, dataclasses
import os, json
from typing import Dict
import math, torch

from transformers import TrainingArguments
import transformers
from peft import PeftModel

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    GPT2Tokenizer,
    GPTNeoXTokenizerFast,
    LlamaTokenizer,
    LlamaTokenizerFast,
)

from trl import SFTConfig, SFTTrainer, DataCollatorForCompletionOnlyLM
from peft import PromptTuningConfig

model_name = "Maykeye/TinyLLama-v0"
tempdir = "tmp"
use_flash_attn = False

model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

""" Function to Pad tokenizer """ 
tokenizer_padding()

#################### DATA PREPROCESSING ###################

response_template = "\n### Label:"
response_template_ids = tokenizer.encode(
    response_template, add_special_tokens=False
)[2:]

data_collator = DataCollatorForCompletionOnlyLM(
    response_template_ids,
    tokenizer=tokenizer,
    ignore_index= -100,
)

format_dataset = lambda example: {  
    "output": example["output"] + tokenizer.eos_token
}

data_files = {"train": "/data/twitter_complaints_small.json"}

json_dataset = datasets.load_dataset("json", data_files=data_files)

formatted_train_dataset = json_dataset["train"].map(format_dataset)

##########################################################

################# PEFT CONFIG ############################

peft_config = PromptTuningConfig(
    task_type="CAUSAL_LM",
    tokenizer_name_or_path="Maykeye/TinyLLama-v0",
    num_virtual_tokens=8,
    prompt_tuning_init="RANDOM",
    prompt_tuning_init_text="hello"
)

##########################################################

################# Trainer CONFIG ############################

train_args = TrainingArguments(
    num_train_epochs=5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=0.00001,
    weight_decay=0,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    logging_steps=1,
    include_tokens_per_second=True,
    save_strategy="epoch",
    output_dir=tempdir,
    use_cpu = True
)
transformer_train_arg_fields = [x.name for x in dataclasses.fields(SFTConfig)]
transformer_kwargs = {
    k: v
    for k, v in train_args.to_dict().items()
    if k in transformer_train_arg_fields
}
training_args = SFTConfig(**transformer_kwargs)

##########################################################

####################### TUNING #####################

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=formatted_train_dataset,
    packing=False,
    data_collator=data_collator,
    dataset_text_field="output",
    args=training_args,
    max_seq_length=2048,
    peft_config=peft_config,
)
trainer.train()

print("### MODEL TUNING DONE !!")

####################### INFERENCE ##################

### LOAD peft tuning configs
checkpoint_path = os.path.join(tempdir, "checkpoint-5")
with open(os.path.join(checkpoint_path, "adapter_config.json"), encoding="utf-8") as f:
    adapter_config = json.load(f)

################ Load the model ###################

tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    attn_implementation="flash_attention_2"
    if use_flash_attn
    else None,
    torch_dtype=torch.bfloat16 if use_flash_attn else None,
)

### FUNCTION to Resize tokenizer and embedding
tokenizer_and_embedding_resize(
    {}, tokenizer=tokenizer, model=base_model
)
model = PeftModel.from_pretrained(
    base_model,
    checkpoint_path,
    attn_implementation="flash_attention_2"
    if use_flash_attn
    else None,
    torch_dtype=torch.bfloat16 if use_flash_attn else None,
)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)


######## Run inference on the text #########

text = "### Text: @NortonSupport Thanks much.\n\n### Label:"
tok_res = tokenizer(text, return_tensors="pt")
input_ids = tok_res.input_ids.to(device)

peft_outputs = model.generate(
    input_ids=input_ids, max_new_tokens=50
)

decoded_result = tokenizer.batch_decode(
    peft_outputs, skip_special_tokens=False
)[0]
decoded_result = decoded_result

print("### MODEL INFERENCE DONE !!")
```

This throws the error:

```
Traceback (most recent call last):
  File "/Users/abhishekm/Documents/Codes/fms-hf-tuning/prompt_tune_test.py", line 214, in <module>
    peft_outputs = model.generate(
                   ^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/transf_test/lib/python3.12/site-packages/peft/peft_model.py", line 1493, in generate
    outputs = self.base_model.generate(**kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/transf_test/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/transf_test/lib/python3.12/site-packages/transformers/generation/utils.py", line 1914, in generate
    result = self._sample(
             ^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/transf_test/lib/python3.12/site-packages/transformers/generation/utils.py", line 2648, in _sample
    model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/transf_test/lib/python3.12/site-packages/peft/peft_model.py", line 1522, in prepare_inputs_for_generation
    if model_kwargs["past_key_values"][0][0].shape[-2] >= model_kwargs["input_ids"].shape[1]:
       ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^
  File "/opt/anaconda3/envs/transf_test/lib/python3.12/site-packages/transformers/cache_utils.py", line 314, 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'
```

_Originally posted by @Abhishek-TAMU in https://github.com/huggingface/transformers/issues/27985#issuecomment-2231281845_
            

Comments:
Hi @amyeroberts  @gante @ArthurZucker @zucchini-nlp 👋
Running above code (with same above mentioned specifications and transformers version used : 4.42.4) works for LoRA tuning and Full-fine tuning. I just get the error in Prompt-tuning as mentioned above. 
Any inputs from my side is needed for repro of this issue ?

@Abhishek-TAMU hey! Can you share the inference code only with the hub checkpoints you are using. Cache objects currently are not uniform, so sharing a small reproducible code would be helpful

Hi @zucchini-nlp, Thanks for attending this issue. Sharing the inference code which pulls the prompt-tuned model from hub and this I think this would help in reproducing the issue. This code works with transformers version `4.40.x` but doesn't work and gives error with version `4.42.x` and `4.43.x`

```
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
)
import torch
import math
import transformers
from typing import Dict
from peft import PeftModel

use_flash_attn=False
model_name="Maykeye/TinyLLama-v0"

########################## LOAD peft tuning configs ##################

checkpoint_path = "abhishek-maurya/tinnyllama-prompt-tuning"

################ Load the model ###################

tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    attn_implementation="flash_attention_2"
    if use_flash_attn
    else None,
    torch_dtype=torch.bfloat16 if use_flash_attn else None,
)

## FUNCTION to Resize tokenizer and embedding
tokenizer_and_embedding_resize(
    {}, tokenizer=tokenizer, model=base_model
)
model = PeftModel.from_pretrained(
    base_model,
    checkpoint_path,
    attn_implementation="flash_attention_2"
    if use_flash_attn
    else None,
    torch_dtype=torch.bfloat16 if use_flash_attn else None,
)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)


######## Run inference on the text #########

text = "### Text: @NortonSupport Thanks much.\n\n### Label:"
tok_res = tokenizer(text, return_tensors="pt")
input_ids = tok_res.input_ids.to(device)

peft_outputs = model.generate(
    input_ids=input_ids, max_new_tokens=50
)

decoded_result = tokenizer.batch_decode(
    peft_outputs, skip_special_tokens=False
)[0]
decoded_result = decoded_result

print("### MODEL INFERENCE DONE !!")
print("Decoded_result:", decoded_result)


```

Hi @Abhishek-TAMU 👋 

In the snippet you shared, the definition of `tokenizer_and_embedding_resize` is missing -- could you share it with us? 🤗  

Thanks @gante for diving in this one. 
I skipped the definition of `tokenizer_and_embedding_resize` thinking it might not be relevant in reproducing the issue, making the minimal reproducible script long.

Sharing the definition of function `tokenizer_and_embedding_resize` [here](https://github.com/foundation-model-stack/fms-hf-tuning/blob/55ca61219264a26ce68dbbcdcc463ae6af3ec3fa/tuning/data/tokenizer_data_utils.py#L23)

CC: @zucchini-nlp 

Hi @Abhishek-TAMU 👋 

The script was crashing on my end too at first, but then I updated peft to 0.12. With peft 0.12, it runs :)

Updating peft to 0.12 worked!!! Thank you @gante and @zucchini-nlp for the fix. 

try close use_unsloth