fix chunk id of recv_v2 by find the next non -1 recusively (#69201)

zhiqiu · web-flow · commit e7a1936377e8 · 2024-11-07T23:58:36.000+08:00
* fix chunk id of recv_v2 by find the next non -1 recusively

* reduce recursion
diff --git a/python/paddle/distributed/passes/pass_utils.py b/python/paddle/distributed/passes/pass_utils.py
@@ -23,9 +23,6 @@
 from paddle.distributed.auto_parallel.static.dist_attribute import (
     OperatorDistAttr,
 )
-from paddle.distributed.auto_parallel.static.mix_to_dist_pass import (
-    dist_skip_op_list,
-)
 from paddle.distributed.auto_parallel.static.utils import (
     get_logger,
     is_backward_op,
@@ -803,16 +800,24 @@ def get_chunk_id(op_idx):
 
 
 def find_var_used_op_chunk_id(var):
-    all_used_ops = var.all_used_ops()
-    for used_op in all_used_ops:
-        if used_op.name() in dist_skip_op_list:
-            for output_var in used_op.results():
-                chunk_id = find_var_used_op_chunk_id(output_var)
-                if chunk_id != -1:
-                    return chunk_id
-        elif used_op.dist_attr and used_op.dist_attr.chunk_id != -1:
-            return used_op.dist_attr.chunk_id
-    return -1
+    visited = set()
+
+    def dfs(var):
+        all_used_ops = var.all_used_ops()
+        for used_op in all_used_ops:
+            if used_op in visited:
+                return -1
+            visited.add(used_op)
+            if used_op.dist_attr and used_op.dist_attr.chunk_id != -1:
+                return used_op.dist_attr.chunk_id
+            else:
+                for output_var in used_op.results():
+                    chunk_id = dfs(output_var)
+                    if chunk_id != -1:
+                        return chunk_id
+        return -1
+
+    return dfs(var)
 
 
 def _split_program_into_forward_backward_optimize(
