BartForConditionalGeneration breaks with label smoothing loss
Author: mingruimingrui
Pull request: no

## Environment info

- `transformers` version: 4.3.3
- The other parameters are irrelevant

### Who can help

@patrickvonplaten @sgugger 

## Information

I apologize for not using the provided template for this issue.

By generating entries with `PreTrainedTokenizer.prepare_seq2seq_batch`, collating with `DataCollatorForSeq2Seq` and
training a `BartForConditionalGeneration` with `Seq2SeqTrainer`, I ran into this particular error message.

```txt
Traceback (most recent call last):
  File "playground.py", line 42, in <module>
    trainer.train()
  File "/Users/mingrui.wang/miniconda3/envs/bert/lib/python3.7/site-packages/transformers/trainer.py", line 940, in train
    tr_loss += self.training_step(model, inputs)
  File "/Users/mingrui.wang/miniconda3/envs/bert/lib/python3.7/site-packages/transformers/trainer.py", line 1304, in training_step
    loss = self.compute_loss(model, inputs)
  File "/Users/mingrui.wang/miniconda3/envs/bert/lib/python3.7/site-packages/transformers/trainer.py", line 1341, in compute_loss
    loss = self.label_smoother(outputs, labels)
  File "/Users/mingrui.wang/miniconda3/envs/bert/lib/python3.7/site-packages/transformers/trainer_pt_utils.py", line 398, in __call__
    nll_loss = log_probs.gather(dim=-1, index=labels)
RuntimeError: Size does not match at dimension 1 expected index [1, 7, 1] to be smaller than src [1, 5, 50265] apart from dimension 2
```

Provided is a script to replicate the error

```python
from torch.utils.data import Dataset
from transformers import (BartForConditionalGeneration, BartTokenizer,
                          BatchEncoding, DataCollatorForSeq2Seq,
                          PreTrainedTokenizer, Seq2SeqTrainer,
                          TrainingArguments)


class DummySeq2SeqDataset(Dataset):

    def __init__(self, tokenizer: PreTrainedTokenizer):
        self.tokenizer = tokenizer
        self.data = [
            ("Hello world!", "Hallo welt!"),
        ]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index: int) -> BatchEncoding:
        src_text, tgt_text = self.data[index]
        return self.tokenizer.prepare_seq2seq_batch(
            src_text, tgt_text,
            return_token_type_ids=False
        )


train_args = TrainingArguments(output_dir='tmp', label_smoothing_factor=0.1)

tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')
model = BartForConditionalGeneration.from_pretrained('facebook/bart-base')

train_dataset = DummySeq2SeqDataset(tokenizer)
data_collator = DataCollatorForSeq2Seq(tokenizer)

trainer = Seq2SeqTrainer(
    model=model,
    args=train_args,
    data_collator=data_collator,
    train_dataset=train_dataset,
    tokenizer=tokenizer,
)
trainer.train()
```

## Source of problem

The problem lies with the interaction between `BartForConditionalGeneration` and how label-smoothing is implemented in `Trainer`.

`BartForConditionalGeneration.forward` is highly tied to `labels`, since it's also used to generate `decoder_input_ids`.

In the behavior of label-smoothing as implemented in the `Trainer` class, the following is currently being done.

https://github.com/huggingface/transformers/blob/3c733f320870261ea948049505a30c30fd6ea23a/src/transformers/trainer.py#L1447-L1469

The whoopsie is that `labels` is removed from the arguments passed to `BartForConditionalGeneration.forward`. Computation of logits then defaulted to using `input_ids` as `decoder_input_ids`.

## Possible solutions

A possible way to fix this would be to shift smooth label loss into the loss computation of each model rather than in `Trainer`. Doing it this way comes with its own set of pros and cons.

Pros
- Backward compatibility can be completely maintained
- Removes the little bit of code smell where true training loss is not reflected in `model.forward` when `label_smoothing > 0`.

Cons
- Complicates configuration
  - label_smoothing loss defined in model config rather than training args
- Requires changes in many places in this repository (albeit, they are the same exact set of changes)

Comments:
Another possible solution would be to change the behavior of `PreTrainedTokenizer.prepare_seq2seq_batch`, by adding something like a `return_decoder_input_ids` flag.

Pros
- Backward compatibility also maintained
- Only a handful of changes has to be made

Cons
- This change can complicate matters for non-autoregressive sequence generation
- A additional cost on serialization and I/O is incurred by having more tensors to pass from data loader workers and send to GPU.

Hi @mingruimingrui 

You are right. But this issue has already been fixed on master. The default `DataCollatorForSeq2Seq` now prepares `decoder_input_ids` when label smoothing is enabled. All seq2seq models now have `prepare_decoder_input_ids_from_labels` method, which let's you prepare `decoder_input_ids` outside of the model. So you can directly pass `decoder_input_ids` to the model and drop `labels` when using label smoothing.

Hi @patil-suraj 

I see but just a little bit of nit-picking regarding this solution...
Mainly because I really can't excuse the passing of the entire model to DataCollator.
Under PyTorch, both Dataset and DataCollator should be passable to python subprocesses.
But a trainable model moved to the GPU would likely not work well with python multiprocessing.
Also worth mentioning is the breaking of backward compatibility.

That's a great point! Pinging @sgugger  here.

> Also worth mentioning is the breaking of backward compatibility.

I'm not sure what's breaking backward compatibility here.

And regarding the previous comments

> Another possible solution would be to change the behavior of PreTrainedTokenizer.prepare_seq2seq_batch, by adding something like a return_decoder_input_ids flag.

This method is now deprecated and will be removed in v5 so we don't encourage using this method.

> A possible way to fix this would be to shift smooth label loss into the loss computation of each model rather than in Trainer

We generally tend to avoid any training-specific code in model files, the model classes are just responsible for doing a forward pass and computing the loss.  Most of the training-related functionality will be handled by `Trainer` or the training scripts.

> I'm not sure what's breaking backward compatibility here.

Ah, this is my bad! I was thinking that since `DataCollatorForSeq2Seq` may not work well with multiprocessing.
Scripts currently using `transformers<=4.3.3` may be affected when updating to `transformers>4.3.3`.
But in `transformers<=4.3.3`, there is no model attribute.

> We generally tend to avoid any training-specific code in model files, the model classes are just responsible for doing a forward pass and computing the loss. Most of the training-related functionality will be handled by `Trainer` or the training scripts.

I agree, adding unrelated functionalities makes a class overloaded and code hard to read.
Though it can be argued that the computation of loss is also a training-related function.

I don't think the multiprocessing will create copies of the models, it will just pass along the reference to it. Did you see GPU memory usage be multiplied by number of processes? I jsut tried your snippet of code and added the model to the data collator:
```
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
```
and I didn't see any change in GPU memory use with 1 or 4 workers in the dataloader.

@sgugger 

The reason I discourage is practice is because it seemed to encourage the usage of model-related components in the `DataCollator`. Passing references of model parameters to python subprocesses is completely fine but special care has to be taken to not use them (in the subprocess) or the terribly descriptive `RuntimeError: CUDA error: initialization error` can be encountered.

GPU memory usage was not my concern in this issue. Similar to what you had mentioned, pickling of PyTorch tensors passes only a reference to the original tensor so GPU memory usage would not increase due to this behavior.

We only use the method of the model ot generate decoder input IDs, not the actual model, so I think it's completely fine in this case. Passing the method from the model would be way weirder in terms of user API.

I see, it's a fair point. Implementing the feature this way also ensures that DataCollator performs all required preprocessing for training data input. Closing issue.