[distributed] Remove recordStream for callback that ends a profiler event (#109933)

davidberard98 · pytorchmergebot · commit 4069d1de59bb · 2023-10-03T14:40:43.000Z
**Background**: recordStreams can result in memory spikes, so we don't want them to appear in FSDP (https://dev-discuss.pytorch.org/t/fsdp-cudacachingallocator-an-outsider-newb-perspective/1486). @ awgu is working on fixing this, but it turns out profiler was causing recordStream to get called when it is enabled. Why profiler was causing recordStream to get called: NCCL calls add profiler events manually; they register a callback to be executed when the future for the collective is completed; this indicates the end of the CPU-side profiler event for the callback: https://github.com/pytorch/pytorch/blob/c2c7c4035f41f90116aadf2df3f5d5b4834f838b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L1822-L1824 In order to guarantee safety, ivalue::Future::invokeCallback calls `recordStream` on the future's storage buffers; this marks the fact that other streams (e.g. the one that the callback runs on) may need to use the storage. https://github.com/pytorch/pytorch/blob/c2c7c4035f41f90116aadf2df3f5d5b4834f838b/aten/src/ATen/core/ivalue_inl.h#L1171-L1173 **Change**: The end-profiler-event callback doesn't actually use the future, so we don't need to recordStream on it. This PR introduces an optional parameter `uses_future` for adding callbacks; a user can set this variable to "false" to unsafely skip the recordStream, if the user knows that the future will not be used in the lambda. **Tests**: (a) unit tests; (b) added an assert in recordStream: https://github.com/pytorch/pytorch/blob/c2c7c4035f41f90116aadf2df3f5d5b4834f838b/c10/cuda/CUDACachingAllocator.cpp#L3260 and verified that it doesn't get triggered when running basic distributed tests w/ profiler enabled Pull Request resolved: #109933 Approved by: https://github.com/wconstab
diff --git a/aten/src/ATen/core/ivalue_inl.h b/aten/src/ATen/core/ivalue_inl.h
@@ -846,6 +846,15 @@ struct C10_EXPORT ivalue::Future final : c10::intrusive_ptr_target {
 
   friend c10::intrusive_ptr<Future>;
 
+  struct FutureCallback {
+    std::function<void(Future&)> callback;
+    bool uses_future; // whether the Future& passed in is actually used
+
+    template <typename T>
+    FutureCallback(T callback, bool uses_future)
+        : callback(std::move(callback)), uses_future(uses_future) {}
+  };
+
  public:
   Future(const Future&) = delete;
   Future(Future&&) = delete;
@@ -942,13 +951,13 @@ struct C10_EXPORT ivalue::Future final : c10::intrusive_ptr_target {
       events_.push_back(std::move(event));
     }
 
-    std::vector<std::function<void(Future&)>> cbs;
+    std::vector<FutureCallback> cbs;
     cbs.swap(callbacks_);
     lock.unlock();
 
     finished_cv_.notify_all();
     for (auto& callback : cbs) {
-      invokeCallback(std::move(callback));
+      invokeCallback(std::move(callback.callback), callback.uses_future);
     }
   }
 
@@ -1023,19 +1032,20 @@ struct C10_EXPORT ivalue::Future final : c10::intrusive_ptr_target {
    * this function will execute the callback immediately.
    */
   template <typename T>
-  void addCallback(T callback) {
+  void addCallback(T callback, bool uses_future = true) {
 #if __cpp_lib_is_invocable >= 201703
     static_assert(
         std::is_invocable_r<void, T, Future&>::value,
         "The callback must have signature void(Future&)");
 #endif
+
     std::unique_lock<std::mutex> lock(mutex_);
     if (completed()) {
       lock.unlock();
-      invokeCallback(std::move(callback));
+      invokeCallback(std::move(callback), uses_future);
       return;
     }
-    callbacks_.emplace_back(std::move(callback));
+    callbacks_.emplace_back(std::move(callback), uses_future);
   }
 
   /**
@@ -1153,24 +1163,29 @@ struct C10_EXPORT ivalue::Future final : c10::intrusive_ptr_target {
   // set up before running the callback, as in, it will set up the CUDA streams,
   // synchronize them with the value, and so on (if needed).
   template<typename T>
-  void invokeCallback(T callback) {
+  void invokeCallback(T callback, bool uses_future) {
 #if __cpp_lib_is_invocable >= 201703
     static_assert(
         std::is_invocable_r<void, T, Future&>::value,
         "The callback must have signature void(Future&)");
 #endif
 
-    c10::OptionalDeviceGuard deviceGuard(currentDevice_);
+    // The synchronization performed below shouldn't be needed when the future
+    // is not used by the callback.
+    if (uses_future) {
+      c10::OptionalDeviceGuard deviceGuard(currentDevice_);
 
-    std::vector<c10::Stream> streams;
-    streams.reserve(devices_.size());
-    for (const c10::Device& device : devices_) {
-      streams.push_back(impl_.getStreamFromGlobalPool(device));
+      std::vector<c10::Stream> streams;
+      streams.reserve(devices_.size());
+      for (const c10::Device& device : devices_) {
+        streams.push_back(impl_.getStreamFromGlobalPool(device));
+      }
+      c10::MultiStreamGuard streamGuard(streams);
+      synchronizeWithCurrentStreams();
+      callback(*this);
+    } else {
+      callback(*this);
     }
-    c10::MultiStreamGuard streamGuard(streams);
-    synchronizeWithCurrentStreams();
-
-    callback(*this);
   }
 
   // This method should be called before this future's value is used, as it
@@ -1206,13 +1221,13 @@ struct C10_EXPORT ivalue::Future final : c10::intrusive_ptr_target {
     completed_ = true;
     eptr_ = std::move(eptr);
 
-    std::vector<std::function<void(Future&)>> cbs;
+    std::vector<FutureCallback> cbs;
     cbs.swap(callbacks_);
     lock.unlock();
 
     finished_cv_.notify_all();
     for (auto& callback : cbs) {
-      invokeCallback(std::move(callback));
+      invokeCallback(std::move(callback.callback), callback.uses_future);
     }
   }
 
@@ -1353,7 +1368,7 @@ struct C10_EXPORT ivalue::Future final : c10::intrusive_ptr_target {
 
   IValue value_; // when finished the value
   TypePtr type_;
-  std::vector<std::function<void(Future&)>> callbacks_;
+  std::vector<FutureCallback> callbacks_;
   std::exception_ptr eptr_;
 
   // An upcast pointer to a virtual class which allows us to manipulate events,
diff --git a/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
@@ -1819,9 +1819,14 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::collective(
     // future blocks the stream this callback runs on the corresponding
     // ncclEndEvents_ ensuring appropriate synchronization.
     if (work->recordFunctionEndCallback_) {
-      work->future_->addCallback([work](at::ivalue::Future& /* unused */) {
-        work->recordFunctionEndCallback_();
-      });
+      work->future_->addCallback(
+          [work](at::ivalue::Future& /* unused */) {
+            work->recordFunctionEndCallback_();
+          },
+          // uses_future = false allows us to skip synchronization in
+          // ivalue::Future, but is only valid as long as the lambda doesn't use
+          // the "Future" argument.
+          /*uses_future=*/false);
     }
     work->future_->markCompleted(at::IValue(*work->outputs_));
   }
@@ -1999,9 +2004,14 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::pointToPoint(
   // future blocks the stream this callback runs on the corresponding
   // ncclEndEvents_ ensuring appropriate synchronization.
   if (work->recordFunctionEndCallback_) {
-    work->future_->addCallback([work](at::ivalue::Future& /* unused */) {
-      work->recordFunctionEndCallback_();
-    });
+    work->future_->addCallback(
+        [work](at::ivalue::Future& /* unused */) {
+          work->recordFunctionEndCallback_();
+        },
+        // uses_future = false allows us to skip synchronization in
+        // ivalue::Future, but is only valid as long as the lambda doesn't use
+        // the "Future" argument.
+        /*uses_future=*/false);
   }
 
   return work;
