Fix nebula in save_16bit_model issue (deepspeedai#3023)

FreyaRao · CodeSkull-1 · web-flow · commit a78d6b89e046 · 2023-03-23T09:45:42.000-07:00
Co-authored-by: Qinghuan Rao &lt;qinghuanrao@microsoft.com&gt;
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
@@ -3534,9 +3534,15 @@ def save_16bit_model(self, save_dir, save_filename="pytorch_model.bin"):
         else:
             state_dict = self.module.state_dict()
 
+        tag = f"global_step{self.global_steps}"
+        tag = str(tag)
+        self.checkpoint_engine.create(tag)
+
         if dist.get_rank() == 0:
             self.checkpoint_engine.makedirs(save_dir, exist_ok=True)
-            logger.info(f"Saving model weights to {path}")
+            logger.info(f"Saving model weights to {path}, tag: {tag}")
             self.checkpoint_engine.save(state_dict, path)
 
+        self.checkpoint_engine.commit(tag)
+
         return True
