return early from save_checkpoint

Author: adammoody

deepspeed/runtime/engine.py

@@ -2396,6 +2396,10 @@ def _save_checkpoint(self, save_dir, tag, client_state={}):
         module = self.module_state_dict()
         self._curr_ckpt_path = None
 
+        # Only a subset of procs may need to save the general model params
+        if not self.save_non_zero_checkpoint:
+            return
+
         state = dict(module=module,
                      buffer_names=self._get_buffer_names(),
                      optimizer=self.optimizer.state_dict()
@@ -2412,10 +2416,9 @@ def _save_checkpoint(self, save_dir, tag, client_state={}):
                      ds_version=version)
         state.update(client_state)
 
-        if self.save_non_zero_checkpoint:
-            log_dist(message=f'Saving model checkpoint: {save_path}', ranks=[0, 1])
-            #logger.info('Saving model checkpoint: {}'.format(save_path))
-            torch.save(state, save_path)
+        log_dist(message=f'Saving model checkpoint: {save_path}', ranks=[0, 1])
+        #logger.info('Saving model checkpoint: {}'.format(save_path))
+        torch.save(state, save_path)
 
     def _get_buffer_names(self):
         buffer_names = []