PaddlePaddle
diff --git a/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/static/engine.py
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎python/paddle/distributed/auto_parallel/static/engine.py
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎python/paddle/distributed/passes/auto_parallel_recompute_pir.py
Copy file name to clipboard
+151-16Lines changed: 151 additions & 16 deletions b/‎python/paddle/distributed/passes/auto_parallel_recompute_pir.py
Copy file name to clipboard
+151-16Lines changed: 151 additions & 16 deletions
diff --git a/‎test/auto_parallel/pir/CMakeLists.txt
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎test/auto_parallel/pir/CMakeLists.txt
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
@@ -957,6 +957,10 @@ void BindOperation(py::module *m) {
                  attr_name,
                  pir::Int32Attribute::get(pir::IrContext::Instance(), val));
            })
+      .def("erase_attr",
+           [](Operation &self, std::string &attr_name) {
+             self.erase_attribute(attr_name);
+           })
       .def("attrs",
            [](Operation &self) -> py::dict {
              py::dict attrs_dict;
 
@@ -806,8 +806,9 @@ def _parallel_pir(self, mode):
         apply_mix2dist_pass(dist_program)
 
         if mode == "train" and self._strategy.recompute.enable:
+            config = copy.deepcopy(self._strategy.recompute.to_dict())
             auto_parallel_recompute_pir_pass = new_pass(
-                "auto_parallel_recompute_pir", {}
+                "auto_parallel_recompute_pir", config
             )
             auto_parallel_recompute_pir_pass.apply(
                 [dist_program], [startup_program]
 
@@ -33,17 +33,18 @@
 class AutoParallelRecomputePIRPass(PassBase):
     def __init__(self):
         super().__init__()
+        self.program_ops = []
 
     def _check_self(self):
         return True
 
     def _check_conflict(self, other_pass):
         return True
 
-    def get_fwd_bwd_ops(self, program):
+    def get_fwd_bwd_ops(self):
         fwd_ops = []
         bwd_ops = []
-        for op in program.global_block().ops:
+        for op in self.program_ops:
             if op.op_role == int(OpRole.Forward):
                 fwd_ops.append(op)
             elif op.op_role == int(OpRole.Backward):
@@ -72,81 +73,215 @@ def is_seed_used_by_dropout(self, seed_op):
             if used_op.name() in dropout_ops
         )
 
-    def get_segments(self, program):
+    def remove_outgoing_op(self, segment):
+        # An OP is considered an outgoing OP if all of results' user OPs are not in segment.
+        # These OPs do not participate in the backward gradient computation and therefore
+        # do not need to have a recomputation during backward.
+        segment_ops = [self.program_ops[idx] for idx in segment]
+        segment_len = len(segment)
+        for idx in range(segment_len - 1, 0, -1):
+            op = segment_ops[idx]
+            user_ops = set()
+            for res in op.results():
+                user_ops = user_ops | set(res.all_used_ops())
+
+            if user_ops & set(segment_ops):
+                continue
+            segment.pop(idx)
+            logger.info(
+                f"Remove outgoing OP '{op.name()}' from the segment for recomputation, as it does not participate in the backward."
+            )
+        return segment
+
+    def get_segments(self):
         # `fwd_recompute_id` indicates the ID assigned to the segment for
         # which the OP requires recompute.
         # A segment comprises all OPs within a program, ranging from the OP
         # with the minimum index to the OP with the maximum index, and all
         # these operations share the same `fwd_recompute_id`.
         segment_beg = {}
         segment_end = {}
-        max_op_id = len(program.global_block().ops)
-        for idx, op in enumerate(program.global_block().ops):
+        max_op_id = len(self.program_ops)
+        for idx, op in enumerate(self.program_ops):
+            # 1. Find the OPs marked with `fwd_recompute_id`.
             if not op.has_attr("fwd_recompute_id"):
                 continue
+            # 2. Delineate the segment range marked by `fwd_recompute_id`.
+            # Note: there may be some unmarked OPs in between.
             rc_id = op.attrs()["fwd_recompute_id"]
             if rc_id not in segment_beg:
                 segment_beg[rc_id] = max_op_id
                 segment_end[rc_id] = 0
             segment_beg[rc_id] = min(segment_beg[rc_id], idx)
             segment_end[rc_id] = max(segment_end[rc_id], idx)
 
+        # 3. Aggregate all segment information into a dictionary.
+        # The key is the id of the segment, which is used to uniquely identify each segment.
+        # The value is a list of indices of the segment OPs in `self.program_ops`.
         segments = {}
-        idx = 0
         assert len(segment_beg.keys()) == len(segment_end.keys())
         for segment_id, beg_id in segment_beg.items():
             assert segment_id in segment_end.keys()
             end_id = segment_end[segment_id]
             assert beg_id <= end_id
-            segment = []
-            for p_id in range(beg_id, end_id - 1):
-                segment.append(p_id)
-            segments[idx] = segment
-            idx += 1
+            segment = list(range(beg_id, end_id + 1))
+            # 4. Remove the outgoing OPs from the segment, as these OPs
+            # do not participate in the backward gradient computation.
+            segments[segment_id] = self.remove_outgoing_op(segment)
+            logger.info(
+                f"Segment ID {segment_id} contains {len(segment)} OPs, all of which will be recomputed."
+            )
         return segments
 
+    def get_op_name(self, op):
+        return op.name().split('.')[1]
+
+    def match_pattern(
+        self,
+        op,
+        visit,
+        fetch_id,
+        fetch_pattern,
+        target_pattern,
+        pre_len,
+        main_len,
+        count,
+        max_count,
+    ):
+        if count >= max_count:
+            return max_count
+        if len(fetch_pattern) > len(target_pattern):
+            return count
+        if self.get_op_name(op) != target_pattern[fetch_id]:
+            return count
+        if fetch_id == len(target_pattern) - 1:
+            for idx in range(pre_len, pre_len + main_len):
+                fetch_op = fetch_pattern[idx]
+                visit[fetch_op] = -1
+            refined_segment = list(set(visit.values()))
+            refined_segment.sort()
+            refined_segment = [idx for idx in refined_segment if idx != -1]
+            return count + 1
+        for res_val in op.results():
+            for user_op in res_val.all_used_ops():
+                fetch_pattern[fetch_id + 1] = user_op
+                count = self.match_pattern(
+                    op=user_op,
+                    visit=visit,
+                    fetch_id=fetch_id + 1,
+                    fetch_pattern=fetch_pattern,
+                    target_pattern=target_pattern,
+                    pre_len=pre_len,
+                    main_len=main_len,
+                    count=count,
+                    max_count=max_count,
+                )
+        return count
+
     def _apply_single_impl(self, main_program, startup_program, context=None):
-        segments = self.get_segments(main_program)
+        self.program_ops = list(main_program.global_block().ops)
+        # 1. Get the recompute segments information form program.
+        segments = self.get_segments()
         if len(segments) == 0:
-            logger.info("No segments found in PIR recompite pass.")
+            logger.info("No segments found in PIR recompite pass, skip pass.")
             return
+        # 2. Get the forward and backward OPs from program.
+        fwd_ops, bwd_ops = self.get_fwd_bwd_ops()
 
-        fwd_ops, bwd_ops = self.get_fwd_bwd_ops(main_program)
+        # 3. Refine the segments based on the patterns.
+        refined_ops_patterns = self.get_attr("refined_ops_patterns")
+        for refined_ops_pattern in refined_ops_patterns:
+            # 3.1 get the refined pattern information.
+            # refined_ops_patterns = pre_ops + main_ops + suf_ops
+            # `main_ops` pattern: it does not participate in backward recomputation
+            #                     and needs to be removed from the segment.
+            # `pre_ops` pattern: it serve only as markers and do require recomputation.
+            # `suf_ops` pattern: it serve only as markers and do require recomputation.
+            # `num` : it limits the maximum number of `main_ops` patterns identified
+            #         within each segment. A value of -1 represents all patterns.
+            num = int(refined_ops_pattern['num'])
+            num = num if num >= 0 else len(fwd_ops)
+            main_ops = refined_ops_pattern['main_ops']
+            pre_ops = refined_ops_pattern['pre_ops']
+            suf_ops = refined_ops_pattern['suf_ops']
+            pattern_ops = pre_ops + main_ops + suf_ops
 
+            for rc_id in segments.keys():
+                # 3.2 Identify and mark the first 'num' patterns in each segment.
+                # The dictionary 'op_idx_map' has keys as OP information.
+                # If an OP belongs to a pattern, its value in the dictionary is marked as -1.
+                op_idx_map = {
+                    self.program_ops[idx]: idx for idx in segments[rc_id]
+                }
+                pattern_count = 0
+                fetch_pattern = [None] * len(pattern_ops)
+                for idx in segments[rc_id]:
+                    op = self.program_ops[idx]
+                    fetch_pattern[0] = op
+                    pattern_count = self.match_pattern(
+                        op=self.program_ops[idx],
+                        visit=op_idx_map,
+                        fetch_id=0,
+                        fetch_pattern=fetch_pattern,
+                        target_pattern=pattern_ops,
+                        pre_len=len(pre_ops),
+                        main_len=len(main_ops),
+                        count=pattern_count,
+                        max_count=num,
+                    )
+                # 3.3 Refined segment to exclude the specified pattern.
+                refined_segment = list(set(op_idx_map.values()))
+                refined_segment.sort()
+                refined_segment = [idx for idx in refined_segment if idx != -1]
+                segments[rc_id] = refined_segment
+
+        # 4. Construct the segment for backward recomputation.
+        # 4.1 Build IrMapping to eplace forward value with backward recompute value.
         input_value = main_program.list_vars()
         value_map = paddle.pir.IrMapping()
         for val in input_value:
             value_map.add(val, val)
 
         for rc_id, segment in segments.items():
+            # 4.2 Find the insertion position for the backward segment,
+            # which should be before backward gradient computation.
             first_bwd_used_op = bwd_ops[-1]
             for idx in segment:
-                op = main_program.global_block().ops[idx]
+                op = self.program_ops[idx]
                 bwd_used_op = self.get_first_bwd_used_op(op, bwd_ops)
                 if first_bwd_used_op.id() > bwd_used_op.id():
                     first_bwd_used_op = bwd_used_op
 
             ori_segment_outputs = backward_utils.ValueSet()
             paddle.pir.set_insertion_point(first_bwd_used_op)
 
+            # 4.3 Clone the segment OPs and replace the forward
+            # value with backward recompute value.
             for idx in segment:
-                op = main_program.global_block().ops[idx]
+                op = self.program_ops[idx]
                 ori_segment_outputs.update(op.results())
 
+                # Random OPs should produce the same output before and after recomputation.
                 if self.is_seed_used_by_dropout(op):
                     continue
 
                 rc_op = op.clone(
                     value_map, paddle.pir.CloneOptions(False, True, True)
                 )
+                # The forward segment and the backward segment have the same segment ID.
+                if rc_op.has_attr("fwd_recompute_id"):
+                    rc_op.erase_attr("fwd_recompute_id")
+
                 rc_op.set_int_attr("bwd_recompute_id", rc_id)
 
+                # Updtate attributes.
                 if first_bwd_used_op.has_attr('op_role'):
                     rc_op.set_int_attr("op_role", first_bwd_used_op.op_role)
 
                 if first_bwd_used_op.has_attr('chunk_id'):
                     rc_op.set_int_attr("chunk_id", first_bwd_used_op.chunk_id)
 
+            # 4.4 Replace the forward value with backward recompute value.
             for ori_value in ori_segment_outputs:
                 rc_value = value_map.look_up(ori_value)
                 ori_value.replace_grad_users_with(rc_value, set(bwd_ops))
@@ -51,7 +51,7 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
     test_auto_parallel_replace_with_parallel_cross_entropy_pass
     PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 60)
   set_tests_properties(test_auto_parallel_recompute_pir_pass
-                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 60)
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 200)
   py_test_modules(
     test_eliminate_transpose_pass MODULES test_eliminate_transpose_pass ENVS
     FLAGS_enable_pir_in_executor=1)