[Distributed] Support param and moment sync for pp shared layer (#10832)

* [Distributed] Support param and moment sync for pp shared layer

* fix name

Author: zhangyuqin1998

paddlenlp/trainer/training_args.py

@@ -1037,9 +1037,7 @@ class TrainingArguments:
     )
     zcc_ema_loss_threshold: Optional[float] = field(
         default=None,
-        metadata={
-            "help": "If set not None, only do EMA when the training loss is smaller than the threshold value"
-        },
+        metadata={"help": "If set not None, only do EMA when the training loss is smaller than the threshold value"},
     )
     save_tokenizer: Optional[bool] = field(
         default=True,
@@ -1354,6 +1352,8 @@ def __post_init__(self):
                                 "enable_offload_queue",
                                 "use_dualpipev",
                                 "enable_dynamic_shape",
+                                "sync_moment",
+                                "sync_param",
                             ]:
                                 raise ValueError(
                                     f"Found unknown pipeline mode config {x}, accpet config is disable_p2p_cache_shape, disable_partial_send_recv."
@@ -1404,6 +1404,18 @@ def __post_init__(self):
                         "use_dualpipev": "use_dualpipev" in pipeline_parallel_config,
                         "enable_dynamic_shape": "enable_dynamic_shape" in pipeline_parallel_config,
                     }
+
+                    pp_sync_param = "sync_param" in pipeline_parallel_config
+                    pp_sync_moment = "sync_moment" in pipeline_parallel_config
+
+                    if pp_sync_param:
+                        logger.info("setting pp sync_param")
+                        strategy.hybrid_configs["pp_configs"].sync_param = True
+
+                    if pp_sync_moment:
+                        logger.info("setting pp sync_moment")
+                        strategy.hybrid_configs["pp_configs"].sync_moment = True
+
                     if dygraph_pp_configs["dp_comm_overlap"]:
                         raise ValueError("overlap has accuracy issue")  # TODO: fix `overalap` + `delay_scale` issue