Repository: huggingface/transformers

Issue:
run_summarization script breaks with label_smoothing_factor and pad_to_max_length true
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: '4.5.0.dev0' (from source)
- Platform: Linux
- Python version: 3.6.9
- PyTorch version (GPU?): '1.8.0' (yes)

## Information

I am running the `examples/seq2seq/run_summarization.py` script with BartForConditionalGeneration. 

The script breaks whenever these two parameters are passed together:
- label_smoothing_factor
- pad_to_max_length

It seems that the source of this behaviour is setting collator to `default_data_collator` if `pad_to_max_length` is defined:
https://github.com/huggingface/transformers/blob/5f19c07a704eca4db376b56f950b729dcaa73039/examples/seq2seq/run_summarization.py#L469-L477

while `prepare_decoder_input_ids_from_labels` is only handled by DataCollatorForSeq2Seq:
https://github.com/huggingface/transformers/blob/5f19c07a704eca4db376b56f950b729dcaa73039/src/transformers/data/data_collator.py#L292-L294

It seems to be related with: [10452](https://github.com/huggingface/transformers/issues/10452), where passing a model argument to DataCollatorForSeq2Seq solves the problem 
`data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)`


This is more of a question than an issue as it is work in progress. A more general one would be:
Is the `default_data_collator` intended for use with seq2seq models (e.g: Bart), with special cases (like label smoothing) to be handled by `DataCollatorForSeq2Seq`?
Or should `DataCollatorForSeq2Seq` always be used with Seq2SeqTrainer in the future?


The problem arises when using:
* [x ] the official example scripts: (give details below)
examples/seq2seq/run_summarization.py
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [x ] an official GLUE/SQUaD task: (give the name) (xsum)
* [ ] my own task or dataset: (give details below)


## To reproduce

Steps to reproduce the behavior:

```
python examples/seq2seq/run_summarization.py \
    --model_name_or_path sshleifer/distilbart-xsum-12-3  \
    --do_train \
    --do_eval \
    --dataset_name xsum \
    --output_dir /tmp/output_dir \
    --per_device_train_batch_size=4 \
    --per_device_eval_batch_size=4 \
    --overwrite_output_dir \
    --predict_with_generate \
    --max_train_samples 500 \
    --max_val_samples 500 \
    --max_source_length 128 \
    --max_target_length 64 \
    --label_smoothing_factor 0.1 \
    --pad_to_max_length true    
```

Output:

```
Traceback (most recent call last):
  File "examples/seq2seq/run_summarization.py", line 595, in <module>
    main()
  File "examples/seq2seq/run_summarization.py", line 533, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/opt/anaconda3/envs/tensorflow2/lib/python3.6/site-packages/transformers/trainer.py", line 1082, in train
    tr_loss += self.training_step(model, inputs)
  File "/opt/anaconda3/envs/tensorflow2/lib/python3.6/site-packages/transformers/trainer.py", line 1472, in training_step
    loss = self.compute_loss(model, inputs)
  File "/opt/anaconda3/envs/tensorflow2/lib/python3.6/site-packages/transformers/trainer.py", line 1511, in compute_loss
    loss = self.label_smoother(outputs, labels)
  File "/opt/anaconda3/envs/tensorflow2/lib/python3.6/site-packages/transformers/trainer_pt_utils.py", line 439, in __call__
    smoothed_loss.masked_fill_(padding_mask, 0.0)
RuntimeError: The expanded size of the tensor (128) must match the existing size (64) at non-singleton dimension 1.  Target sizes: [4, 128, 1].  Tensor sizes: [4, 64, 1]
  0%|
```
<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior
Script works for a parameter set including:
- label_smoothing_factor
- pad_to_max_length

Or info which collator class should be used in the future
<!-- A clear and concise description of what you would expect to happen. -->



Discussion:
I think the `DataCollatorForSeq2Seq` should be used in all cases as it does more than just padding. If you want to suggest a PR with the fix, that would be more than welcome!
Assuming the goal is:
- using DataCollatorForSeq2Seq in Seq2SeqTrainer as default when no data_collator is provided, while keeping the remaining functionality unchanged, 

the first approach could be:
- providing Seq2SeqTrainer with an `__init__` method: 
    - instantiating a DataCollatorForSeq2Seq if no collator provided, and
    - calling Trainer's `__init__` and passing the instance along with other parameters. 
    
Something like:

```
class Seq2SeqTrainer(Trainer):
    
    def __init__(
        self,
        model: Union[PreTrainedModel, torch.nn.Module] = None,
        args: TrainingArguments = None,
        data_collator: Optional[DataCollator] = None,
        train_dataset: Optional[Dataset] = None,
        eval_dataset: Optional[Dataset] = None,
        tokenizer: Optional["PreTrainedTokenizerBase"] = None,
        model_init: Callable[[], PreTrainedModel] = None,
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
        callbacks: Optional[List[TrainerCallback]] = None,
        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
    ):    
        """
        Setting DataCollatorForSeq2Seq as default if no data_collator is provided.
        """

        if data_collator is None:
            # Perform validation and overwrite model with model_init before passing to collator,
            # as done in Trainer
            if tokenizer is None:
                raise RuntimeError(
                    "`tokenizer` parameter is required by the default `DataCollatorForSeq2Seq`"
                )
            if model is None and model_init is None:
                raise RuntimeError(
                    "`Trainer` requires either a `model` or `model_init` argument"
                )
            model_collator = model
            if model_init is not None:
                # No parameter handling for hyper-parameter search (trial)
                # Only passing the prepare_decoder_input_ids_from_labels function
                model_collator = model_init()

            data_collator = DataCollatorForSeq2Seq(tokenizer, model=model_collator)

        super().__init__(
            model,
            args,
            data_collator,
            train_dataset,
            eval_dataset,
            tokenizer,
            model_init,
            compute_metrics,
            callbacks,
            optimizers,
        )                          
 
```

Of course, I would need to look further into the code and the handling of other DataCollatorForSeq2Seq 
parameters like: `pad_to_multiple_of=8 if training_args.fp16 else None`

@sgugger, Thanks for the suggestion, it is very interesting;)
Mmm, I was thinking of an easier fix to just use that in the example script without necessary changing the default in `Seq2SeqTrainer`.

Resolving patch:
diff --git a/examples/seq2seq/run_summarization.py b/examples/seq2seq/run_summarization.py
--- a/examples/seq2seq/run_summarization.py
+++ b/examples/seq2seq/run_summarization.py
@@ -38,7 +38,6 @@
     HfArgumentParser,
     Seq2SeqTrainer,
     Seq2SeqTrainingArguments,
-    default_data_collator,
     set_seed,
 )
 from transformers.file_utils import is_offline_mode
@@ -466,15 +465,12 @@ def preprocess_function(examples):
 
     # Data collator
     label_pad_token_id = -100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
-    if data_args.pad_to_max_length:
-        data_collator = default_data_collator
-    else:
-        data_collator = DataCollatorForSeq2Seq(
-            tokenizer,
-            model=model,
-            label_pad_token_id=label_pad_token_id,
-            pad_to_multiple_of=8 if training_args.fp16 else None,
-        )
+    data_collator = DataCollatorForSeq2Seq(
+        tokenizer,
+        model=model,
+        label_pad_token_id=label_pad_token_id,
+        pad_to_multiple_of=8 if training_args.fp16 else None,
+    )
 
     # Metric
     metric = load_metric("rouge")


Test patch: