fix empty state_dict

DesmonDay · Nov 6, 2024 · 16857c1 · 16857c1
1 parent 2f0b407
commit 16857c1
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/paddlenlp/trainer/unified_checkpoint/async_handler.py b/paddlenlp/trainer/unified_checkpoint/async_handler.py
@@ -77,6 +77,11 @@ def _file_save_async_or_sync(
                     state_dict[k] = state_dict.pop(k).cpu().numpy()
             safe_save_file(state_dict, path, metadata={"format": "np"})
         else:
+            if len(state_dict.keys()) == 0:
+                saved_signal_path = os.path.join(signal_path, f".{state_dict_type}.done.{self.global_rank}")
+                paddle.save(self.global_rank, saved_signal_path)
+                return
+
             if state_dict_type == "model_weight":
                 if self._shm_model_weight is None:
                     self._meta_dict_model, buffer_size = create_meta_dict(state_dict)