Training stage error with batch mode on conditional generation for multimodal models
### System Info

- `transformers` version: 4.29.2
- Platform: Linux-5.15.90.1-microsoft-standard-WSL2-x86_64-with-glibc2.17
- Python version: 3.8.16
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cpu (False)
- 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
- Using distributed or parallel set-up in script?: no

### Who can help?

@younesbelkada @gante @sgugger 

### Information

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

### Tasks

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

### Reproduction

```
import torch
from transformers import VisionEncoderDecoderModel, TrOCRProcessor

loc = "microsoft/trocr-small-handwritten"

processor = TrOCRProcessor.from_pretrained(loc)
model = VisionEncoderDecoderModel.from_pretrained(loc)

decoder_input_ids = torch.tensor([[0,  7344,  2159,    12, 345], [0,  7344,  2159,    12, 346]]) # a batch_size of 2 Prefixes for each examples 
decoder_attention_mask = None
encoder_hidden_states = torch.randn(2, 578, 384) # a random encoder input to the decoder
encoder_attention_mask = None
decoder_inputs_embeds = None
output_attentions = None
output_hidden_states = None
use_cache = None
past_key_values = None
return_dict = True
kwargs_decoder = {}

decoder_outputs = model.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            inputs_embeds=decoder_inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            use_cache=use_cache,
            past_key_values=past_key_values,
            return_dict=return_dict,
            **kwargs_decoder,
        )
decoder_outputs['logits'].shape, decoder_input_ids.shape
```

### Expected behavior

The decoder default max_target_length is set to 128. So I expect the logits with shape [2, 128, 64044], where batch_size is 2. But I only get shape of [2, 5, 64044], where prefix length is 5.
This is an issue for[ this merged PR](https://github.com/huggingface/transformers/pull/22424)<|||||>Hey @cramraj8 👋 

The output shape is correct -- the sequence length of the logits is as long as the input sequence length. 

Perhaps you're interested in using the model to complete the prefix, in which case you'll need to use generative tools like `.generate()` or the `Seq2SeqTrainer`. Have a look at the documentation for these terms :)<|||||>Hi @gante , Thanks for the reply. Yes, I tried using `.generate() `and `Seq2SeqTrainer` to do the training. But if I provide prefix `input_ids` and parse them as `decoder_input_ids ` to the `.generate()` function, the Trainer throws error during the training stage claiming the loss calculation end up getting different shape for prediction and label. That is why I looked up `model.decode() `function. It is true that output sequence length must be same as input sequence length for the `decode() `function. But for the prefix completion task, how do we adapt it for `generate`() and `decode`() function.

I do have this question - I am interested in conditional decoding where the a multimodal completes a prefix or a prompt. Do we train the model with complete text (prefix and completion text) during the training stage, and only provide prefix as an additional input during inference stage ?

Or we can still provide prefix as additional input to both training and inference stages ?<|||||>@cramraj8 I see, now I understand what your goal is :) 

AFAIK We do not support passing a prefix at train time, I'm afraid you'll have to build a custom solution. In any case, it must be based on `.generate()` and `Seq2SeqTrainer`, as your task relies on auto-regressive text generation!

You can also train without a prefix at all, even if you expect a certain prefix (or set of prefixes) at inference time. For instance, Whisper does this (see [section 2.3 of its paper](https://arxiv.org/pdf/2212.04356.pdf)). At train time, treat the prefixes as variables. At inference time, starts generating from the prefix.<|||||>Thank you! This is helpful. Looks like not doing anything during train time, and applying prefix during test time works better. <|||||>Hi @gante, during my implementation I found that I am getting an error of device mismatch at the following location.
```
File "/mnt/azureml/cr/.../exe/wd/trainer_seq2seq.py", line 296, in prediction_step
    generated_tokens = self.model.generate(**inputs, **gen_kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/generation/utils.py", line 1328, in generate
    input_ids, model_kwargs = self._prepare_decoder_input_ids_for_generation(
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/generation/utils.py", line 676, in _prepare_decoder_input_ids_for_generation
    decoder_input_ids = torch.cat([decoder_input_ids_start, decoder_input_ids], dim=-1)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:4 and cpu! (when checking argument for argument tensors in method wrapper_cat)
```

This happens when I do `Seq2SeqTrainer `evaluation with prefix (`decoder_input_ids`) provided. In CPU machine, the code works perfectly fine. But at the GPU presence, it was throwing this device error. After debugging I found that the `decoder_input_ids` were never placed in the corresponding cuda device IDs even though other input values were placed in GPU. I did the following change, and it worked fine now. I am not sure if it's a bug or not, but I am bringing this up to your attention if anyone face similar issues in future. 

In addition, I had to overwrite Seq2SeqTrainer to separate `decoder_input_ids` from `inputs `and assign it with `self._gen_kwargs` so that the code works. Otherwise, I was getting complex errors.

```
decoder_input_ids = inputs.pop("decoder_input_ids")
self._gen_kwargs["decoder_input_ids"] = decoder_input_ids
```

(adding @wgx998877 for reference)
<|||||>@cramraj8 would you be able to share a short reproducer? (like the one you shared at the top)