Inefficient zero checkpoint loading

As it was originally reported in https://github.com/huggingface/transformers/issues/12680 a user can easily train and save checkpoints, but don't have enough RAM to subsequently load that same checkpoint into memory.

One discussed approach is to use staggered loading so multiple processed won't try to use CPU memory at once, but the main issue appears in the fact that currently each process loads zero checkpoints for all ranks in deepspeed, even though most of the time each process uses just their own rank's checkpoint.

In zero2 this is needed when load_from_fp32_weights is True in load_state_dict to recover original fp32 weights which are spread out through multiple processes. In zero3 this option is ignored, but if I'm not mistaken it always recovers fp32 weights here.

The other intention for having more than just this rank's state is to be able to dynamically move from say DP=4 to DP=2 and so each process will then need 2 zero checkpoints.

But I think there should be also a support for a case where a user is doing the straightforward save/load, is ok to recover from fp16 weights and it shouldn't take much more additional memory than it took to save the checkpoint.

Not sure what is the best way to approach this, so would be happy to hear your ideas.

Based on profiling the main CPU memory-hungry call is:

self._load_zero_checkpoint(...)

it does return all of it when it's done. So we are talking only about a high tmp memory need.

Thank you.

@tjruwase

Start at self._load_zero_checkpoint(...) and load_state_dict, then profile the checkpoint-loading path described in the issue to understand its temporary CPU memory use. Done should support straightforward save/load with reduced extra memory while preserving the stated zero2 and zero3 checkpoint-loading behavior.