[Unified Checkpoint] Fix load best checkpoint (PaddlePaddle#8935)

* fix load best

* update

paddlenlp/trainer/trainer.py

@@ -149,6 +149,7 @@
 from .utils import reshard as reshard_util
 from .utils.async_save import AsyncSaver
 from .utils.helper import (  # nested_truncate,
+    broadcast_dataset_rank0_model,
     broadcast_dp_optimizer,
     broadcast_moe_optimizer,
     distributed_concat,
@@ -1162,6 +1163,8 @@ def fused_allreduce_gradients_no_sync(paramlist, hcg):
                         self.state.best_model_checkpoint,
                         safe_serialization=True,
                     )
+                    if self.args.sharding_parallel_degree > 1 or self.args.data_parallel_degree > 1:
+                        broadcast_dataset_rank0_model(self.model)
                 else:
                     weight_name = PADDLE_WEIGHTS_NAME
                     best_model_path = os.path.join(
@@ -1204,6 +1207,8 @@ def _load_best_model_from_peft_checkpoint(self):
                 self.state.best_model_checkpoint,
                 safe_serialization=True,
             )
+            if self.args.sharding_parallel_degree > 1 or self.args.data_parallel_degree > 1:
+                broadcast_dataset_rank0_model(self.model)
             return
 
         convert_tp = False

paddlenlp/trainer/utils/helper.py

@@ -23,6 +23,7 @@
 import paddle
 import paddle.distributed as dist
 from paddle.distributed import fleet
+from paddle.distributed.parallel import sync_params_buffers
 
 from paddlenlp.utils.log import logger
 from paddlenlp.utils.nested import nested_broadcast_tensor_with_empty  # noqa: F401
@@ -309,3 +310,29 @@ def _broadcast_moe_optimizer_state(state_dict):
         state_dict = base_state_dict
         del base_state_dict
     return state_dict
+
+
+def broadcast_dataset_rank0_model(model):
+    if paddle.distributed.get_world_size() <= 1:
+        return
+
+    logger.info("Start broadcast model in sharding group or data parallel group.")
+    hcg = fleet.get_hybrid_communicate_group()
+    sharding_group = hcg.get_sharding_parallel_group()
+    dp_group = hcg.get_data_parallel_group()
+    if sharding_group.nranks > 1:
+        sync_params_buffers(
+            model,
+            sharding_group,
+            hcg.get_sharding_parallel_group_src_rank(),
+            is_model_parallel=False,
+            fuse_params=False,
+        )
+    if dp_group.nranks > 1:
+        sync_params_buffers(
+            model,
+            dp_group,
+            hcg.get_data_parallel_group_src_rank(),
+            is_model_parallel=False,
+            fuse_params=False,
+        )