[BUG] load_checkpoint should load directly to gpu

Describe the bug

Currently HF Transformers integration users can finetune a model and save the checkpoint with given resources. However resuming from that same checkpoint requires much more CPU peak memory - which can be huge for large models, which prevents users from resuming their finetuning. (The current workaround is to add a huge swap file)

To Reproduce

I reproduced it as part of this bug report: https://github.com/huggingface/transformers/issues/17258

The full reproduction steps are here: https://github.com/huggingface/transformers/issues/17258#issuecomment-1133492187

I also verified that torch.load doesn't load everything in CPU memory when map_location="cpu" https://github.com/huggingface/transformers/issues/17258#issuecomment-1133522602

and I tracked the issue down to deepspeed loading those potentially huge zero checkpoints (70GB for gpt-j-6) into cpu memory first:

https://github.com/microsoft/DeepSpeed/blob/5208eb73da5269034ded69c4dd7c4bff81df81e7/deepspeed/runtime/engine.py#L2748

Expected behavior

save_checkpoint and load_checkpoint should require approximately the same amount of memory and should be lean and not need any CPU memory other than the size of the largest param or optim state since torch.load copies params via cpu.

With upcoming models like 176B the current implementation just won't work as it would require several TBs of CPU memory to load a zero checkpoint.

@tjruwase, @jeffra