Update unified_checkpoint.py

PaddlePaddle · Dec 20, 2024 · 3c0059e · 3c0059e
1 parent 297dbce
commit 3c0059e
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/paddlenlp/trainer/unified_checkpoint/unified_checkpoint.py b/paddlenlp/trainer/unified_checkpoint/unified_checkpoint.py
@@ -435,7 +435,10 @@ def load_unified_optimizer(self, model, optimizer, resume_from_checkpoint):
             os.path.join(resume_from_checkpoint, SAFE_OPTIMIZER_INDEX_NAME)
         )
         if has_merge_optimizer_safetensors:
-            with open(os.path.join(resume_from_checkpoint, SAFE_OPTIMIZER_INDEX_NAME), "r") as f:
+            optimizer_index_file = os.path.join(resume_from_checkpoint, SAFE_OPTIMIZER_INDEX_NAME)
+            if distributed_isfile(optimizer_index_file):
+                distributed_file(optimizer_index_file)
+            with open(optimizer_index_file, "r") as f:
                 index = json.loads(f.read())
 
         # get quant ckpt info `ckpt_quant_stage` and `quant_ckpt_resume_times`