Trajepl: nebula load fix

deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py

@@ -37,11 +37,14 @@ def create(self, tag):
         self.checkpoint = torch_nebula.Checkpoint(tag, -2)
 
     def save(self, state_dict, path: str):
+        log_dist(f"[Nebula] Create dummy files for loading.")
+        torch.save("", path)
+
         tag = _get_tag_from_path(path)
         partititon_name = os.path.basename(path)
-        logger.info(f"[Nebula] Saving {partititon_name} under tag{tag}...")
+        logger.info(f"[Nebula] Saving {partititon_name} under tag {tag}...")
         self.checkpoint.save(partititon_name, state_dict)
-        logger.info(f"[Nebula] Saved {partititon_name} under tag{tag}.")
+        logger.info(f"[Nebula] Saved {partititon_name} under tag {tag}.")
         return None
 
     def load(self, path: str, map_location=None):
@@ -58,21 +61,24 @@ def load(self, path: str, map_location=None):
 
         partititon_name = os.path.basename(path)
         logger.info(
-            f"[Nebula] Loading {path} under tag{tag} from {self.nebula_load_path}...")
+            f"[Nebula] Loading {path} under tag {tag} from nebula path {self.nebula_load_path}..."
+        )
 
         checkpoint = None
         if tag is None:
             checkpoint = torch_nebula.get_latest_checkpoint(
                 persist_path=self.nebula_load_path)
-            if checkpoint is None or (checkpoint is not None and checkpoint.tag == ''):
-                logger.warning(f"Unable to find latest valid checkpoint from Nebula!")
-                return None
         else:
             checkpoint = torch_nebula.get_checkpoint(tag=tag,
                                                      persist_path=self.nebula_load_path)
+
+        if checkpoint is None or (checkpoint is not None and checkpoint.tag == ''):
+            logger.warning(f"Unable to find latest valid checkpoint from Nebula!")
+            return None
+
         partition = checkpoint.load(partititon_name, map_location=map_location)
         logger.info(
-            f"[Nebula] Loaded {path} under tag{tag} from {self.nebula_load_path}.")
+            f"[Nebula] Loaded {path} under tag {tag} from {self.nebula_load_path}.")
         return partition
 
     def commit(self, tag):