cherry-pick from (PaddlePaddle#58338)

Fix nccl trace * fix nccl_async_trace destruct problem when train finished * update * format code style
hitywt · Nov 14, 2023 · cfefaa1 · cfefaa1
1 parent 3d44f57
commit cfefaa1
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 2 deletions.
diff --git a/paddle/fluid/distributed/collective/process_group_nccl.cc b/paddle/fluid/distributed/collective/process_group_nccl.cc
@@ -130,6 +130,11 @@ ProcessGroupNCCL::ProcessGroupNCCL(
       pg_timeout_(timeout) {
   LOG(INFO) << "ProcessGroupNCCL pg_timeout_ " << pg_timeout_;
 }
+ProcessGroupNCCL::~ProcessGroupNCCL() {
+  LOG(INFO) << "ProcessGroupNCCL destruct ";
+  auto& comm_task_manager = phi::distributed::CommTaskManager::GetInstance();
+  comm_task_manager.Stop();
+}
 
 void ProcessGroupNCCL::GroupStart() {
   NCCL_CHECK(phi::dynload::ncclGroupStart());

diff --git a/paddle/fluid/distributed/collective/process_group_nccl.h b/paddle/fluid/distributed/collective/process_group_nccl.h
@@ -82,7 +82,8 @@ class ProcessGroupNCCL final : public ProcessGroupWithStream {
                    int rank,
                    int size,
                    int gid,
-                   int64_t timeout = 30 * 60 * 1000);
+                   int64_t timeout = 20 * 1000);
+  ~ProcessGroupNCCL();
 
   std::string GetBackendName() const override { return "NCCL"; }
 

diff --git a/paddle/phi/core/distributed/comm_task_manager.cc b/paddle/phi/core/distributed/comm_task_manager.cc
@@ -74,6 +74,17 @@ void CommTaskManager::CommTaskEnqueue(std::shared_ptr<CommTask> comm_task) {
   }
 }
 
+void CommTaskManager::Stop() {
+  terminated_.store(true);
+
+  LOG(INFO) << "CommTaskManager stopped begin.";
+  if (comm_task_loop_thread_.joinable()) {
+    comm_task_loop_thread_.join();
+    comm_task_list_cv_.notify_one();
+  }
+  LOG(INFO) << "CommTaskManager stopped.";
+}
+
 void CommTaskManager::CommTaskLoop() {
   bool done = false;
   while (!terminated_.load() || !done) {
@@ -100,7 +111,11 @@ void CommTaskManager::CommTaskLoop() {
         }
         iter = comm_task_list_.erase(iter);
       } else {
-        ++iter;
+        if (task->IsStarted() && task->IsCompleted()) {
+          iter = comm_task_list_.erase(iter);
+        } else {
+          ++iter;
+        }
       }
     }
 
@@ -131,6 +146,8 @@ void CommTaskManager::CommTaskLoop() {
     if (comm_task_list_.empty() && init_comm_task_map_.empty() &&
         start_comm_task_map_.empty()) {
       done = true;
+    } else {
+      done = false;
     }
   }
 }

diff --git a/paddle/phi/core/distributed/comm_task_manager.h b/paddle/phi/core/distributed/comm_task_manager.h
@@ -46,6 +46,7 @@ class CommTaskManager {
   }
 
   void CommTaskEnqueue(std::shared_ptr<CommTask> comm_task);
+  void Stop();
 
  private:
   void CommTaskLoop();