[cherry-pick]【Auto-Parallel|Comm】fix set_skip_gc_vars to adapt to all scenarios and fix communication hang issue on GPU-H (#70687)

* 【Auto-Parallel】Refactor set_skip_gc_vars to adapt to all scenarios (#70615)

* 【Auto-Parallel | Comm】fix communication hang issue on GPU-H (#70360)

Author: liym27

paddle/fluid/pybind/pir.cc

@@ -946,6 +946,19 @@ void BindOperation(py::module *m) {
                                                 phi::IntArray(val));
              self.set_attribute(attr_name, attr);
            })
+      .def("set_str_array_attr",
+           [](Operation &self,
+              std::string &attr_name,
+              const std::vector<std::string> &val) {
+             std::vector<Attribute> val_attr;
+             for (auto &str : val) {
+               val_attr.emplace_back(
+                   StrAttribute::get(pir::IrContext::Instance(), str));
+             }
+             auto attr =
+                 pir::ArrayAttribute::get(pir::IrContext::Instance(), val_attr);
+             self.set_attribute(attr_name, attr);
+           })
       .def("set_str_attr",
            [](Operation &self, std::string &attr_name, std::string &val) {
              self.set_attribute(

python/paddle/distributed/auto_parallel/static/engine.py

@@ -34,7 +34,7 @@
 from paddle.distributed.passes.pass_base import new_pass
 from paddle.distributed.passes.pass_utils import (
     _split_program_into_forward_backward_optimize,
-    set_pir_skip_gc_vars,
+    set_skip_gc_vars,
 )
 from paddle.framework import (
     IrGraph,
@@ -943,7 +943,7 @@ def _parallel_pir(self, mode):
             opt_job.set_micro_batch_id(0)
             jobs.append(opt_job)
 
-            type_to_program = set_pir_skip_gc_vars(
+            type_to_program = set_skip_gc_vars(
                 self._strategy.gradient_merge.k_steps,
                 job_types,
                 sub_programs,

python/paddle/distributed/passes/pass_utils.py

@@ -268,6 +268,21 @@ def set_skip_gc_vars(num_micro_batches, job_types, sub_programs, jobs):
     thus a sub_program's vars might be used as the op's input of the later sub_program,
     and these vars cannot be gc after executing current sub_program.
     """
+    if paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
+        "FLAGS_enable_pir_api"
+    ]:
+        return _set_skip_gc_vars_in_pir(
+            num_micro_batches, job_types, sub_programs, jobs
+        )
+    else:
+        return _set_skip_gc_vars_in_old_ir(
+            num_micro_batches, job_types, sub_programs, jobs
+        )
+
+
+def _set_skip_gc_vars_in_old_ir(
+    num_micro_batches, job_types, sub_programs, jobs
+):
     assert num_micro_batches >= 1, "num_micro_batches needs to be >= 1"
     type_to_program = dict(zip(job_types, sub_programs))
 
@@ -300,32 +315,54 @@ def set_skip_gc_vars(num_micro_batches, job_types, sub_programs, jobs):
     return type_to_program
 
 
-def set_pir_skip_gc_vars(num_micro_batches, job_types, sub_programs, jobs):
+def _set_skip_gc_vars_in_pir(num_micro_batches, job_types, sub_programs, jobs):
     assert num_micro_batches >= 1, "num_micro_batches needs to be >= 1"
-    type_to_var_names = {}
     type_to_program = dict(zip(job_types, sub_programs))
-    for job_type, program in type_to_program.items():
-        type_to_var_names[job_type] = set()
-        ops = program.global_block().ops
-        for op in ops:
-            if op.name() == "builtin.shadow_output":
-                # if a value is renamed by shadow_output,
-                # it will be used by other sub_programs
-                type_to_var_names[job_type].add(op.attrs()["output_name"])
-        if job_type in ["backward", "backward_w"]:
-            assert (
-                len(type_to_var_names[job_type]) == 0
-            ), f"The {job_type} sub_program can't have skip_gc_vars. But it is {type_to_var_names[job_type]}."
 
+    # step1: Get all required vars of every sub_program that are non-persistable and not in op's no_need_buffer.
+    type_to_required_vars = {}
     no_need_buffer_vars = core.get_no_need_buffer_values(type_to_program)
+    for job_type, program in type_to_program.items():
+        required_vars = set()
+        persistable_vars = set()
+        for key in program.global_block().kwargs():
+            required_vars.add(key)
+        for op in program.global_block().ops:
+            for var in op.operands_source():
+                if var.has_name:
+                    required_vars.add(var.name)
+                    if var.persistable:
+                        persistable_vars.add(var.name)
+            for var in op.results():
+                if var.has_name:
+                    required_vars.add(var.name)
+                    if var.persistable:
+                        persistable_vars.add(var.name)
+        if job_type in no_need_buffer_vars:
+            required_vars -= no_need_buffer_vars[job_type]
+        required_vars -= persistable_vars
+        type_to_required_vars[job_type] = required_vars
 
-    for job_type, var_set in no_need_buffer_vars.items():
-        if len(var_set) > 0:
-            type_to_var_names[job_type] = type_to_var_names[job_type] - var_set
-
-    for job in jobs:
+    # step2: Set `skip_gc_vars` for each job
+    suffixed_required_vars = [set() for i in range(num_micro_batches)]
+    num_jobs = len(jobs)
+    for job_id in reversed(range(num_jobs)):
+        job = jobs[job_id]
         job_type = job.type()
-        job.set_skip_gc_vars(type_to_var_names[job_type])
+        required_vars = type_to_required_vars[job_type]
+        micro_batch_id = job.micro_batch_id()
+        skip_gc_vars = required_vars & suffixed_required_vars[micro_batch_id]
+        logger.debug(
+            f"Skip gc vars for {job_type}-({micro_batch_id}): {skip_gc_vars}"
+        )
+
+        if job_type in ["send_backward", "backward_w"]:
+            assert (
+                len(skip_gc_vars) == 0
+            ), f"When enabling pipeline parallelism strategy, the skip_gc_vars for {job_type} subprogram must be empty, but it is {skip_gc_vars}."
+
+        job.set_skip_gc_vars(skip_gc_vars)
+        suffixed_required_vars[micro_batch_id] |= required_vars
 
     return type_to_program