More fixes to how FP16_Optimizer loads state from a checkpoint

Author: igor0

deepspeed/runtime/fp16/fused_optimizer.py

@@ -457,7 +457,7 @@ def load_state_dict(self, state_dict, load_optimizer_states=True):
             optimizer.load_state_dict(checkpoint['optimizer'])
         """
 
-        if state_dict is None:
+        if state_dict is None or "dynamic_loss_scale" not in state_dict:
             state_dict = self.state_dict()
             self.refresh_fp32_params()
             return
@@ -492,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."

setup.py

@@ -167,7 +167,7 @@ def op_enabled(op_name):
 
 start_time = time.time()
 
-setup(name='deepspeed',
+setup(name='deepspeed-igor',
       version=version_str,
       description='DeepSpeed library',
       long_description=readme_text,