CUDA error when training roBERTa from scratch with data parallel.
Author: balazik
Pull request: no

## Environment info
- `transformers` version: 3.5.0
- Platform:  Linux-5.4.0-52-generic-x86_64-Ubuntu 18.04.5 LTS
- Python version: 3.8.3
- PyTorch version (GPU?): 1.7.0+cu110
- Tensorflow version (GPU?): 2.3.1
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes

### Who can help
@LysandreJik  @sgugger

## Information

Model I am using (Bert, XLNet ...): roberta-base

The problem arises when using:
* [x] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Run transformers\examples\language-modeling\run_language_modeling.py as MLM task with WikiText-2 dataset (as mentioned in official README.md):
```shell
export CUDA_VISIBLE_DEVICES=0,1,2,3
export TRAIN_FILE=data/wiki.train.raw
export TEST_FILE=data/wiki.test.raw
python run_language_modeling.py \
    --output_dir=output \
    --model_type=roberta \
    --model_name_or_path=roberta-base \
    --do_train \
    --train_data_file=$TRAIN_FILE \
    --do_eval \
    --eval_data_file=$TEST_FILE \
    --mlm \
    --whole_word_mask
```
When we enable (with export CUDA_VISIBLE_DEVICES=0,1,2,3) all four GPUs (4xNVIDIA A100) training will throw RuntimeError: CUDA error: device-side assert triggered.
If we keep only one GPU enabled (with export CUDA_VISIBLE_DEVICES=0) training works flawlessly.

What we tried:
- run run_language_modeling.py with WikiText-2 dataset
- run run_language_modeling.py with custom sentenced text dataset
- official example for EsperBERTo in jupyter

All of the above mentioned have failed with the same error as soon as we enabled multiple GPUs!

Full error for run_language_modeling.py:
```console
./run_mlm_wiki.sh
11/18/2020 15:29:20 - WARNING - __main__ -   Process rank: -1, device: cuda:0, n_gpu: 4, distributed training: False, 16-bits training: False
11/18/2020 15:29:20 - INFO - __main__ -   Training/evaluation parameters TrainingArguments(output_dir='output', overwrite_output_dir=False, do_train=True, do_eval=True, do_predict=False, evaluate_during_training=False, evaluation_strategy=<EvaluationStrategy.NO: 'no'>, prediction_loss_only=False, per_device_train_batch_size=8, per_device_eval_batch_size=8, per_gpu_train_batch_size=None, per_gpu_eval_batch_size=None, gradient_accumulation_steps=1, eval_accumulation_steps=None, learning_rate=5e-05, weight_decay=0.0, adam_beta1=0.9, adam_beta2=0.999, adam_epsilon=1e-08, max_grad_norm=1.0, num_train_epochs=3.0, max_steps=-1, warmup_steps=0, logging_dir='runs/Nov18_15-29-20_a4000-20an1', logging_first_step=False, logging_steps=500, save_steps=500, save_total_limit=None, no_cuda=False, seed=42, fp16=False, fp16_opt_level='O1', local_rank=-1, tpu_num_cores=None, tpu_metrics_debug=False, debug=False, dataloader_drop_last=False, eval_steps=500, dataloader_num_workers=0, past_index=-1, run_name='output', disable_tqdm=False, remove_unused_columns=True, label_names=None, load_best_model_at_end=False, metric_for_best_model=None, greater_is_better=None)
/home/aime/.local/lib/python3.8/site-packages/transformers/modeling_auto.py:845: FutureWarning: The class `AutoModelWithLMHead` is deprecated and will be removed in a future version. Please use `AutoModelForCausalLM` for causal language models, `AutoModelForMaskedLM` for masked language models and `AutoModelForSeq2SeqLM` for encoder-decoder models.
  warnings.warn(
Some weights of RobertaForMaskedLM were not initialized from the model checkpoint at roberta-base and are newly initialized: ['lm_head.decoder.bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
/home/aime/.local/lib/python3.8/site-packages/transformers/tokenization_utils_base.py:1541: FutureWarning: The `max_len` attribute has been deprecated and will be removed in a future version, use `model_max_length` instead.
  warnings.warn(
/home/aime/.local/lib/python3.8/site-packages/transformers/data/datasets/language_modeling.py:40: FutureWarning: This dataset will be removed from the library soon, preprocessing should be handled with the ðŸ¤— Datasets library. You can have a look at this example script for pointers: https://github.com/huggingface/transformers/blob/master/examples/language-modeling/run_mlm.py
  warnings.warn(
11/18/2020 15:29:27 - INFO - filelock -   Lock 139760254204416 acquired on data/cached_lm_RobertaTokenizerFast_510_wiki.train.raw.lock
11/18/2020 15:29:27 - INFO - filelock -   Lock 139760254204416 released on data/cached_lm_RobertaTokenizerFast_510_wiki.train.raw.lock
11/18/2020 15:29:27 - INFO - filelock -   Lock 139760254204848 acquired on data/cached_lm_RobertaTokenizerFast_510_wiki.test.raw.lock
11/18/2020 15:29:27 - INFO - filelock -   Lock 139760254204848 released on data/cached_lm_RobertaTokenizerFast_510_wiki.test.raw.lock
/home/aime/.local/lib/python3.8/site-packages/transformers/trainer.py:277: FutureWarning: Passing `prediction_loss_only` as a keyword argument is deprecated and won't be possible in a future version. Use `args.prediction_loss_only` instead. Setting `args.prediction_loss_only=True
  warnings.warn(
  0%|                                                                                                                                                                                                                | 0/447 [00:00<?, ?it/s]/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:64: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.
  warnings.warn('Was asked to gather along dimension 0, but all '
  0%|â–                                                                                                                                                                                                       | 1/447 [00:07<58:40,  7.89s/it]Traceback (most recent call last):
  File "run_language_modeling.py", line 351, in <module>
    main()
  File "run_language_modeling.py", line 315, in main
    trainer.train(model_path=model_path)
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/trainer.py", line 775, in train
    tr_loss += self.training_step(model, inputs)
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/trainer.py", line 1112, in training_step
    loss = self.compute_loss(model, inputs)
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/trainer.py", line 1136, in compute_loss
    outputs = model(**inputs)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 161, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 171, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
    output.reraise()
  File "/home/aime/.local/lib/python3.8/site-packages/torch/_utils.py", line 428, in reraise
    raise self.exc_type(msg)
RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/modeling_roberta.py", line 894, in forward
    outputs = self.roberta(
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/modeling_roberta.py", line 686, in forward
    encoder_outputs = self.encoder(
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/modeling_roberta.py", line 421, in forward
    layer_outputs = layer_module(
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/modeling_roberta.py", line 341, in forward
    self_attention_outputs = self.attention(
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/modeling_roberta.py", line 273, in forward
    self_outputs = self.self(
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/modeling_roberta.py", line 203, in forward
    attention_probs = nn.Softmax(dim=-1)(attention_scores)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/modules/activation.py", line 1198, in forward
    return F.softmax(input, self.dim, _stacklevel=5)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/functional.py", line 1512, in softmax
    ret = input.softmax(dim)
RuntimeError: CUDA error: device-side assert triggered
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:108: cunn_ClassNLLCriterion_updateOutput_kernel: block: [0,0,0], thread: [0,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:108: cunn_ClassNLLCriterion_updateOutput_kernel: block: [0,0,0], thread: [1,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:108: cunn_ClassNLLCriterion_updateOutput_kernel: block: [0,0,0], thread: [2,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:108: cunn_ClassNLLCriterion_updateOutput_kernel: block: [0,0,0], thread: [3,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:108: cunn_ClassNLLCriterion_updateOutput_kernel: block: [0,0,0], thread: [4,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:108: cunn_ClassNLLCriterion_updateOutput_kernel: block: [0,0,0], thread: [5,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:108: cunn_ClassNLLCriterion_updateOutput_kernel: block: [0,0,0], thread: [6,0,0] Assertion `t >= 0 && t < n_classes` failed.
...
```

## Expected behavior
Training should start in parallel on all four Nvidia A100 GPUs without errors.

Comments:
The flag `whole_word_mask` cannot work wirth RoBERTa as it's only compatible with the BERT tokenizer and RoBERTa uses a different tokenizer. I'm surprised it was working on one GPU.

In any case, the `run_language_modeling.py` script is not maintained anymore, it has been replaced by new versions (`run_clm`, `run_mlm`, `run_plm`) that you can find in the `language-modeling` folder. Those new scripts are tested on a multi-GPU setup.

@sgugger Thanks for quick reply.

As I mentioned we tried other examples sa well.
#### EsperBERTo with Trainer class: 
https://github.com/huggingface/blog/blob/master/how-to-train.md.

Please don't mind the whole_word_mask flag because we also tried the model with & without it. Every time as soon as we use multiple GPUs we get above mentioned error:
```console
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:108: cunn_ClassNLLCriterion_updateOutput_kernel: block: [0,0,0], thread: [0,0,0] Assertion `t >= 0 && t < n_classes` failed.
```

#### Here we installed the newest transformers master with new version of language-modeling\run_mlm.py
Run parameters exact from README.md:
```shell
python run_mlm.py \
    --model_name_or_path roberta-base \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --do_train \
    --do_eval \
    --output_dir /tmp/test-mlm
```
And again we get the same error log (I shortened the error THCUNN part):
```console
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:108: cunn_ClassNLLCriterion_updateOutput_kernel: block: [0,0,0], thread: [29,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:108: cunn_ClassNLLCriterion_updateOutput_kernel: block: [0,0,0], thread: [30,0,0] Assertion `t >= 0 && t < n_classes` failed.
/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:108: cunn_ClassNLLCriterion_updateOutput_kernel: block: [0,0,0], thread: [31,0,0] Assertion `t >= 0 && t < n_classes` failed.
Traceback (most recent call last):
  File "run_mlm.py", line 392, in <module>
    main()
  File "run_mlm.py", line 362, in main
    trainer.train(model_path=model_path)
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/trainer.py", line 747, in train
    tr_loss += self.training_step(model, inputs)
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/trainer.py", line 1075, in training_step
    loss = self.compute_loss(model, inputs)
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/trainer.py", line 1099, in compute_loss
    outputs = model(**inputs)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 162, in forward
    return self.gather(outputs, self.output_device)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 174, in gather
    return gather(outputs, output_device, dim=self.dim)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 68, in gather
    res = gather_map(outputs)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 61, in gather_map
    return type(out)(((k, gather_map([d[k] for d in outputs]))
  File "<string>", line 7, in __init__
  File "/home/aime/.local/lib/python3.8/site-packages/transformers/file_utils.py", line 1305, in __post_init__
    for element in iterator:
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 61, in <genexpr>
    return type(out)(((k, gather_map([d[k] for d in outputs]))
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 55, in gather_map
    return Gather.apply(target_device, dim, *outputs)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/_functions.py", line 71, in forward
    return comm.gather(inputs, ctx.dim, ctx.target_device)
  File "/home/aime/.local/lib/python3.8/site-packages/torch/nn/parallel/comm.py", line 230, in gather
    return torch._C._gather(tensors, dim, destination)
RuntimeError: CUDA error: device-side assert triggered
  0%|▍                                                                                                                                                                                                     | 1/450 [00:09<1:07:51,  9.07s/it]
```

#### We even tried a custom simplified training code inspired by your docs.
https://huggingface.co/transformers/custom_datasets.html#fine-tuning-with-native-pytorch-tensorflow

```python
from transformers import LineByLineTextDataset
from transformers import DataCollatorForLanguageModeling, Trainer, TrainingArguments, RobertaConfig
from transformers import RobertaForMaskedLM
from transformers import AdamW
from torch.utils.data import DataLoader

dataset = LineByLineTextDataset(
    tokenizer=tokenizer,
    file_path=r'./data/oscar.eo.txt',
    block_size=512,
)

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=True, mlm_probability=0.15
)

training_args = TrainingArguments(
    output_dir=r'./EsperBERTo',
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=4,
    save_steps=10_000,
    save_total_limit=2,
    prediction_loss_only=True
)

print(training_args.n_gpu)

config = RobertaConfig(
    vocab_size=52_000,
    max_position_embeddings=514,
    num_attention_heads=12,
    num_hidden_layers=6,
    type_vocab_size=1,
)

args = training_args
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

model = RobertaForMaskedLM(config=config)
model = torch.nn.DataParallel(model)
model = model.to(device)

model.train()
train_loader = DataLoader(dataset, batch_size=4, shuffle=True)
optim = AdamW(model.parameters(), lr=5e-5)

for epoch in range(3):
    for batch in train_loader:
        optim.zero_grad()
        input_ids = batch['input_ids'].to(device)
        outputs = model(input_ids)
        loss = outputs[0]
        
        if args.n_gpu > 1:
            loss = loss.mean()  # mean() to average on multi-gpu parallel training
        
        loss.backward()
        optim.step()

model.eval()
```
But we always get the same error.

I don't have the error on my side with two GPUs and the same command, so I think the bug comes from something in your enviromnent. The fact the simple training loop also fails encourages me in the same direction.
If you try to use just two GPUs with `CUDA_VISIBLE_DEVICES`, does the problem persist? Maybe one of your GPUs is in a bad state?

After some laborious debugging we figured out that the problem was indeed in our HW configuration.

For others if your machine has an AMD EPYC 7402 you will probably needed to **disable IOMMU** (AMD I/O Virtualization Technology) in BIOS. After disabling all examples work.

I apologise for the inconvenience.