[Unified Checkpoint] update async save logic (#9274) (#9275)

* update async save signal

* fix async save hang

Author: DesmonDay

paddlenlp/trainer/plugins/unified_checkpoint.py

@@ -136,21 +136,25 @@ def __init__(self, args):
         self._process_master_weight = None
         self._process_optimizer_weight = None
         self._lock = None
-        self._shared_save_path = None
         self._shared_save_model_flag = None
         self._shared_save_master_weight_flag = None
         self._shared_save_optimizer_flag = None
 
         if "async_save" in self.args.unified_checkpoint_config:
             self._lock = multiprocessing.Lock()
             self._shared_save_model_path = multiprocessing.Array("c", 100000)
+            self._shared_save_model_signal_path = multiprocessing.Array("c", 100000)
             self._shared_save_master_weight_path = multiprocessing.Array("c", 100000)
+            self._shared_save_master_weight_signal_path = multiprocessing.Array("c", 100000)
             self._shared_save_optimizer_path = multiprocessing.Array("c", 100000)
+            self._shared_save_optimizer_signal_path = multiprocessing.Array("c", 100000)
             self._shared_save_model_flag = multiprocessing.Array("i", 1)
             self._shared_save_master_weight_flag = multiprocessing.Array("i", 1)
             self._shared_save_optimizer_flag = multiprocessing.Array("i", 1)
 
-    def _file_save_async_or_sync(self, state_dict, path, is_sync=True, state_dict_type="model_weight"):
+    def _file_save_async_or_sync(
+        self, state_dict, path, signal_path=None, is_sync=True, state_dict_type="model_weight"
+    ):
         if is_sync:
             for k in list(state_dict.keys()):
                 if isinstance(state_dict[k], paddle.Tensor):
@@ -165,6 +169,7 @@ def _file_save_async_or_sync(self, state_dict, path, is_sync=True, state_dict_ty
                 meta_dict = self._meta_dict_model
                 shared_save_flag = self._shared_save_model_flag
                 shared_save_path = self._shared_save_model_path
+                shared_save_signal_path = self._shared_save_model_signal_path
                 if self._process_model_weight is None:
                     self._process_model_weight = multiprocessing.Process(
                         target=self._save_file_async_in_process,
@@ -173,12 +178,14 @@ def _file_save_async_or_sync(self, state_dict, path, is_sync=True, state_dict_ty
                             self._shm_model_weight.name,
                             self._shared_save_model_flag,
                             self._shared_save_model_path,
+                            self._shared_save_model_signal_path,
                             self._lock,
                             state_dict_type,
                             self.global_rank,
                         ),
                     )
                     self._process_model_weight.start()
+                process = self._process_model_weight
             elif state_dict_type == "master_weight":
                 if self._shm_master_weight is None:
                     self._meta_dict_master_weight, buffer_size = create_meta_dict(state_dict)
@@ -187,6 +194,7 @@ def _file_save_async_or_sync(self, state_dict, path, is_sync=True, state_dict_ty
                 meta_dict = self._meta_dict_master_weight
                 shared_save_flag = self._shared_save_master_weight_flag
                 shared_save_path = self._shared_save_master_weight_path
+                shared_save_signal_path = self._shared_save_master_weight_signal_path
                 if self._process_master_weight is None:
                     self._process_master_weight = multiprocessing.Process(
                         target=self._save_file_async_in_process,
@@ -195,6 +203,7 @@ def _file_save_async_or_sync(self, state_dict, path, is_sync=True, state_dict_ty
                             self._shm_master_weight.name,
                             self._shared_save_master_weight_flag,
                             self._shared_save_master_weight_path,
+                            self._shared_save_master_weight_signal_path,
                             self._lock,
                             "model_weight"
                             if "skip_save_model_weight" in self.args.unified_checkpoint_config
@@ -203,6 +212,7 @@ def _file_save_async_or_sync(self, state_dict, path, is_sync=True, state_dict_ty
                         ),
                     )
                     self._process_master_weight.start()
+                process = self._process_master_weight
             elif state_dict_type == "optimizer_weight":
                 if self._shm_optimizer_weight is None:
                     self._meta_dict_optim, buffer_size = create_meta_dict(state_dict)
@@ -211,6 +221,7 @@ def _file_save_async_or_sync(self, state_dict, path, is_sync=True, state_dict_ty
                 meta_dict = self._meta_dict_optim
                 shared_save_flag = self._shared_save_optimizer_flag
                 shared_save_path = self._shared_save_optimizer_path
+                shared_save_signal_path = self._shared_save_optimizer_signal_path
                 if self._process_optimizer_weight is None:
                     self._process_optimizer_weight = multiprocessing.Process(
                         target=self._save_file_async_in_process,
@@ -219,21 +230,26 @@ def _file_save_async_or_sync(self, state_dict, path, is_sync=True, state_dict_ty
                             self._shm_optimizer_weight.name,
                             self._shared_save_optimizer_flag,
                             self._shared_save_optimizer_path,
+                            self._shared_save_optimizer_signal_path,
                             self._lock,
                             state_dict_type,
                             self.global_rank,
                         ),
                     )
                     self._process_optimizer_weight.start()
+                process = self._process_optimizer_weight
 
             while True:  # wait until no process is saving.
                 flag_value = shared_save_flag[0]
                 if flag_value == 0:
                     break
+                if not process.is_alive():
+                    raise RuntimeError(f"The process that saves {state_dict_type} has been killed unexpectedly.")
                 time.sleep(0.5)
                 logger.info(f"Wait for the previous save process to finish saving {state_dict_type}")
             # only save model weight or save master weight, we enter this loop.
             self._reset_and_update(shared_save_path, path)
+            self._reset_and_update(shared_save_signal_path, signal_path)
             _traverse_copy_to_shm(state_dict, meta_dict, shm_state_dict.buf)
             with self._lock:
                 shared_save_flag[0] = 1
@@ -244,6 +260,7 @@ def _save_file_async_in_process(
         shm_name,
         shared_save_flag,
         shared_save_path,
+        shared_save_signal_path,
         lock,
         state_dict_type,
         global_rank,
@@ -257,11 +274,12 @@ def _save_file_async_in_process(
                 continue
             if flag_value == 1:  # need to save
                 path = shared_save_path[:].decode("utf-8").rstrip("\x00")
+                signal_path = shared_save_signal_path[:].decode("utf-8").rstrip("\x00")
                 logger.info(f"Start to async save {path}")
                 state_dict = _read_state_dict_from_shm(meta_dict, shm)  # numpy array
                 safe_save_file(state_dict, path, {"format": "np"})
                 del state_dict
-                saved_signal_path = os.path.join(os.path.dirname(path), f".{state_dict_type}.done.{global_rank}")
+                saved_signal_path = os.path.join(signal_path, f".{state_dict_type}.done.{global_rank}")
                 paddle.save(global_rank, saved_signal_path)
                 with lock:
                     shared_save_flag[0] = 0
@@ -276,7 +294,7 @@ def _reset_and_update(self, shared_array, new_value):
         encoded_value = new_value.encode("utf-8")
         shared_array[: len(encoded_value)] = encoded_value
 
-    def save_unified_checkpoint(self, model, optimizer, output_dir):
+    def save_unified_checkpoint(self, model, optimizer, output_dir, signal_dir=None):
         """save unified checkpoint
 
         Args:
@@ -313,6 +331,8 @@ def save_unified_checkpoint(self, model, optimizer, output_dir):
 
         save_directory = output_dir
         os.makedirs(save_directory, exist_ok=True)
+        if signal_dir is not None:
+            os.makedirs(signal_dir, exist_ok=True)  # only for async save
 
         # save model weights
         if not skip_save_model_weight:
@@ -325,6 +345,7 @@ def save_unified_checkpoint(self, model, optimizer, output_dir):
             self._file_save_async_or_sync(
                 state_dict,
                 path=os.path.join(save_directory, shard_file),
+                signal_path=signal_dir,
                 is_sync=is_sync_save,
                 state_dict_type="model_weight",
             )
@@ -393,7 +414,7 @@ def load_unified_checkpoint(self, model, optimizer, resume_from_checkpoint: str)
         if self.args.dataset_rank == 0:
             load_unified_checkpoint_locally(self.args, model, resume_from_checkpoint, safe_serialization=True)
 
-    def save_non_merge_optimizer(self, model, optimizer, output_dir):
+    def save_non_merge_optimizer(self, model, optimizer, output_dir, signal_dir):
         paddle.device.cuda.empty_cache()
         optim_state_dict = nested_copy(optimizer.state_dict())
         master_weights = None
@@ -432,12 +453,14 @@ def save_non_merge_optimizer(self, model, optimizer, output_dir):
         self._file_save_async_or_sync(
             optim_state_dict,
             path=os.path.join(output_dir, optimizer_name),
+            signal_path=signal_dir,
             is_sync=is_sync_save,
             state_dict_type="optimizer_weight",
         )
         self._file_save_async_or_sync(
             master_weights,
             path=os.path.join(output_dir, master_weights_name),
+            signal_path=signal_dir,
             is_sync=is_sync_save,
             state_dict_type="master_weight",
         )
@@ -484,22 +507,23 @@ def load_non_merge_optimizer(self, model, optimizer, resume_from_checkpoint):
 
         return returned_optim_state_dict
 
-    def save_unified_optimizer(self, model, optimizer, output_dir):
+    def save_unified_optimizer(self, model, optimizer, output_dir, signal_dir):
         """save unified optimizer
 
         Args:
             model (PretrainedModel): model used to get key mapping.
             optimizer (Optimizer): optimizer to save
             output_dir (str): Save directory.
+            signal_dir (str): Asynchronous saving signal directory.
 
         """
 
         if "ignore_merge_optimizer" in self.args.unified_checkpoint_config:
-            self.save_non_merge_optimizer(model, optimizer, output_dir)
+            self.save_non_merge_optimizer(model, optimizer, output_dir, signal_dir)
             return
 
         if paddle.distributed.get_world_size() <= 1:
-            self.save_single_card_optimizer(model, optimizer, output_dir)
+            self.save_single_card_optimizer(model, optimizer, output_dir)  # no need to save signal
             return
 
         # Split into naive optimizer params and master weights.
@@ -515,20 +539,24 @@ def save_unified_optimizer(self, model, optimizer, output_dir):
 
         save_directory = output_dir
         os.makedirs(save_directory, exist_ok=True)
+        if signal_dir is not None:
+            os.makedirs(signal_dir, exist_ok=True)
 
         is_sync_save = True
         if "async_save" in self.args.unified_checkpoint_config:
             is_sync_save = False
         self._file_save_async_or_sync(
             optim_state_dict,
             path=os.path.join(save_directory, shard_optim_file),
+            signal_path=signal_dir,
             is_sync=is_sync_save,
             state_dict_type="optimizer_weight",
         )
         if master_weight_state_dict is not None:
             self._file_save_async_or_sync(
                 master_weight_state_dict,
                 path=os.path.join(save_directory, shard_master_weight_file),
+                signal_path=signal_dir,
                 is_sync=is_sync_save,
                 state_dict_type="master_weight",
             )
@@ -716,14 +744,20 @@ def unlink_shared_memory(self):
 
         if self._shared_save_model_flag is not None:
             while self._shared_save_model_flag[0] > 0:  # async process is saving
+                if not self._process_model_weight.is_alive():
+                    raise RuntimeError("The process that saves model_weight has been killed unexpectedly.")
                 time.sleep(0.5)
             self._shared_save_model_flag[0] = -1
         if self._shared_save_master_weight_flag is not None:
             while self._shared_save_master_weight_flag[0] > 0:
+                if not self._process_master_weight.is_alive():
+                    raise RuntimeError("The process that saves master_weight has been killed unexpectedly.")
                 time.sleep(0.5)
             self._shared_save_master_weight_flag[0] = -1
         if self._shared_save_optimizer_flag is not None:
             while self._shared_save_optimizer_flag[0] > 0:
+                if not self._process_optimizer_weight.is_alive():
+                    raise RuntimeError("The process that saves optimizer_weight has been killed unexpectedly.")
                 time.sleep(0.5)
             self._shared_save_optimizer_flag[0] = -1
 
@@ -740,7 +774,8 @@ def unlink_shared_memory(self):
             self._shm_optimizer_weight.unlink()
             self._shm_optimizer_weight = None
 
-        dist.barrier()
+        if paddle.distributed.get_world_size() > 1:
+            dist.barrier()
 
 
 def load_unified_checkpoint_locally(args, model, resume_from_checkpoint: str, safe_serialization=False):