support for model.generate with assistant_model / model being load_in_8bit and PeftModel (LoRA)
Author: achibb
Pull request: no

### System Info

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

- `transformers` version: 4.29.2
- Platform: Windows-10-10.0.22621-SP0
- Python version: 3.9.16
- Huggingface_hub version: 0.14.1
- Safetensors version: not installed
- PyTorch version (GPU?): 1.13.1+cu116 (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, RTX3090ti
- Using distributed or parallel set-up in script?: idk, I use accelerate via: device_map="auto", and load_in_8bit=True

### Who can help?

@gante 

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

Reproduce:

1: Setup load PEFT model as assistant_model (bloom with lora, load_in_8bit):
```python

import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

peft_model_id = "aari1995/GermanGPT_dolly_lora_1b5"
config = PeftConfig.from_pretrained(peft_model_id)
assistant_model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path).to("cuda:0")#, return_dict=True, load_in_8bit=True, device_map='auto')
# Load the Lora model
assistant_model = PeftModel.from_pretrained(assistant_model, peft_model_id)
```

2: Load Bloom model (load_in_8bit, with or without lora):
```python
model = AutoModelForCausalLM.from_pretrained("malteos/bloom-6b4-clp-german-oasst-v0.1",load_in_8bit=True, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained("malteos/bloom-6b4-clp-german-oasst-v0.1")
```

3. Generate using PeftModel as :
```python
prompt = "
Hallo

"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")
outputs = assistant_model.generate(**inputs, assistant_model=model,use_cache=True)
print(tokenizer.batch_decode(outputs, skip_special_tokens=False))
```

### Expected behavior

I expected to get an assisted generation, however it gets stuck, for example here: 
generation/utils.py
-> 4253     prev_seq_len = model_kwargs["assistant_past_key_values"][0][assistant_kv_indexing].shape[-2]
'NoneType' object is not subscriptable

i suspect it maybe due to the fact that i use load_in_8bit=True and also use PeftModel.

Thanks!

Comments:
cc @younesbelkada since @gante  is on vacation.

Same thing here, did not realize this was because of 8-bit.

Following the instructions in the blog post for assisted generation, I run into some issues. (FYI, both the longform_model and assistant_model are 8-bit finetuned versions of OPT, which is the exact same model used in the blog post.)

First, when I do exactly what's in the post:
```
  prompt = prompt + "\nAnswer:"
  inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
  outputs = longform_model.generate(**inputs, assistant_model=assistant_model)
  print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
```

I get an error telling me that assisted generation requires `use_cache=True`. Hmm... weird, and the blog post didn't seem to need to use that argument, but okay, let's try it!

```
  prompt = prompt + "\nAnswer:"
  inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
  outputs = longform_model.generate(**inputs, assistant_model=assistant_model, use_cache=True)
  print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
  ```

Then this happens:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-10-e9645bbc79d4> in <module>
----> 1 generate_from_prompt("Which is a species of fish? Tope or rope?")

<ipython-input-9-14fc80d284ea> in generate_from_prompt(prompt)
      2     prompt = prompt + "\nAnswer:"
      3     inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
----> 4     outputs = longform_model.generate(**inputs, assistant_model=assistant_model, use_cache=True)
      5     print(tokenizer.batch_decode(outputs, skip_special_tokens=True))

/usr/lib/python3/dist-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     25         def decorate_context(*args, **kwargs):
     26             with self.clone():
---> 27                 return func(*args, **kwargs)
     28         return cast(F, decorate_context)
     29 

~/.local/lib/python3.8/site-packages/transformers/generation/utils.py in generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, **kwargs)
   1493 
   1494             # 12. run assisted generate
-> 1495             return self.assisted_decoding(
   1496                 input_ids,
   1497                 assistant_model=assistant_model,

~/.local/lib/python3.8/site-packages/transformers/generation/utils.py in assisted_decoding(self, input_ids, assistant_model, do_sample, logits_processor, logits_warper, stopping_criteria, pad_token_id, eos_token_id, output_attentions, output_hidden_states, output_scores, return_dict_in_generate, synced_gpus, streamer, **model_kwargs)
   4253                 # 1.1. use the assistant model to obtain the next candidate logits
   4254                 if "assistant_past_key_values" in model_kwargs:
-> 4255                     prev_seq_len = model_kwargs["assistant_past_key_values"][0][assistant_kv_indexing].shape[-2]
   4256                     # `new_token_len` can be 1 or 2 (next token in assistant + last token picked by the larger model)
   4257                     new_token_len = candidate_input_ids.shape[1] - prev_seq_len

TypeError: 'NoneType' object is not subscriptable
```

I'm using bleeding edge version of Transformers, so I'm curious what I'm doing wrong here, or else maybe this is just a bug.

Hey @achibb @andersonbcdefg 👋  First of all, apologies for the delay :)

I looked at your script, and the root cause for the exception on my end (with transformers v4.30, peft 0.3.0, and torch 2.0.0) was in the execution of the PEFT model -- it had caching set to False. Assisted generation needs caching on both models, so manually setting the config fixed it. This means the `use_cache` argument in `generate()` is not being piped correctly, for which I'll open a PR 🤗 

______________________________________
(temporary fix until the PR gets merged)
```py
import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

peft_model_id = "aari1995/GermanGPT_dolly_lora_1b5"
config = PeftConfig.from_pretrained(peft_model_id)
assistant_model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path).to("cuda:0")#, return_dict=True, load_in_8bit=True, device_map='auto')
# Load the Lora model
assistant_model = PeftModel.from_pretrained(assistant_model, peft_model_id).to("cuda:0")
assistant_model.config.use_cache = True

model = AutoModelForCausalLM.from_pretrained("malteos/bloom-6b4-clp-german-oasst-v0.1",load_in_8bit=True, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained("malteos/bloom-6b4-clp-german-oasst-v0.1")

prompt = "
Hallo

"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")
outputs = assistant_model.generate(**inputs, assistant_model=model, use_cache=True)
print(tokenizer.batch_decode(outputs, skip_special_tokens=False))
```

@gante thanks, however I see a mistake on my side, I accidently switched the models. My generation model is a regular model and my assistant is a peft, so:

```python
import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

peft_model_id = "aari1995/GermanGPT_dolly_lora_1b5"
config = PeftConfig.from_pretrained(peft_model_id)
assistant_model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path).to("cuda:0")#, return_dict=True, load_in_8bit=True, device_map='auto')
# Load the Lora model
assistant_model = PeftModel.from_pretrained(assistant_model, peft_model_id).to("cuda:0")
assistant_model.config.use_cache = True

model = AutoModelForCausalLM.from_pretrained("malteos/bloom-6b4-clp-german-oasst-v0.1",load_in_8bit=True, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained("malteos/bloom-6b4-clp-german-oasst-v0.1")

prompt = "
Hallo

"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")
outputs = model.generate(**inputs, assistant_model=assistant_model, use_cache=True)
print(tokenizer.batch_decode(outputs, skip_special_tokens=False))
```
Here i get the error:

RuntimeError: The expanded size of the tensor (11) must match the existing size (12) at non-singleton dimension 2. 
Target sizes: [16, 0, 11].  Tensor sizes: [16, 1, 12]



I hope your vacation was nice!

(@achibb  looking at your newest comment to determine whether the issue needs to be reopened :) )

@achibb the root cause was the different class name, when the model gets loaded with PEFT. See the PR description in #24198 to see how it impacted the script you were trying to run :p 

After the PR gets merged, you will be able to run the script you pasted above!

It should be sorted now -- try running your script from `main` :)

perfect, works now! Thanks :)