CUDA out of memory when resuming from checkpoint (not always)
Author: lucadinidue
Pull request: no

## Environment info
- `transformers` version: 4.17.0.dev0
- Platform: Linux-4.18.0-147.51.2.el8_1.ppc64le-ppc64le-with-glibc2.17
- Python version: 3.8.2
- PyTorch version (GPU?): 1.7.0a0+e85d494 (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


### Who can help
@sgugger 

Model I am using BERT:

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: MLM
* [ ] my own task or dataset: (give details below)


_____________________________________________________________________

I'm currently pretraining BERT in MLM using the run_mlm.py with trainer in a distributed settig. To execute the script I run:

python -m torch.distributed.launch --nproc_per_node 4 run_mlm.py --model_type bert --train_file XXX --validation_file XXX --max_seq_length 128 --line_by_line True --output_dir XXX --do_train True --do_eval True --seed 2147483647 --save_steps 4000 --num_train_epochs 3 --per_device_train_batch_size 64 --logging_steps 100 --eval_steps 100 --remove_unused_columns=False --cache_dir XXX --logging_first_step True --evaluation_strategy steps --eval_accumulation_steps 1 --resume_from_checkpoint XXX/checkpoint-129000

I'm on a GPU cluster so, after 24h my script is stopped and I have to resume the pretraining. I successfully did it two times, then the third time I get this error:

Error message:

INFO|trainer.py:1277] 2022-03-10 17:37:39,275 >> ***** Running training *****
[INFO|trainer.py:1278] 2022-03-10 17:37:39,275 >>   Num examples = 51943464
[INFO|trainer.py:1279] 2022-03-10 17:37:39,275 >>   Num Epochs = 3
[INFO|trainer.py:1280] 2022-03-10 17:37:39,276 >>   Instantaneous batch size per device = 64
[INFO|trainer.py:1281] 2022-03-10 17:37:39,276 >>   Total train batch size (w. parallel, distributed & accumulation) = 256
[INFO|trainer.py:1282] 2022-03-10 17:37:39,276 >>   Gradient Accumulation steps = 1
[INFO|trainer.py:1283] 2022-03-10 17:37:39,276 >>   Total optimization steps = 608715
[INFO|trainer.py:1303] 2022-03-10 17:37:39,345 >>   Continuing training from checkpoint, will skip to saved global_step
[INFO|trainer.py:1304] 2022-03-10 17:37:39,345 >>   Continuing training from epoch 0
[INFO|trainer.py:1305] 2022-03-10 17:37:39,345 >>   Continuing training from global step 129000
[INFO|trainer.py:1307] 2022-03-10 17:37:39,346 >>   Will skip the first 0 epochs then the first 129000 batches in the first epoch. If this takes a lot of time, you can add the `--ignore_data_skip` flag to your launch command, but you will resume the training on data already seen by your model.
Skipping the first batches: 100%|██████████| 129000/129000 [8:05:02<00:00,  4.43it/s]   
  0%|          | 0/608715 [00:00<?, ?it/s]                  Traceback (most recent call last):
  File "run_mlm.py", line 566, in <module>:04:40,  4.43it/s]
    main()
  File "run_mlm.py", line 515, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/m100_work/IscrC_HiNLM/idini000/modules/transformers/src/transformers/trainer.py", line 1398, in train
    tr_loss_step = self.training_step(model, inputs)
  File "/m100_work/IscrC_HiNLM/idini000/modules/transformers/src/transformers/trainer.py", line 2000, in training_step
    loss.backward()
  File "/cineca/prod/opt/libraries/pytorch/1.7/cuda--10.2/lib/python3.8/site-packages/torch/tensor.py", line 221, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/cineca/prod/opt/libraries/pytorch/1.7/cuda--10.2/lib/python3.8/site-packages/torch/autograd/__init__.py", line 130, in backward
    Variable._execution_engine.run_backward(
RuntimeError: CUDA out of memory. Tried to allocate 972.00 MiB (GPU 3; 15.78 GiB total capacity; 11.59 GiB already allocated; 1.02 GiB free; 13.75 GiB reserved in total by PyTorch)


I've read that this was a problem with older versions of trainer but I updated it recently.

Thanks in advance!



Comments:
It might be that some of the older processes have left some junk in the GPU memory? Do you have a way to insure everything is clean before launching the training resuming from checkpoint?

This was probably what happened. I can't check if everything is clean but I launched it again and it worked fine.