fsmt translation with use_cache=True bug
Author: lihaoxin2020
Pull request: no

### System Info

I'm running [run_translation.py](https://github.com/huggingface/transformers/tree/main/examples/pytorch/translation) to train from scratch with fsmt artifacts. 

I'm using 
- `transformers` version: 4.27.0.dev0
- Platform: Linux-3.10.0-862.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.8.13
- Huggingface_hub version: 0.12.0
- PyTorch version (GPU?): 1.13.1+cu117 (True)

### Who can help?

@ArthurZucker @younesbelkada 

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

1. Detailed configuration `run_translation_config.json` as follow:
```json
{
    "config_name": "allenai/wmt19-de-en-6-6-base",
    "tokenizer_name": "allenai/wmt16-en-de-12-1",
    "use_fast_tokenizer": false,
    "dataset_name": "wmt16",
    "dataset_config_name": "de-en",
    "source_lang": "en",
    "target_lang": "de",
    "max_source_length": 512,
    "max_target_length": 512,
    "generation_max_length": 512,
    "preprocessing_num_workers": 32,
    "output_dir": "./tmp/transformer_base_wmt16_6_6_en-de_fsmt",
    "run_name": "transformer_base_wmt16_6_6_en-de_fsmt",
    "deepspeed": "./ds_stage2_config.json", (NOT important)
    "sortish_sampler": true,
    "overwrite_output_dir": true,
    "do_train": true,
    "do_eval": true,
    "do_predict": false,
    "evaluation_strategy": "steps",
    "eval_steps": 2000,
    "predict_with_generate": true,
    "load_best_model_at_end": true,
    "logging_strategy": "steps",
    "logging_steps": 1000,
    "logging_first_step" :true,
    "eval_accumulation_steps": 16,
    "dataloader_num_workers": 32,
    "per_device_train_batch_size": 64,
    "per_device_eval_batch_size": 64,
    "gradient_accumulation_steps": 4,
    "fp16": true,
    "adam_beta1": 0.9,
    "adam_beta2":0.998,
    "adam_epsilon": 1e-08,
    "learning_rate": 5e-4,
    "weight_decay": 0.01,
    "label_smoothing_factor": 0.1,
    "lr_scheduler_type": "linear",
    "warmup_ratio": 0.04,
    "num_train_epochs": 25,
    "save_strategy": "steps",
    "save_steps": 2000,
    "save_total_limit": 10,
    "seed": 42
}
```

2. run command 
```shell
python -m torch.distributed.launch --nproc_per_node 8 \
    run_translation.py \
    run_translation_config.json
```

### Expected behavior

```
Traceback (most recent call last): 
  File "run_translation.py", line 675, in <module>
    main()
  File "run_translation.py", line 592, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/home/ace14487oj/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/transformers/trainer.py", line 1570, in train
    return inner_training_loop(
  File "/home/ace14487oj/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/transformers/trainer.py", line 1835, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/home/ace14487oj/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/transformers/trainer.py", line 2583, in training_step
    loss = self.compute_loss(model, inputs)
  File "/home/ace14487oj/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/transformers/trainer.py", line 2625, in compute_loss
    loss = self.label_smoother(outputs, labels)
  File "/home/ace14487oj/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/transformers/trainer_pt_utils.py", line 498, in __call__
    nll_loss = log_probs.gather(dim=-1, index=labels)
RuntimeError: Size does not match at dimension 1 expected index [64, 104, 1] to be smaller than self [64, 1, 43536] apart from dimension 2
```

After some investigation, I found the error was about use_cache implementation and only happens if I leave use_cache = true in fsmt_model_config. When I tried use_cache = false, the error disappear. I also tried to train BART with run_translation.py on exactly same parameters (use_cache = true), and similar issue does not happen on BART. 

Has anybody run into this issue as well or otherwise maybe confirm this is a bug? 

Thanks! 

Comments:
Hi @lihaoxin2020 
Thanks for the issue! 
I tried to run the script with the command:
```bash
python -m torch.distributed.launch --nproc_per_node 2 run_translation.py run_translation_config.json 
```
and getting:
```
run_translation.py: error: the following arguments are required: --model_name_or_path, --output_dir
```
If I put the correct values for those flags I get `ValueError: Need either a dataset name or a training/validation file.`

I will try to see if I can reproduce locally without having to run this script

Based on your traceback, here is a script I made to try to reproduce the issue, 

```python
import torch
from transformers import AutoModelForSeq2SeqLM
from transformers.trainer_pt_utils import LabelSmoother

model = AutoModelForSeq2SeqLM.from_pretrained("allenai/wmt19-de-en-6-6-base")
dummy_input = torch.LongTensor([[0, 1, 1, 2, 3]])
labels = torch.LongTensor([[0, 1, 1, 2, 3]])

outputs = model(input_ids=dummy_input, use_cache=True)

label_smoother = LabelSmoother()
loss = label_smoother(outputs, labels, shift_labels=True)
```
Does this script works for you? 

> Based on your traceback, here is a script I made to try to reproduce the issue,
> 
> ```python
> import torch
> from transformers import AutoModelForSeq2SeqLM
> from transformers.trainer_pt_utils import LabelSmoother
> 
> model = AutoModelForSeq2SeqLM.from_pretrained("allenai/wmt19-de-en-6-6-base")
> dummy_input = torch.LongTensor([[0, 1, 1, 2, 3]])
> labels = torch.LongTensor([[0, 1, 1, 2, 3]])
> 
> outputs = model(input_ids=dummy_input, use_cache=True)
> 
> label_smoother = LabelSmoother()
> loss = label_smoother(outputs, labels, shift_labels=True)
> ```
> 
> Does this script works for you?

Hi @younesbelkada ! For this script, I think you need to try `outputs = model(input_ids=dummy_input, decoder_input_ids=labels, use_cache=True)` to make it equivalent to the context I referred to. 

I got the same error with the new `outputs` line:
```
Traceback (most recent call last):
  File "./playground.py", line 16, in <module>
    loss = label_smoother(outputs, labels, shift_labels=True)
  File "/mmfs1/home/lihaoxin/workspace/lihaoxin/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/transformers/trainer_pt_utils.py", line 498, in __call__
    nll_loss = log_probs.gather(dim=-1, index=labels)
RuntimeError: Size does not match at dimension 1 expected index [1, 4, 1] to be smaller than self [1, 0, 43536] apart from dimension 2
```


Hi @lihaoxin2020 
I managed to reproduce the issue, thanks a lot! 
As a quick fix I propose you to not use `use_cache` for now while we investigate what is happening ! Thanks!

I also had the same issue (albeit with a custom training script), here's what I think is happening:

In the run_translation config you've set `label_smoothing_factor` to greater than 0. As a result, the `labels` field is removed from the model forward call on the [Trainer file](https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py) when computing loss:
```
if self.label_smoother is not None and "labels" in inputs:
          labels = inputs.pop("labels")
```

This is done because if labels is provided to the FSMT model when called, the loss will be calculated twice. (Once in the FSMT model call, and once in the label smoothing class).

In [modeling_fsmt](https://github.com/huggingface/transformers/blob/main/src/transformers/models/fsmt/modeling_fsmt.py), we can see that use_cache is only implicitly disabled when labels is provided, which is not the case with label smoothing enabled:
```
if labels is not None:
    use_cache = False
```

In FSMTDecoder, we can see that when use_cache is enabled, even during training, it will slice off all of the input IDs except for the last whenever called:
```
if use_cache:
    input_ids = input_ids[:, -1:]
    positions = positions[:, -1:]  # happens after we embed them
```

I don't know how BART handles it, but in [Marian](https://github.com/huggingface/transformers/blob/main/src/transformers/models/marian/modeling_marian.py) (which copies a lot of BART code), they only slice off the input_ids when preparing inputs for generation (past_key_values is probably only passed when use_cache is on):
```
def prepare_inputs_for_generation(
    self, input_ids, past_key_values=None, attention_mask=None, use_cache=None, **kwargs
):
    # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
    if attention_mask is None:
        attention_mask = input_ids.new_ones(input_ids.shape)

    if past_key_values:
        input_ids = input_ids[:, -1:]
```

[younesbelkada](https://github.com/younesbelkada)'s solution definitely works for now.

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.