PaddlePaddle
diff --git a/‎python/paddle/distributed/passes/pass_utils.py
Copy file name to clipboard
+8-156Lines changed: 8 additions & 156 deletions b/‎python/paddle/distributed/passes/pass_utils.py
Copy file name to clipboard
+8-156Lines changed: 8 additions & 156 deletions
@@ -1043,164 +1043,16 @@ def used_by_grad_merge_add(value):
     return ["backward_w"]
 
 
-def _split_program_for_vpp(
-    program, num_model_chunks, oprole_names, split_bw=False
-):
-    place = _get_device()
-    if isinstance(place, paddle.framework.CUDAPlace):
-        place = paddle.framework.CUDAPlace(
-            paddle.distributed.ParallelEnv().dev_id
-        )
-    cur_place = paddle.base.libpaddle.Place()
-    cur_place.set_place(place)
-
-    def get_var_name(op_idx, result_idx):
-        result_value = all_ops[op_idx].result(result_idx)
-        all_used_ops = result_value.all_used_ops()
-        shadow_output_op_used = None
-        for op in all_used_ops:
-            if op.name() == "builtin.shadow_output":
-                shadow_output_op_used = op
-
-        if shadow_output_op_used is not None:
-            var_name = shadow_output_op_used.attrs()["output_name"]
-        else:
-            var_name = f"var_{op_idx}_{all_ops[op_idx].name()}_{result_idx}"
-        return var_name
-
-    def add_persistable_var(op_idx, program_type):
-        all_program_types = list(type_to_program.keys())
-        following_program_types = all_program_types[
-            all_program_types.index(program_type) + 1 :
-        ]
-        op_num_results = type_to_ops[program_type][op_idx].num_results()
-        op_name = type_to_ops[program_type][op_idx].name()
-
-        for idx in range(op_num_results):
-            var_name = None
-            for type in reversed(following_program_types):
-                op_result = type_to_ops[type][op_idx].result(idx)
-                if op_result.use_empty():
-                    continue
-
-                # if this op's output is used, create the persistable
-                # var to be used in other programs.
-                if var_name is None:
-                    if op_name in ["pd_op.data", "builtin.parameter"]:
-                        var_name = op_result.name
-                    else:
-                        var_name = get_var_name(op_idx, idx)
-                        if "var_" in var_name:
-                            paddle.pir.set_insertion_point_after(
-                                type_to_ops[program_type][op_idx]
-                            )
-                            paddle._C_ops.set_persistable_value(
-                                type_to_ops[program_type][op_idx].result(idx),
-                                var_name,
-                            )
-                            # type_to_ops[program_type][op_idx].result(idx).persistable = True
-
-                program_block = type_to_program[type].global_block()
-                new_result_var = program_block.add_kwarg(
-                    var_name, op_result.type()
-                )
-                new_result_var.place_attr = cur_place
-                new_result_var.persistable = op_result.persistable
-                type_to_ops[type][op_idx].result(idx).replace_all_uses_with(
-                    new_result_var
-                )
-
-        for type in following_program_types:
-            type_to_ops[type][op_idx].erase()
-
-    type_to_program = OrderedDict()
-    type_to_ops = OrderedDict()
-
-    # Step1: create programs and ops for each type
-    for type in oprole_names:
-        if type == "optimizer":
-            type_to_program["optimizer"] = program.clone()
-            type_to_ops["optimizer"] = (
-                type_to_program["optimizer"].global_block().ops
-            )
-        else:
-            chunk_ids = list(range(num_model_chunks))
-            if "backward" in type:
-                chunk_ids.reverse()
-            for chunk_id in chunk_ids:
-                type_to_program[type + str(chunk_id)] = program.clone()
-                type_to_ops[type + str(chunk_id)] = (
-                    type_to_program[type + str(chunk_id)].global_block().ops
-                )
-
-    # Step2: delete the ops not belong to the type
-    # 1. delete ops
-    # 2. add persistable var used between multiple programs
-    all_ops = program.global_block().ops
-    chunk_ids = list(range(num_model_chunks))
-
-    bwd_pattern_ops_type = []
-    for idx in range(len(all_ops) - 1, -1, -1):
-        op = all_ops[idx]
-        op_role = op.op_role
-        op_chunk_id = op.chunk_id
-        # Step2.1: infer chunk_id for ops that don't have chunk_id
-        if op_role != int(OpRole.Optimize) and op_chunk_id == -1:
-            op_chunk_id = infer_chunk_id(idx, all_ops, False)
-            if op_chunk_id == -1:
-                raise ValueError(
-                    f"Cannot infer chunk_id for op {op.name()} at index {idx}"
-                )
-
-        # Step2.2: identify the job_type of the op
-        if op_role == int(OpRole.Optimize):
-            job_type = "optimizer"
-        elif op_role == int(OpRole.Backward) and split_bw:
-            if len(bwd_pattern_ops_type) == 0:
-                bwd_pattern_ops_type = _pir_get_backward_op_type(all_ops, idx)
-            job_type = bwd_pattern_ops_type.pop()
-        elif op_role == int(OpRole.Backward) and (not split_bw):
-            job_type = "backward"
-        elif op_role == int(OpRole.Forward):
-            job_type = "forward"
-        else:
-            raise ValueError(
-                f"The op[{op.name()}]'s op role: {op_role} isn't one of Forward, Backward or Optimizer."
-            )
-
-        # Step2.3: delete ops not belong to the type
-        for type in oprole_names:
-            if type == job_type:
-                break
-            for chunk_id in chunk_ids:
-                type_to_ops[type + str(chunk_id)][idx].erase()
-
-        chunk_order = range(0, op_chunk_id)
-        if "backward" in job_type:
-            chunk_order = range(num_model_chunks - 1, op_chunk_id, -1)
-        for chunk_id in chunk_order:
-            type_to_ops[job_type + str(chunk_id)][idx].erase()
-
-        # Step2.4: add persistable var used between multiple programs
-        if job_type != "optimizer":
-            add_persistable_var(idx, job_type + str(op_chunk_id))
-
-    return list(type_to_program.keys()), list(type_to_program.values())
-
-
-def _pir_program_for_vpp(
-    program, num_model_chunks, split_bw=False, enable_send_recv_overlap=False
-):
-    _pir_overlap_send_recv(program)
+def _create_program_and_ops(program, job_type, chunk_id=None):
+    if chunk_id is not None:
+        program_name = f"{job_type}{chunk_id}"
+    else:
+        program_name = job_type
 
-    oprole_names = ["forward", "backward", "optimizer"]
-    if split_bw:
-        oprole_names = ["forward", "backward_b", "backward_w", "optimizer"]
+    cloned_program = program.clone()
+    ops = cloned_program.global_block().ops
 
-    program_types, programs = _split_program_for_vpp(
-        program, num_model_chunks, oprole_names, split_bw=split_bw
-    )
-    return program_types, programs
+    return program_name, cloned_program, ops
 
 
 def _program_for_vpp(