add eliminate_transpose arg

PaddlePaddle · Apr 28, 2024 · 3318bb5 · 3318bb5
1 parent ea2926c
commit 3318bb5
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/llm/llama/auto_parallel/run_pretrain_auto.py b/llm/llama/auto_parallel/run_pretrain_auto.py
@@ -86,6 +86,12 @@ class PreTrainingArguments(TrainingArguments):
             "help": "Enable fused_linear_param_grad pass, which should replace add_n_op with add_op for gradients accumulation."
         },
     )
+    eleminate_transpose: bool = field(
+        default=False,
+        metadata={
+            "help": "Enable eleminate_transpose pass, which should replace transpose with reshape when sequence parallel is enabled."
+        },
+    )
     job_schedule_profiler_start: int = field(
         default=-1,
         metadata={"help": "The step to start job_schedule_profiler."},
@@ -132,6 +138,11 @@ def __post_init__(self):
             fused_passes.enable = True
             fused_passes.fused_passes_list.append("fused_linear_param_grad_add_pass")
 
+        if self.eliminate_transpose:
+            fused_passes = self.strategy.fused_passes
+            fused_passes.enable = True
+            fused_passes.fused_passes_list.append("eliminate_transpose")
+
         logger.info(self.strategy)