OOM error when resuming training from a checkpoint
# 🐛 Bug

Some previous issues #2954 described a memory leak when resuming training from a checkpoint.  I still get an OOM error when resuming training from a checkpoint.
Hi, unfortunately, we have no way of helping without having more information. What scripts are you using? What model? What Python version? What transformers version? 

It would be wonderful if you could use the issue template and describe exactly the issue so that we may help.<|||||>Hi, I'm sorry. I'm using [examples/run_language_modeling.py](https://github.com/huggingface/transformers/blob/master/examples/run_language_modeling.py). After some more experiments, I noticed as #2954 did, that the OOM error only happens when resuming on a checkpoint in multi GPU training. When resuming using a single GPU there's no error.

Command example:
```
python -m torch.distributed.launch --nproc_per_node 8 run_language_modeling.py --output_dir=./output/ --model_type=gpt2 --model_name_or_path=gpt2-large --do_train --train_data_file=./data/training.txt --per_gpu_train_batch_size 1 --num_train_epochs 3 --fp16
```

Error:

```python
Traceback (most recent call last):
  File "run_language_modeling.py", line 992, in <module>
    main()
  File "run_language_modeling.py", line 942, in main
    global_step, tr_loss = train(args, train_dataset, model, tokenizer)
  File "run_language_modeling.py", line 428, in train
    optimizer.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "optimizer.pt")))
  File "/opt/conda/lib/python3.6/site-packages/torch/serialization.py", line 590, in load
    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
  File "/opt/conda/lib/python3.6/site-packages/torch/serialization.py", line 764, in _legacy_load
    result = unpickler.load()
  File "/opt/conda/lib/python3.6/site-packages/torch/serialization.py", line 726, in persistent_load
    deserialized_objects[root_key] = restore_location(obj, location)
  File "/opt/conda/lib/python3.6/site-packages/torch/serialization.py", line 190, in default_restore_location
    result = fn(storage, location)
  File "/opt/conda/lib/python3.6/site-packages/torch/serialization.py", line 170, in _cuda_deserialize
    return storage_type(obj.size())
  File "/opt/conda/lib/python3.6/site-packages/torch/cuda/__init__.py", line 478, in _lazy_new
    return super(_CudaBase, cls).__new__(cls, *args, **kwargs)
RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 31.72 GiB total capacity; 1.89 GiB already allocated; 10.88 MiB free; 1.92 GiB reserved in total by PyTorch)
```

@LysandreJik Can you please reopen the issue?<|||||>I ran into this issue as well when restarting from a checkpoint. 
I think this is a bug in [trainer.py](https://github.com/huggingface/transformers/blob/3e0f06210646a440509efa718b30d18322d6a830/src/transformers/trainer.py#L334) :
```
optimizer.load_state_dict(torch.load(os.path.join(model_path, "optimizer.pt")))
```
Loading from `optimizer.pt` causes `optimizer` to be mapped to the same device as the saved `optimizer.pt`. In this case it's always `cuda:0`(saved by local master), which puts all optimizers on gpu0, causing OOM.

Changing it to 
```
optimizer.load_state_dict(torch.load(os.path.join(model_path, "optimizer.pt"), map_location=self.args.device))
```
solved it for me.<|||||>This looks correct, can you open a PR?<|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
