[AutoParallel][PIR] Fix 1F1B and VPP hang (#68141)

AndSonder · web-flow · commit 57587d4d476f · 2024-10-12T16:21:44.000+08:00
* add enable send recv

* Update pir.cc

* Update pass_utils.py

* fix hang

* remove useless codes

* fix

* remove useless codes

* remove useless codes

* fix vpp hang
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -3058,6 +3058,57 @@ void RemoveRedundantMemcpyAfterShadowFeed(pir::Block* block,
           VLOG(6) << *it;
         }
       }
+
+      pir::Value shadow_source = it->operand_source(0);
+      auto val_src_place =
+          shadow_source.type()
+              .dyn_cast<paddle::dialect::AllocatedDenseTensorType>()
+              .place();
+
+      if (shadow_value.use_count() >= 1 || val_src_place == phi::CPUPlace()) {
+        bool all_use_is_scalar = true;
+        for (auto use_it = shadow_value.use_begin();
+             use_it != shadow_value.use_end();
+             ++use_it) {
+          auto use_op = use_it->owner();
+
+          if (!use_op->isa<PhiKernelOp>()) {
+            all_use_is_scalar = false;
+            break;
+          }
+
+          auto op_info = ctx->GetRegisteredOpInfo(
+              use_op->dyn_cast<PhiKernelOp>().op_name());
+
+          if (!op_info) {
+            all_use_is_scalar = false;
+            break;
+          }
+
+          auto* op_info_concept =
+              op_info.GetInterfaceImpl<dialect::OpYamlInfoInterface>();
+          auto [input_infos, _1, _2, _3, _4] =
+              op_info_concept->get_op_info_(op_info.name());
+
+          uint32_t val_index = 0;
+          for (uint32_t index = 0; index < use_op->num_operands(); index++) {
+            if (use_op->operand_source(index) == shadow_value) {
+              val_index = index;
+              break;
+            }
+          }
+
+          if (!input_infos[val_index].is_mutable_attribute) {
+            all_use_is_scalar = false;
+            break;
+          }
+        }
+        if (all_use_is_scalar) {
+          // set dst_place_type for shadow_feed, 0 for cpu_place
+          VLOG(6) << "Reset shadow_feed dst_place_type to 0 for scalar use";
+          it->set_attribute("dst_place_type", pir::Int32Attribute::get(ctx, 0));
+        }
+      }
     }
   }
 }
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -1294,7 +1294,6 @@ def _init_comm(self):
                 all_process_groups = get_all_process_groups()
                 for process_group in all_process_groups:
                     process_group.instantiate()
-                pass
                 return
 
             # Traverse different rank programs and traverse each op of them,
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -162,7 +162,7 @@ def apply_partition_pass(program):
                 assign_out.get_defining_op().dist_attr = (
                     copy_op_attr_with_new_member(
                         assign_out.get_defining_op().dist_attr,
-                        new_chunk_id=prev_op.dist_attr.chunk_id,
+                        new_chunk_id=op.dist_attr.chunk_id,
                     )
                 )
 
diff --git a/python/paddle/distributed/auto_parallel/static/reshard_funcs/same_status_reshard_func.py b/python/paddle/distributed/auto_parallel/static/reshard_funcs/same_status_reshard_func.py
@@ -45,20 +45,24 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
         all_process_ids = sorted(all_process_ids)
 
         cur_global_rank = paddle.distributed.get_rank()
-        comm_group = new_process_group(all_process_ids)
+
+        for src, dst in zip(src_mesh.process_ids, dst_mesh.process_ids):
+            if src != dst:
+                new_process_group([src, dst], group_type="p2p")
+                new_process_group([dst, src], group_type="p2p")
 
         is_send = True
         for src, dst in zip(src_mesh.process_ids, dst_mesh.process_ids):
             if src == cur_global_rank:
-                dst_local_rank = all_process_ids.index(dst)
                 chunk_id = -1
                 if src_value.get_defining_op().dist_attr:
                     chunk_id = src_value.get_defining_op().dist_attr.chunk_id
 
+                comm_group = new_process_group([src, dst], group_type="p2p")
                 paddle._C_ops.send_v2(
                     src_value,
                     comm_group.id,
-                    dst_local_rank,
+                    comm_group.ranks.index(dst),
                     True,
                     False,
                 )
@@ -81,15 +85,15 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
                     if var.dist_attr().process_mesh == dst_mesh:
                         chunk_id = find_var_used_op_chunk_id(var)
 
-                src_local_rank = all_process_ids.index(src)
                 assert (
                     -1 not in dst_type.shape
                 ), "dynamic shape is not supported by pir-auto parallel yet."
 
+                comm_group = new_process_group([src, dst], group_type="p2p")
                 recv_value = paddle._C_ops.recv_v2(
                     dst_type._local_shape,
                     dst_type.dtype,
-                    src_local_rank,
+                    comm_group.ranks.index(src),
                     comm_group.id,
                     True,
                     False,
diff --git a/python/paddle/distributed/passes/pass_utils.py b/python/paddle/distributed/passes/pass_utils.py
@@ -466,6 +466,29 @@ def _create_program(src_block, dst_block, src_op, force_create=False):
             _create_var(src_block, dst_block, output_varname, force_create)
 
 
+def _pir_overlap_send_recv(program):
+    """
+    This function is used to replace the function '_insert_sync_for_fthenb_1f1b'.
+    The finally target of this function is as follows:
+        1. no need to insert the 'c_sync_calc' and 'c_sync_calc' operators
+        2. 'send_v2' operator uses 'dist_attr.execution_stream' to set stream of its own.
+        3. 'recv_v2' operator uses 'dist_attr.execution_stream' to set stream of its own.
+    """
+    for block in program.blocks:
+        for op in block.ops:
+            if op.name() == "pd_op.send_v2":
+                op.set_bool_attr("dynamic_shape", False)
+                op.set_bool_attr("use_calc_stream", True)
+                ring_id = op.attrs()["ring_id"]
+                op.set_execution_stream(f"send_stream_{ring_id}")
+                op.set_scheduling_priority(0)
+            elif op.name() == "pd_op.recv_v2":
+                op.set_bool_attr("dynamic_shape", False)
+                op.set_bool_attr("use_calc_stream", True)
+                op.set_execution_stream("recv_stream")
+                op.set_scheduling_priority(0)
+
+
 def _insert_sync_for_fthenb_1f1b(program, dist_context=None):
     """
     This implementation refers to lots of Paddle/python/paddle/base/optimizer.py.
@@ -805,6 +828,8 @@ def find_var_used_op_chunk_id(var):
 def _split_program_into_forward_backward_optimize(
     main_program, enable_send_recv_overlap=False
 ):
+    _pir_overlap_send_recv(main_program)
+
     forward_complete_op_role(main_program)
     complete_ops = main_program.global_block().ops
 
@@ -837,22 +862,25 @@ def _split_program_into_forward_backward_optimize(
                 # if this op's output is used, create the persistable
                 # var to be used in other programs.
                 result_in_opt = opt_ops[op_idx].result(idx)
+
                 if result_in_opt.use_empty() is False:
                     name = f"var_{op_idx}_{complete_ops[op_idx].name()}_{idx}"
                     paddle.pir.set_insertion_point_after(bwd_ops[op_idx])
                     paddle._C_ops.set_persistable_value(
                         bwd_ops[op_idx].result(idx), name
                     )
-                    # bwd_ops[op_idx].result(idx).persistable = True
+
                     new_result_var_in_opt = opt_block.add_kwarg(
                         name, result_in_opt.type()
                     )
                     new_result_var_in_opt.persistable = (
                         result_in_opt.persistable
                     )
+
                     opt_ops[op_idx].result(idx).replace_all_uses_with(
                         new_result_var_in_opt
                     )
+
             opt_ops[op_idx].erase()
         else:
             # in backward program, only the forward ops should be removed

Original file line number	Diff line number	Diff line change
`@@ -162,7 +162,7 @@ def apply_partition_pass(program):`
`162`	`162`	`assign_out.get_defining_op().dist_attr = (`
`163`	`163`	`copy_op_attr_with_new_member(`
`164`	`164`	`assign_out.get_defining_op().dist_attr,`
`165`		`- new_chunk_id=prev_op.dist_attr.chunk_id,`
	`165`	`+ new_chunk_id=op.dist_attr.chunk_id,`
`166`	`166`	`)`
`167`	`167`	`)`
`168`	`168`