Revert "Remove timer training argument" (#4505)

Author: huangjiyi

examples/experiments/ernie_pretrain/ernie/pretrain.py

@@ -347,11 +347,12 @@ def main():
     if getattr(config.trainer_args, "dp_comm_overlap", False):
         logger.warning("Pipeline dp_comm_overlap and FusedLinearWithGradAdd can not be used at the same time.")
 
-    from paddle.distributed.fleet.meta_parallel.pipeline_parallel import (
-        PipelineParallel,
-    )
+    if getattr(config.trainer_args, "timer", False):
+        from paddle.distributed.fleet.meta_parallel.pipeline_parallel import (
+            PipelineParallel,
+        )
 
-    PipelineParallel.timer_printer = lambda _: None
+        PipelineParallel.timer_printer = lambda _: None
 
     def formatv(v):
         if isinstance(v, ListConfig):

paddleformers/cli/train/ernie_pretrain/workflow.py

@@ -356,11 +356,12 @@ def run_ernie_pretrain(model_args, data_args, generating_args, training_args):
     if getattr(training_args, "dp_comm_overlap", False):
         logger.warning("Pipeline dp_comm_overlap and FusedLinearWithGradAdd can not be used at the same time.")
 
-    from paddle.distributed.fleet.meta_parallel.pipeline_parallel import (
-        PipelineParallel,
-    )
+    if getattr(training_args, "timer", False):
+        from paddle.distributed.fleet.meta_parallel.pipeline_parallel import (
+            PipelineParallel,
+        )
 
-    PipelineParallel.timer_printer = lambda _: None
+        PipelineParallel.timer_printer = lambda _: None
 
     def formatv(v):
         if isinstance(v, ListConfig):

paddleformers/trainer/trainer.py

@@ -82,7 +82,6 @@
 from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
     DygraphShardingOptimizerV2,
 )
-from paddle.distributed.fleet.meta_parallel.pipeline_parallel import PipelineParallel
 from paddle.distributed.fleet.utils.hybrid_parallel_util import (
     fused_allreduce_gradients,
 )
@@ -446,7 +445,6 @@ def __init__(
 
             set_profile_timers(self.timers)
         self.runtime_timer = RuntimeTimer("RuntimeTimer")
-        PipelineParallel.timer_printer = lambda _: None
 
         self.model_wrapped = model
         self.model = model

paddleformers/trainer/training_args.py

@@ -36,7 +36,7 @@
 from ..utils.import_utils import is_paddlefleet_available
 from ..utils.log import logger
 from ..utils.pdc_sdk import FLASH_DEVICE
-from ..utils.tools import get_env_device, paddle_device
+from ..utils.tools import paddle_device
 from .trainer_utils import (
     IntervalStrategy,
     OptimizerNames,
@@ -1528,6 +1528,12 @@ class TrainingArguments:
             "help": "Enable splitting backward pass into stages to balance computation and reduce peak memory usage in model parallelism."
         },
     )
+    timer: bool = field(
+        default=False,
+        metadata={
+            "help": "Enable timing for pipeline parallel stages to profile and optimize communication/computation overlap."
+        },
+    )
     stage1_tensor_fusion: bool = field(
         default=False,
         metadata={
@@ -1945,6 +1951,7 @@ def __post_init__(self):
                                 "enable_delay_scale_loss",
                                 "enable_dp_comm_overlap",
                                 "enable_sharding_comm_overlap",
+                                "enable_timer",
                                 "enable_release_grads",
                                 "enable_clear_every_step_cache",
                                 "enable_overlap_p2p_comm",
@@ -1997,7 +2004,7 @@ def __post_init__(self):
                         "delay_scale_loss": True,  # TODO[Waynezee]: remove this config in the future
                         "dp_comm_overlap": enable_dp_comm_overlap,
                         "sharding_comm_overlap": self.enable_sharding_comm_overlap,
-                        "enable_timer": get_env_device() != "xpu",
+                        "enable_timer": self.timer,
                         "release_gradients": self.pp_release_grads or self.release_grads,
                         "overlap_p2p_comm": self.overlap_p2p_comm,
                         "clear_every_step_cache": self.clear_every_step_cache,
@@ -2428,6 +2435,7 @@ def is_context_parallel_supported():
                             "enable_delay_scale_loss",
                             # "enable_dp_comm_overlap",       # no implementation for auto_parallel
                             # "enable_sharding_comm_overlap", # no implementation for auto_parallel
+                            # "enable_timer",                 # no implementation for auto_parallel
                             # "disable_batch_p2p_comm",       # no implementation for auto_parallel
                             "enable_split_backward",
                             "auto_parallel_sync_shared_params",