RAM Out-Of-Memory error with `run_mlm.py` when loading a 6Gb json dataset
Author: RomanCast
Pull request: no

### System Info

- `transformers` version: 4.23.0.dev0
- Platform: Linux-4.18.0-305.65.1.el8_4.x86_64-x86_64-with-glibc2.28
- Python version: 3.9.13
- Huggingface_hub version: 0.9.1
- PyTorch version (GPU?): 1.12.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes (DDP)

### Who can help?

@sgugger because it might be an error with the Trainer

### Information

- [ ] The official example scripts
- [X] 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

When using the `run_mlm.py` script with almost no modifications using a Json dataset of ~6Gb, my job gets killed by SLURM. The stack trace looks like the following : 

```
Traceback (most recent call last):
  File "/gpfswork/rech/rax/commun/miniconda3/envs/artificial_data/lib/python3.9/runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/gpfswork/rech/rax/commun/miniconda3/envs/artificial_data/lib/python3.9/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/gpfswork/rech/rax/commun/miniconda3/envs/artificial_data/lib/python3.9/site-packages/torch/distributed/launch.py", line 193, in <module>
    main()
  File "/gpfswork/rech/rax/commun/miniconda3/envs/artificial_data/lib/python3.9/site-packages/torch/distributed/launch.py", line 189, in main
    launch(args)
  File "/gpfswork/rech/rax/commun/miniconda3/envs/artificial_data/lib/python3.9/site-packages/torch/distributed/launch.py", line 174, in launch
    run(args)
  File "/gpfswork/rech/rax/commun/miniconda3/envs/artificial_data/lib/python3.9/site-packages/torch/distributed/run.py", line 752, in run
    elastic_launch(
  File "/gpfswork/rech/rax/commun/miniconda3/envs/artificial_data/lib/python3.9/site-packages/torch/distributed/launcher/api.py", line 131, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/gpfswork/rech/rax/commun/miniconda3/envs/artificial_data/lib/python3.9/site-packages/torch/distributed/launcher/api.py", line 245, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
=======================================================
src/training/run_mlm.py FAILED
-------------------------------------------------------
Failures:
[1]:
  time      : 2023-01-06_11:50:42
  host      : r10i0n8-ib0
  rank      : 1 (local_rank: 1)
  exitcode  : -9 (pid: 870220)
  error_file: <N/A>
  traceback : Signal 9 (SIGKILL) received by PID 870220
-------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2023-01-06_11:50:42
  host      : r10i0n8-ib0
  rank      : 0 (local_rank: 0)
  exitcode  : -9 (pid: 870219)
  error_file: <N/A>
  traceback : Signal 9 (SIGKILL) received by PID 870219
=======================================================
slurmstepd: error: Detected 14 oom-kill event(s) in StepId=275597.batch. Some of your processes may have been killed by the cgroup out-of-memory handler.
```

What's weird is that the dataset is not that large (unfortunately I can't share it), and it worked fine with other datasets of similar size (for instance using ~4Gb of OSCAR text datasets). It also worked fine with another Json dataset of the same type but 50 times smaller.

I am putting the issue here because the last messages logged are the following : 

```
[INFO|trainer.py:502] 2023-01-06 11:50:25,673 >> max_steps is given, it will override any value given in num_train_epochs
[INFO|trainer.py:556] 2023-01-06 11:50:25,673 >> Using cuda_amp half precision backend
[INFO|trainer.py:725] 2023-01-06 11:50:25,674 >> The following columns in the training set don't have a corresponding argument in `XLMRobertaForMaskedLM.forward` and have been ignored: special_tokens_mask. If special_tokens_mask are not expected by `XLMRobertaForMaskedLM.forward`,  you can safely ignore this message.
```

which seems to indicate that it happens inside the Trainer method `_inner_training_loop`.

For reference, here is the command I'm running : 

```shell
n_gpus=4

python -m torch.distributed.launch --nproc_per_node $n_gpus \
    src/training/run_mlm.py \
    --model_type xlm-roberta \
    --config_overrides max_position_embeddings=512 \
    --tokenizer_name tokenizers/br_tokenizer_30k \
    --train_file ${TRAIN_FILE} \
    --is_split_into_words \
    --line_by_line \
    --use_auth_token \
    --validation_split_percentage 5 \
    --max_eval_samples 5000 \
    --max_seq_length 128 \
    --eval_steps 500 \
    --output_dir $OUTPUT_DIR \
    --do_train \
    --do_eval \
    --load_best_model_at_end \
    --metric_for_best_model "loss" \
    --greater_is_better False \
    --evaluation_strategy steps \
    --per_device_train_batch_size $per_device_batch_size \
    --per_device_eval_batch_size $per_device_batch_size \
    --gradient_accumulation_steps $(( $per_device_total_batch_size / $per_device_batch_size )) \
    --fp16 \
    --learning_rate 2e-5 \
    --weight_decay 1e-2 \
    --max_steps 100_000 \
    --warmup_steps 10_000 \
    --logging_dir $TENSORBOARD_DIR \
    --logging_steps 200 \
    --save_strategy steps \
    --save_steps 1000 \
    --save_total_limit 2 \
    --preprocessing_num_workers $(( 8 * $n_gpus )) \
    --report_to tensorboard \
    --seed 42
```

Some of the modifications to `run_mlm.py` involve using pre-tokenized datasets instead of raw text datasets by using the option `--is_split_into_words`.

Any idea why this happens or how to circumvent it ? 

### Expected behavior

I would expect the Trainer to start training without OOM failure, especially given the fact that `run_mlm.py` tokenizes and groups the sentences in the dataset without OOM issues.

Comments:
Debugging seems to indicate that this OOM error happens when wrapping the model with DDP:

https://github.com/huggingface/transformers/blob/48d4e147d824efab97637947709d5aa67c809b3d/src/transformers/trainer.py#L1446-L1451

However, I don't really see why the data size would impact this line...

I found out that this happens with 4 GPUs, but not with 2 or 1 GPUs, so my workaround at the moment is to train with only 2 GPUs which is slower but doable.

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.