[zero] restore fp16 params if no zero ckpts available

[jeffra]

In fine-tuning scenarios the user often wants to only load the model parameter checkpoints and may not have the zero optimizer states on disk since they are not needed. This fixes a bug where the weights are not properly restored from the ckpt if the zero optimizer states are not present on disk.

[stas00]

I confirm that it solved the problem! Thank you, Jeff!

Good to go to merge this into the big-science branch!

[stas00]

BTW, the checkpoint still requires mp_rank_*_model_states.pt and the program crashes without those.

AssertionError
        sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list)sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list)
  File "/gpfsssd/worksf/projects/rech/six/ura81os/code/deepspeed-big-science/deepspeed/runtime/state_dict_factory.py", line 162, in check_ckpt_list
    assert len(self.ckpt_list) > 0

Do we really need those? Or is this another legacy check?

[jeffra]

BTW, the checkpoint still requires mp_rank_*_model_states.pt and the program crashes without those.

AssertionError
        sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list)sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list)
  File "/gpfsssd/worksf/projects/rech/six/ura81os/code/deepspeed-big-science/deepspeed/runtime/state_dict_factory.py", line 162, in check_ckpt_list
    assert len(self.ckpt_list) > 0

Do we really need those? Or is this another legacy check?

I am pretty sure we still need these, since they’re associated with tensor parallelism model weights. @ShadenSmith to confirm wrt PP checkpoints though?

[stas00]

I think it needs those at least for getting the saved args out

[stas00]

So now this needs to be replayed to the big-science branch. Thank you!

[jeffra]

So now this needs to be replayed to the big-science branch. Thank you!

pushed this commit to big-science now :)