Merge pull request #1 from igor0/igor

Igor

Author: igor0

deepspeed/runtime/engine.py

@@ -1596,6 +1596,7 @@ def _load_checkpoint(self,
     def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True):
         zero_sd_list = self._get_all_zero_checkpoints(load_dir, tag)
         if zero_sd_list is None:
+            self.optimizer._restore_from_fp16_weights()
             return
 
         self.optimizer.load_state_dict(

deepspeed/runtime/fp16/fused_optimizer.py

@@ -456,6 +456,12 @@ def load_state_dict(self, state_dict, load_optimizer_states=True):
             model.load_state_dict(checkpoint['model'])
             optimizer.load_state_dict(checkpoint['optimizer'])
         """
+
+        if state_dict is None or "dynamic_loss_scale" not in state_dict:
+            state_dict = self.state_dict()
+            self.refresh_fp32_params()
+            return
+
         # I think it should actually be ok to reload the optimizer before the model.
         self.dynamic_loss_scale = state_dict["dynamic_loss_scale"]
         self.cur_scale = state_dict["cur_scale"]
@@ -486,7 +492,7 @@ def load_state_dict(self, state_dict, load_optimizer_states=True):
                 self.fp32_groups_flat, state_dict["fp32_groups_flat"]
             ):
                 current.data.copy_(saved.data)
-        except RuntimeError as error:
+        except (RuntimeError, KeyError) as error:
             print(error)
             print(
                 "Error in loading fp32 model parameters!\nRefreshing fp32 model params from the model's fp16 params instead. This may incur some precision loss."