Model checkpointing in training

[zarzen]

Hi there,
I am using ZeRO-3 for training, and use engine.save_checkpoint for storing the temporal results.
But I found the saved model states are all empty.
I checked the implementation, and indeed before the save_checkpint, the model is partitioned: https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage3.py#L3114
Thus, it is reasonable to see the model states values are empty tensors (because here

DeepSpeed/deepspeed/runtime/engine.py

Line 2912 in dac9056

def _save_checkpoint(self, save_dir, tag, client_state={}):

, the system stores the module.state_dict.)

My questions are:

1. how can we restore the model parameter from the checkpoint result?
2. should we store the param.ds_tensor instead of param.data?