CUDA OOM when loading optimizer state dict
### System Info

I am finetuning GPT NEO 1.3B with 1 GPU and 24GB VRAM. 

From scratch, model weights load and train fine with <1GB of memory to spare.

When loading from checkpoint, I run into issues loading optimizer states.
https://github.com/huggingface/transformers/blob/04bfac83b793b757e7b33188f88eebe21ac65ef7/src/transformers/trainer.py#L2434-L2436

If I hard-code `map_location ='cpu'`, the OOM goes away and I am able to resume training. After some reading, this was caused by the optimizer state dict being associated with 'gpu:0' at the time of checkpoint saving. I also found another ticket adding device logic to `map_location=self.args.device` in the first place due to CPU OOM for the PR author. I was wondering if the interface needs to be change in order to accommodate both use cases.

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

see above. 

Otherwise from using run_clm.py from the repo without deepspeed. 

### Expected behavior

Resume from checkpoint should fit in the same vram footprint as finetune from scratch. 
cc @sgugger <|||||>I tried to find the reference to the issue mentioning a CPU OOM but couldn't find it. Do you have a link handY?
We could always load the weights on the CPU then move them which would be a tiny bit slower, but this just happens once. With models getting bigger and bigger it's probably makes sens to have this as a default behavior though.<|||||>Sorry I lost track of the issue with the CPU OOM but it was related to GPU memory in aggregate (deepspeed iirc) is greater than CPU. 

https://github.com/huggingface/transformers/issues/3730#issuecomment-629563466 is also relevant. My understanding is that `self.args.device` is where the model lives which is usually `gpu` or `tpu`. This puts optimizer state state_dict loading also to device with limited memory. Is there any reason to not always put this into CPU? My understanding is that if the model is trained using GPU they will be implicitly copied over during optimization but avoids the setup OOM. 