Error in `generate` when wrapping a pre-trained language model in nn.Module class and using custom embeddings: Key and Value must have the same sequence length
Author: jonflynng
Pull request: no
Labels: name: bug

name: Generation

### System Info

Google Colab on L4 GPU:

Transformers version: 4.45.0.dev0
PyTorch version: 2.3.1+cu121
Python version: 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0]

### Who can help?

@ArthurZucker @zucchini-nlp 

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

I'm encountering an issue when attempting to use the `generate` method with `inputs_embeds` instead of `input_ids` in a custom model based on `AutoModelForCausalLM`. 

**Note: this error only occurs when wrapping the language model in a wrapper class with `nn.Module`**

Steps to reproduce:
1. Create a custom model inheriting from `nn.Module` that uses `AutoModelForCausalLM` as its base.
2. Train the model using `Trainer` from transformers or a normal training loop, both cause the issue.
3. Attempt to generate output using `model.language_model.generate()` with `inputs_embeds` instead of `input_ids`.

Additional context:
- This issue occurs only after training the model.
- My custom model requires the use of `inputs_embeds` for specific use cases that can't be handled with `input_ids` alone.

Code:

```
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoProcessor, AutoTokenizer, AutoModelForCausalLM, AutoConfig, MusicgenConfig, MusicgenForConditionalGeneration
from peft import get_peft_model, LoraConfig, TaskType

class CustomModel(nn.Module):
    def __init__(self, language_model_name, model_type, device):
        super(CustomModel, self).__init__()

        self.device = device

        self.tokenizer = AutoTokenizer.from_pretrained(language_model_name)
        self.language_model = AutoModelForCausalLM.from_pretrained(language_model_name, torch_dtype=torch.bfloat16, device_map="auto").to(self.device)

        peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM, inference_mode=False, r=8, lora_alpha=16, lora_dropout=0, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"])

        self.language_model = get_peft_model(
                self.language_model,
                peft_config
            )

        self.language_model.gradient_checkpointing_enable()

model = CustomModel("meta-llama/Meta-Llama-3-8B-Instruct", "cuda")

# Dummy training data
dummy_text = f"This is a dummy text input."
dummy_multimodal_embeddings = torch.randn(1, 300, model.language_model.config.hidden_size, dtype=torch.bfloat16).to("cuda")

optimizer = torch.optim.AdamW(model.language_model.parameters(), lr=1e-5)

def train_step():
    model.language_model.train()
    optimizer.zero_grad()

    # Tokenize text and get embeddings
    text_inputs = model.tokenizer(dummy_text, return_tensors="pt", add_special_tokens=False).to("cuda")
    text_embeds = model.language_model.get_input_embeddings()(text_inputs.input_ids).to("cuda", dtype=torch.bfloat16)

    text_embeds.requires_grad = True

    # Combine embeddings
    combined_embeds = torch.cat([text_embeds, dummy_multimodal_embeddings], dim=1)

    # Create attention mask and position ids
    attention_mask = torch.ones((1, combined_embeds.shape[1]), dtype=torch.long).to("cuda")
    position_ids = torch.arange(combined_embeds.shape[1], dtype=torch.long).unsqueeze(0).to("cuda")

    # Create labels
    labels = torch.full((1, combined_embeds.shape[1]), -100, dtype=torch.long).to("cuda")
    text_length = text_inputs.input_ids.size(1)
    labels[:, :text_length-1] = text_inputs.input_ids[:, 1:]

    outputs = model.language_model(
        inputs_embeds=combined_embeds,
        attention_mask=attention_mask,
        position_ids=position_ids,
        labels=labels
    )
    loss = outputs.loss
    loss.backward()
    optimizer.step()

for _ in range(20):
    train_step()
    
    from transformers import TextStreamer

test_input = "Continue the Fibonacci sequence: 1, 1, 2, 3, 5, 8,"
inputs = model.tokenizer(test_input, return_tensors="pt").to("cuda")
text_embeds = model.language_model.get_input_embeddings()(inputs.input_ids).to("cuda", dtype=torch.bfloat16)

attention_mask = torch.ones((1, text_embeds.shape[1]), dtype=torch.long).to("cuda")

text_streamer = TextStreamer(model.tokenizer)

model.language_model.generate(
        inputs_embeds=text_embeds,
        attention_mask=attention_mask,
        streamer=text_streamer,
        max_new_tokens=256,
    )
```

Stack trace and error:

```
RuntimeError                              Traceback (most recent call last)
<ipython-input-19-52aa2192be3a> in <cell line: 9>()
      7 inputs_embeds = model.language_model.get_input_embeddings()(inputs.input_ids)
      8 text_streamer = TextStreamer(model.tokenizer)
----> 9 _ = model.language_model.generate(inputs_embeds=inputs_embeds, streamer = text_streamer, max_new_tokens = 256, use_cache = True)

21 frames
/usr/local/lib/python3.10/dist-packages/peft/peft_model.py in generate(self, *args, **kwargs)
   1644                 with self._enable_peft_forward_hooks(*args, **kwargs):
   1645                     kwargs = {k: v for k, v in kwargs.items() if k not in self.special_peft_forward_args}
-> 1646                     outputs = self.base_model.generate(*args, **kwargs)
   1647             else:
   1648                 outputs = self.base_model.generate(**kwargs)

/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py in decorate_context(*args, **kwargs)
    113     def decorate_context(*args, **kwargs):
    114         with ctx_factory():
--> 115             return func(*args, **kwargs)
    116 
    117     return decorate_context

/usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py in generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, **kwargs)
   2021 
   2022             # 13. run sample (it degenerates to greedy search when `generation_config.do_sample=False`)
-> 2023             result = self._sample(
   2024                 input_ids,
   2025                 logits_processor=prepared_logits_processor,

/usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py in _sample(self, input_ids, logits_processor, stopping_criteria, generation_config, synced_gpus, streamer, logits_warper, **model_kwargs)
   2979 
   2980             # forward pass to get next token
-> 2981             outputs = self(**model_inputs, return_dict=True)
   2982 
   2983             if synced_gpus and this_peer_finished:

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)
   1530             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531         else:
-> 1532             return self._call_impl(*args, **kwargs)
   1533 
   1534     def _call_impl(self, *args, **kwargs):

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1539                 or _global_backward_pre_hooks or _global_backward_hooks
   1540                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541             return forward_call(*args, **kwargs)
   1542 
   1543         try:

/usr/local/lib/python3.10/dist-packages/transformers/models/llama/modeling_llama.py in forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
   1187 
   1188         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-> 1189         outputs = self.model(
   1190             input_ids=input_ids,
   1191             attention_mask=attention_mask,

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)
   1530             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531         else:
-> 1532             return self._call_impl(*args, **kwargs)
   1533 
   1534     def _call_impl(self, *args, **kwargs):

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1539                 or _global_backward_pre_hooks or _global_backward_hooks
   1540                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541             return forward_call(*args, **kwargs)
   1542 
   1543         try:

/usr/local/lib/python3.10/dist-packages/transformers/models/llama/modeling_llama.py in forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
    987 
    988             if self.gradient_checkpointing and self.training:
--> 989                 layer_outputs = self._gradient_checkpointing_func(
    990                     decoder_layer.__call__,
    991                     hidden_states,

/usr/local/lib/python3.10/dist-packages/torch/_compile.py in inner(*args, **kwargs)
     22             import torch._dynamo
     23 
---> 24             return torch._dynamo.disable(fn, recursive)(*args, **kwargs)
     25 
     26         return inner

/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py in _fn(*args, **kwargs)
    449             prior = set_eval_frame(callback)
    450             try:
--> 451                 return fn(*args, **kwargs)
    452             finally:
    453                 set_eval_frame(prior)

/usr/local/lib/python3.10/dist-packages/torch/_dynamo/external_utils.py in inner(*args, **kwargs)
     34     @functools.wraps(fn)
     35     def inner(*args, **kwargs):
---> 36         return fn(*args, **kwargs)
     37 
     38     return inner

/usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py in checkpoint(function, use_reentrant, context_fn, determinism_check, debug, *args, **kwargs)
    485                 "use_reentrant=False."
    486             )
--> 487         return CheckpointFunction.apply(function, preserve, *args)
    488     else:
    489         gen = _checkpoint_without_reentrant_generator(

/usr/local/lib/python3.10/dist-packages/torch/autograd/function.py in apply(cls, *args, **kwargs)
    596             # See NOTE: [functorch vjp and autograd interaction]
    597             args = _functorch.utils.unwrap_dead_wrappers(args)
--> 598             return super().apply(*args, **kwargs)  # type: ignore[misc]
    599 
    600         if not is_setup_ctx_defined:

/usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py in forward(ctx, run_function, preserve_rng_state, *args)
    260 
    261         with torch.no_grad():
--> 262             outputs = run_function(*args)
    263         return outputs
    264 

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)
   1530             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531         else:
-> 1532             return self._call_impl(*args, **kwargs)
   1533 
   1534     def _call_impl(self, *args, **kwargs):

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1539                 or _global_backward_pre_hooks or _global_backward_hooks
   1540                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541             return forward_call(*args, **kwargs)
   1542 
   1543         try:

/usr/local/lib/python3.10/dist-packages/transformers/models/llama/modeling_llama.py in forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cache_position, position_embeddings, **kwargs)
    732 
    733         # Self Attention
--> 734         hidden_states, self_attn_weights, present_key_value = self.self_attn(
    735             hidden_states=hidden_states,
    736             attention_mask=attention_mask,

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)
   1530             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531         else:
-> 1532             return self._call_impl(*args, **kwargs)
   1533 
   1534     def _call_impl(self, *args, **kwargs):

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1539                 or _global_backward_pre_hooks or _global_backward_hooks
   1540                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541             return forward_call(*args, **kwargs)
   1542 
   1543         try:

/usr/local/lib/python3.10/dist-packages/transformers/models/llama/modeling_llama.py in forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cache_position, position_embeddings, **kwargs)
    658         is_causal = True if causal_mask is None and q_len > 1 else False
    659 
--> 660         attn_output = torch.nn.functional.scaled_dot_product_attention(
    661             query_states,
    662             key_states,

RuntimeError: Key and Value must have the same sequence length
```

### Expected behavior

The generation process fails with the aforementioned RuntimeError, indicating a mismatch in sequence lengths for key and value tensors in the attention mechanism.

Comments:
Hey @jonflynng !

You have ti disable `gradient_checkpointing` after training is done with `model.language_model.gradient_checkpointing_disable()`. Caching mechanism does jot work with gradient checkpointing, which was causing the error

Thanks, that fixed it.

Also, to note, this still seems to happen if `position_ids` are passed in as an argument to `generate`.

Yes, `generate()` currently can give unexpected results if a user passes `position_ids`. We will work on that