fix p2p initialize hang (#58376)

wentaoyu · web-flow · commit 76f773885b1c · 2023-10-25T22:26:57.000+08:00
* fix p2p initialize hang

* fix code style
diff --git a/paddle/fluid/distributed/collective/process_group_nccl.cc b/paddle/fluid/distributed/collective/process_group_nccl.cc
@@ -573,13 +573,13 @@ std::shared_ptr<ProcessGroup::Task> ProcessGroupNCCL::Recv(
                 << NCCLDTypeToString(phi::ToNCCLDataType(tensor->dtype()))
                 << ", src_in_group: " << src_rank
                 << ", ncclcomm: " << comm_context->GetNcclComm()
-                << ", stream: " << stream << ", rank_in_group: " << rank_
-                << ", nranks: " << size_ << ", offset: " << offset
-                << ", sync_op: " << sync_op
+                << ", stream: " << stream
+                << ", rank_in_group: " << rank_in_group << ", nranks: " << size_
+                << ", offset: " << offset << ", sync_op: " << sync_op
                 << ", use_calc_stream: " << use_calc_stream
                 << GetGroupMessage();
 
-        comm_context->Recv(tensor, tensor->numel(), src_rank, stream);
+        comm_context->Recv(tensor, tensor->numel(), rank_in_group, stream);
       },
       src_rank,
       *tensor,
@@ -610,15 +610,15 @@ std::shared_ptr<ProcessGroup::Task> ProcessGroupNCCL::Send(
                        phi::ToNCCLDataType(tensor_maybe_partial.dtype()))
                 << ", dst_in_group: " << dst_rank
                 << ", ncclcomm: " << comm_context->GetNcclComm()
-                << ", stream: " << stream << ", rank_in_group: " << rank_
-                << ", nranks: " << size_ << ", offset: " << offset
-                << ", sync_op: " << sync_op
+                << ", stream: " << stream
+                << ", rank_in_group: " << rank_in_group << ", nranks: " << size_
+                << ", offset: " << offset << ", sync_op: " << sync_op
                 << ", use_calc_stream: " << use_calc_stream
                 << GetGroupMessage();
 
         comm_context->Send(tensor_maybe_partial,
                            tensor_maybe_partial.numel(),
-                           dst_rank,
+                           rank_in_group,
                            stream);
       },
       dst_rank,
@@ -656,11 +656,9 @@ void ProcessGroupNCCL::CreateNCCLEnvCache(const Place& place,
                                           const std::string& store_key,
                                           CommType comm_type,
                                           int p2p_rank) {
-  ncclUniqueId nccl_id;
-
   VLOG(3) << "init nccl rank_in_group: " << rank_ << ", nranks: " << size_
           << ", gid: " << gid_ << ", place key: " << place_key
-          << ", nccl uniqueid: " << SerializeNCCLUniqueId(nccl_id);
+          << ", store_key: " << store_key;
 
   for (size_t i = 0; i < s_group_call_counter; ++i) {
     NCCL_CHECK(phi::dynload::ncclGroupEnd());
diff --git a/paddle/phi/core/distributed/comm_context_manager.cc b/paddle/phi/core/distributed/comm_context_manager.cc
@@ -33,6 +33,7 @@
 #include "paddle/phi/backends/context_pool.h"
 #include "paddle/phi/common/memory_utils.h"
 #include "paddle/phi/core/distributed/nccl_comm_context.h"
+#include "paddle/phi/core/distributed/nccl_tools.h"
 #endif
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
 #include "paddle/phi/core/distributed/xccl_comm_context.h"
@@ -78,12 +79,14 @@ void CommContextManager::CreateNCCLCommContext(
     std::memcpy(&nccl_id, nccl_id_wrapper.data(), nccl_id_wrapper.size());
   }
 
-  if (p2p_opt && p2p_opt->rank > 0) {
+  if (p2p_opt) {
     rank = p2p_opt->rank;
-  }
-  if (p2p_opt && p2p_opt->num_ranks > 0) {
     size = p2p_opt->num_ranks;
   }
+  VLOG(3) << "init NCCLCommContext rank: " << rank << ", size: " << size
+          << ", unique_comm_key: " << unique_comm_key
+          << ", unique_key: " << unique_key
+          << ", nccl_id: " << SerializeNCCLUniqueId(nccl_id);
   auto nccl_comm_context =
       std::make_unique<NCCLCommContext>(rank, size, nccl_id);
   if (CommContextManager::device_id != -1) {
