Optimize memory when loading checkpoint

[NouamaneTazi]

Offload memory states to CPU to avoid memory peak when loading checkpoint, and only move to GPU after first fwd-bwd

Details

It seems that the peak reserved memory is higher when we load a checkpoint because when we do the 1st fwd-bwd, which is the most heavy in terms of memory allocation, we already have the optimizer’s state loaded.
Whereas when we train from scratch we do fwd-bwd (optimizer’s state doesnt load yet), then free some memory, then reallocate the optimizer’s state

case of from scratch:
Memory usage: 9793.18MiB. Peak allocated 10465.18MiB. Peak reserved: 10972.00MiB # before fwd-bwd
Memory usage: 9927.32MiB. **Peak allocated 21428.54MiB**. Peak reserved: 21926.00MiB # after fwd-bwd
Memory usage: 9927.32MiB. Peak allocated 10695.32MiB. Peak reserved: 21926.00MiB # before optim step
Memory usage: 16456.16MiB. Peak allocated 16648.16MiB. Peak reserved: 21926.00MiB # after optim step + zero_grad
>>iter2
Memory usage: 16456.16MiB. Peak allocated 17128.16MiB. Peak reserved: 21926.00MiB # before fwd-bwd


case of resume training:
Memory usage: 16321.23MiB. Peak allocated 17153.23MiB. Peak reserved: 17392.00MiB # before fwd-bwd
Memory usage: 16455.37MiB. **Peak allocated 27956.59MiB**. Peak reserved: 28474.00MiB # after fwd-bwd
Memory usage: 16455.37MiB. Peak allocated 17223.37MiB. Peak reserved: 28474.00MiB # before optim step
Memory usage: 16456.16MiB. Peak allocated 16647.37MiB. Peak reserved: 28474.00MiB # after optim step + zero_grad
>>iter2
Memory usage: 16456.16MiB. Peak allocated 17128.16MiB. Peak reserved: 28428.00MiB # before fwd-bwd

TLDR; The memory demand when loading a checkpoint is larger by Optimizer’s state compared when training from scratch

Peak reserved (when loading checkpoint) = Peak reserved (when training from scratch) + Optimizer's state

In example above: 
28GB = 22GB + 6GB
optimizer's state is 6GB = 0.856*10**9 (local params) * 2 (adam's state) * 4 bytes / 1024 / 1024

[xrsrke]

LGTM. Left a small comment.