[NCCL][CUDA][CUDA Graphs] Flush enqueued work before starting a graph capture (#104487)

eqy · pytorchmergebot · commit db63bf3d7e5e · 2023-09-01T05:42:08.000Z
#103503 addresses the situation where additional work is enqueued for the NCCL watchdog to poll during a graph capture---something we want to avoid as the subsequent polling will query an event and crash the capture. However, there is currently no check that there was not work _already_ enqueued for the watchdog to poll. If there was already work that was enqueued and not cleaned up before the start of a graph capture, then we run into a similar problem where the event query will crash the capture. We've observed this causing crashes on several workloads, although it is somewhat flaky (if the watchdog happens to poll just before the graph capture and cleanup, then we dodge the crash). This is a bit of a tricky issue as it involves making sure that no process group has enqueued work at the start of a capture, and as such the simplest solution is to add a bit of global state to track all process groups. This PR forces the start of the graph capture to wait until all enqueued work is completed and cleaned up or times out. I did consider the alternative of simply having the watchdog skip cleanup if we detect that we are in the middle of a graph capture, but I think deferring the cleanup until later could result in false positive timeouts (e.g., if we defer cleanup on work that has completed long ago, checking timers after the graph capture could yield a "timeout"). CC @Aidyn-A @bottler @kwen2501 @ptrblck Pull Request resolved: #104487 Approved by: https://github.com/kwen2501
diff --git a/aten/src/ATen/cuda/CUDAGraph.cpp b/aten/src/ATen/cuda/CUDAGraph.cpp
@@ -4,6 +4,7 @@
 #include <ATen/Functions.h>
 #include <c10/cuda/CUDACachingAllocator.h>
 #include <c10/cuda/CUDAFunctions.h>
+#include <torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp>
 
 namespace at::cuda {
 
@@ -115,6 +116,12 @@ void CUDAGraph::capture_begin(MempoolId_t pool/*=0*/, cudaStreamCaptureMode capt
   // due to the capture status being updated _after_ a capture had already started.
   c10::cuda::CUDACachingAllocator::beginAllocateStreamToPool(capture_dev_, capture_stream_, mempool_id_);
 
+#ifdef USE_C10D_NCCL
+  // If the watchdog has remaining work enqueued, an event query on the remaining work will crash
+  // the graph capture, so we wait for all pending work to be completed.
+  c10d::ProcessGroupNCCL::waitForAllPendingWorks();
+#endif
+
   // cudaStreamCaptureModeGlobal is the most conservative option to
   // prevent potentially unsafe CUDA API calls during capture.  See
   // https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html#group__CUDART__STREAM_1g9d0535d93a214cbf126835257b16ba85
diff --git a/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
@@ -286,6 +286,9 @@ inline void errorIfCapturingNonCapturableNCCL(c10::cuda::CaptureStatus status) {
 const int64_t ProcessGroupNCCL::kWatchdogThreadSleepMillis = 1000;
 constexpr int64_t kSynchronizeBusyWaitMillis = 10;
 thread_local uint64_t ProcessGroupNCCL::ncclActiveGroupCounter_ = 0;
+std::mutex ProcessGroupNCCL::allProcessGroupsMutex_;
+std::unordered_set<c10d::ProcessGroupNCCL*>
+    ProcessGroupNCCL::all_nccl_process_groups;
 
 std::ostream& operator<<(
     std::ostream& output,
@@ -728,6 +731,10 @@ ProcessGroupNCCL::ProcessGroupNCCL(
     }
   }
 #endif
+  {
+    std::lock_guard<std::mutex> lk(allProcessGroupsMutex_);
+    all_nccl_process_groups.insert(this);
+  }
 }
 
 void ProcessGroupNCCL::runHealthCheck() {
@@ -897,6 +904,11 @@ ProcessGroupNCCL::~ProcessGroupNCCL() {
   // Abort all NCCL Communicators on Process Group Destruction
   std::string abortReason = c10::str("Process Group destroyed on rank ", rank_);
   abort(abortReason);
+
+  {
+    std::lock_guard<std::mutex> lk(allProcessGroupsMutex_);
+    all_nccl_process_groups.erase(this);
+  }
 }
 
 void ProcessGroupNCCL::ncclCommWatchdog() {
@@ -1089,6 +1101,15 @@ void ProcessGroupNCCL::runHookLoop() {
   }
 }
 
+void ProcessGroupNCCL::waitForAllPendingWorks() {
+  std::lock_guard<std::mutex> lk(allProcessGroupsMutex_);
+  for (auto it = ProcessGroupNCCL::all_nccl_process_groups.begin();
+       it != ProcessGroupNCCL::all_nccl_process_groups.end();
+       it++) {
+    (*it)->waitForPendingWorks();
+  }
+}
+
 std::exception_ptr ProcessGroupNCCL::WorkNCCL::checkForNCCLErrors(
     const std::vector<std::shared_ptr<NCCLComm>>& ncclComms) const {
   return checkForNCCLErrorsInternal(ncclComms);
@@ -1649,6 +1670,13 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::endCoalescing() {
   c10::cuda::CaptureStatus capture_status =
       c10::cuda::currentStreamCaptureStatusMayInitCtx();
 
+  if (capture_status != c10::cuda::CaptureStatus::None) {
+    std::lock_guard<std::mutex> lock(workMetaListMutex_);
+    TORCH_INTERNAL_ASSERT(
+        workMetaList_.empty(),
+        "In the middle of a CUDA Graph capture but the enqueued work is not empty. The watchdog will crash the capture when it polls the work.");
+  }
+
   if ((coalescing_state_ & CoalColl) &&
       capture_status == c10::cuda::CaptureStatus::None) {
     workEnqueue(work);
@@ -1823,6 +1851,13 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::collective(
   work->numelIn_ = inputs[0].numel();
   work->numelOut_ = outputs[0].numel();
 
+  if (capture_status != c10::cuda::CaptureStatus::None) {
+    std::lock_guard<std::mutex> lock(workMetaListMutex_);
+    TORCH_INTERNAL_ASSERT(
+        workMetaList_.empty(),
+        "In the middle of a CUDA Graph capture but the enqueued work is not empty. The watchdog will crash the capture when it polls the work.");
+  }
+
   if (!coalescing_state_ && capture_status == c10::cuda::CaptureStatus::None) {
     workEnqueue(work);
   }
diff --git a/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp b/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp
@@ -8,6 +8,7 @@
 #include <mutex>
 #include <thread>
 #include <unordered_map>
+#include <unordered_set>
 
 #include <torch/csrc/distributed/c10d/Backend.hpp>
 #include <torch/csrc/distributed/c10d/NCCLUtils.hpp>
@@ -345,6 +346,11 @@ class TORCH_API ProcessGroupNCCL : public Backend {
 
   ~ProcessGroupNCCL() override;
 
+  // Check that all work is done (no enqueued work).
+  // We use this to avoid uwittingly having watchdogs query work during
+  // CUDA graph captures.
+  static void waitForAllPendingWorks();
+
   c10::intrusive_ptr<Options> getOptions() {
     return options_;
   }
@@ -689,6 +695,9 @@ class TORCH_API ProcessGroupNCCL : public Backend {
   // Mutex to Guard workMetaList_
   std::mutex workMetaListMutex_;
 
+  // Mutex to Guard all_nccl_process_groups
+  static std::mutex allProcessGroupsMutex_;
+
   // Condition Variable for watchdog thread sleep
   std::condition_variable workMetaListCV_;
 
@@ -703,6 +712,9 @@ class TORCH_API ProcessGroupNCCL : public Backend {
 
   std::list<ProcessGroupNCCL::WorkNCCL> completedWorkList_;
 
+  // All process groups for checking Watchdog status
+  static std::unordered_set<c10d::ProcessGroupNCCL*> all_nccl_process_groups;
+
   // Add Work Pointer to workVector
   void workEnqueue(c10::intrusive_ptr<ProcessGroupNCCL::WorkNCCL>);
 
