[REQUEST] An option to only save the model state_dict when save_checkpoint(), and how to manually save & load the model state_dict when using ZERO3

[BlinkDL]

In my training code, I only save & load the model state_dict (no optimizer states). I find this is good enough with a few steps of warmup, and saves lots of space when training a large model (as I have to save very frequently, like once per hour).

However one can't directly save the model state_dict when using ZERO3.

May I know how I can manually request all GPUs to save their partitioned state_dict, And how to manually load these state_dict to each GPU?

[tjruwase]

@BlinkDL, thanks for your question. I think the closest available feature involves rank 0 saving the full model state_dict rather than all ranks saving partitions. This requires enough CPU memory to hold the full model state_dict. Will that work for you?

[tjruwase]

@BlinkDL, are you still interested in this issue?

[BlinkDL]

@BlinkDL, thanks for your question. I think the closest available feature involves rank 0 saving the full model state_dict rather than all ranks saving partitions. This requires enough CPU memory to hold the full model state_dict. Will that work for you?

Can I save different parts of the state_dict in different ranks? That will be perfect.

[Thomas-MMJ]

Can I save different parts of the state_dict in different ranks? That will be perfect.

@tjruwase any more info?