Data collator not working with Masking language model of Bart type
Author: ahadda5
Pull request: no
Labels: name: bug

### System Info

**System Info**

- transformers version : 4.20.1
- Ubuntu : "20.04.3 LTS 
- Python: 3.7 
- PyTorch + GPU:  torch:1.11+    cuda:11.4



### Who can help?

@patil-suraj @sgugger 



### 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

Trying to train a BART model using masking. 
The model type is BartForConditionalGeneration. 
Before trying it on a custom dataset, I wanted to try it on the given example [here](https://huggingface.co/course/chapter7/3?fw=pt#the-dataset).   (which is in fact similar [to here](https://github.com/huggingface/notebooks/blob/main/examples/language_modeling.ipynb) 

My code is an exact match to the provided example except the choice of a model relevant to my task. 

So instead of 
```
model_checkpoint = "distilbert-base-uncased"

model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)
```

We have 
```
model_checkpoint = "memray/bart-wikikp"

model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)
```


Based on the documentation this unsupervised approach is viable if one wants to fine-tune the model for a specific domain. Therefore, before fine-tuning, _Masked language modelling helps  acquaint the model with the new corpus first._  Also, in the documentation _there is no mention of specific-tasks against using masking training_ e.g. text classification or question answering vs my text generation task.  

When running the trainer.train I get the following flavors of CUDA error  (if I use the accelerator the error becomes RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.)

> “***** Running training *****
  Num examples = 10000
  Num Epochs = 3
  Instantaneous batch size per device = 32
  Total train batch size (w. parallel, distributed & accumulation) = 32
  Gradient Accumulation steps = 1
  Total optimization steps = 939
  0%|          | 0/939 [00:00<?, ?it/s]/opt/conda/conda-bld/pytorch_1646755953518/work/aten/src/ATen/native/cuda/Indexing.cu:703: indexSelectLargeIndex: block: [42,0,0], thread: [32,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
/opt/conda/conda-bld/pytorch_1646755953518/work/aten/src/ATen/native/cuda/Indexing.cu:703: indexSelectLargeIndex: block: [42,0,0], thread: [33,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
/opt/conda/conda-bld/pytorch_1646755953518/work/aten/src/ATen/native/cuda/Indexing.cu:703: indexSelectLargeIndex: block: [42,0,0], thread: [34,0,0] Assertion `srcIndex < srcSelectDimSize` failed.

>Traceback (most recent call last):
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3457, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-13-7c826244a095>", line 1, in <module>
    trainer.train()
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/transformers/trainer.py", line 1413, in train
    ignore_keys_for_eval=ignore_keys_for_eval,
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/transformers/trainer.py", line 1651, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/transformers/trainer.py", line 2345, in training_step
    loss = self.compute_loss(model, inputs)
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/transformers/trainer.py", line 2377, in compute_loss
    outputs = model(**inputs)
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/transformers/models/bart/modeling_bart.py", line 1368, in forward
    return_dict=return_dict,
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/transformers/models/bart/modeling_bart.py", line 1229, in forward
    return_dict=return_dict,
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/transformers/models/bart/modeling_bart.py", line 850, in forward
    output_attentions=output_attentions,
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/transformers/models/bart/modeling_bart.py", line 327, in forward
    output_attentions=output_attentions,
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/transformers/models/bart/modeling_bart.py", line 191, in forward
    query_states = self.q_proj(hidden_states) * self.scaling
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/haddad/.conda/envs/hugg/lib/python3.7/site-packages/torch/nn/modules/linear.py", line 103, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: CUDA error: CUBLAS_STATUS_NOT_INITIALIZED when calling `cublasCreate(handle)`


### Expected behavior

IF i remove the standard datacollator from the training args , it works(it runs trainer.train till the end) !  That is the only difference i saw between the two official examples (links above) , but this means we are losing that functionality!
I didn't try with my own custom collator.

Comments:
So commenting the collator argument , otherwise CUDA faces lots of issues whether with the accelator or not.  

```
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=downsampled_dataset["train"],
    eval_dataset=downsampled_dataset["test"],
    # data_collator=data_collator,
)
```

This means there is an indexing problem somewhere. Please use the [forums](https://discuss.huggingface.co/) to help debug your code. You can also look at [this video](https://youtu.be/L-WSwUWde1U) which explains how to debug a training loop step by step to get more information about where the error comes from.

Solution [here](https://discuss.huggingface.co/t/masked-language-model-for-bart-not-bert/19945)