Load model and optimizet states on CPU to void OOMs

[sgugger]

This PR makes sure we load the model and optimizer state dicts (in the load_state method) on the CPU. Otherwise, those might be tensors that are on GPU 0 (since save_state only saves on process 0), which may result in an OOM: at some point we may have the optimizer state properly loaded on GPU 0 (by process 0) but also loaded a second time in GPU 0 by process 1 during load_state_dict, so twice that size on poor GPU 0.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[muellerzr]

LG2M! Nice catch! Looks like we might be able to link #296?

[sgugger]

That's no linked per se since in this issue, the user was using optimizer.load_state_dict and not accelerator.load_state.