Split out reusable CUDAFuture from FutureNCCL

[lw]

Stack from ghstack:

• Add support for async callbacks in ivalue::Future #48790 Add support for async callbacks in ivalue::Future
• Drop FutureNCCL in favor of vanilla CUDAFuture #49014 Drop FutureNCCL in favor of vanilla CUDAFuture
• Make CUDAFuture remember and restore current device in callback #48789 Make CUDAFuture remember and restore current device in callback
• Remove DataPtr extractor from CUDAFuture #48840 Remove DataPtr extractor from CUDAFuture
• Cache the DataPtrs in CUDAFuture #48788 Cache the DataPtrs in CUDAFuture
• Split out reusable CUDAFuture from FutureNCCL #48506 Split out reusable CUDAFuture from FutureNCCL
• Merge common parts of FutureNCCL into at::ivalue::Future #48505 Merge common parts of FutureNCCL into at::ivalue::Future
• Split FutureNCCL's CUDA-specific parts from generic future logic #48504 Split FutureNCCL's CUDA-specific parts from generic future logic
• Support wider range of types in FutureNCCL #48502 Support wider range of types in FutureNCCL
• Don't store device indices separately on FutureNCCL #48501 Don't store device indices separately on FutureNCCL
• Add multi-GPU support to FutureNCCL #48500 Add multi-GPU support to FutureNCCL
• Fix FutureNCCL not recording dataptrs with caching alloc in wait() #48563 Fix FutureNCCL not recording dataptrs with caching alloc in wait()
• Fix FutureNCCL's completed() disagreeing with wait() #48503 Fix FutureNCCL's completed() disagreeing with wait()
• Record CUDA events for "follow-up" FutureNCCL inside markCompleted #48499 Record CUDA events for "follow-up" FutureNCCL inside markCompleted
• Use fresh stream from pool for each FutureNCCL callback #48498 Use fresh stream from pool for each FutureNCCL callback
• Make FutureNCCL record events in current stream #48497 Make FutureNCCL record events in current stream
• Have FutureNCCL record streams w/ allocator in addCallback #48496 Have FutureNCCL record streams w/ allocator in addCallback
• Add some safeguards to FutureNCCL #48562 Add some safeguards to FutureNCCL
• Remove NCCL dependency from PythonFutureWrapper #48495 Remove NCCL dependency from PythonFutureWrapper
• Avoid using FutureNCCL before it's ready #48561 Avoid using FutureNCCL before it's ready

This commit is part of a stack that reworks FutureNCCL in order to extract a generic CUDA-aware Future subclass. The stack deliberately breaks up this transition into elementary changes, to make it easier to verify that the behavior is preserved (or to highlight how it gets changed).

---

FutureNCCL is now a general-purpose type-agnostic multi-device class, so in this commit I extract it from ProcessGroupNCCL to make it available for wider use (notably by the RPC module). We'll call this new class CUDAFuture. We'll keep FutureNCCL as a subclass of CUDAFuture to deal with some NCCL peculiarity, namely the fact that the future becomes complete immediately upon creation. We can clean this up for good once we're done merging Future and Work.

I'm not exactly sure of where to put CUDAFuture. It needs to be available to both c10d and RPC (which lives under torch/csrc). If I figured CMake out correctly (and that's a big if) I think c10d can only depend on ATen (I'll maybe add a comment with how I tracked that down). Hence we cannot put CUDAFuture in torch/csrc. On the other hand, RPC currently depends on c10d, because RPC agents use ProcessGroups internally, so it would be "ok" to put CUDAFuture in c10d. However, we want to get rid of ProcessGroups in RPC, and at that point RPC should in principle not depend on c10d. In that case, the only shared dep between the two that I see is ATen itself.

While I'm a bit wary of putting it right in ATen, I think it might actually make sense. CUDAFuture is intended to be a general-purpose component that can be reused in all settings and is not particularly tied to c10d or RPC. Moreover, ATen already contains ivalue::Future, and it contains a lot of CUDA helpers, so CUDAFuture definitely belongs to the "closure" of what's already there.

Differential Revision: D25180532

[lw]

Here we go:

• c10d links agains torch
  

  pytorch/torch/lib/c10d/CMakeLists.txt

  Line 58 in c5ce995

  set(C10D_LIBS torch)

  
  

  pytorch/torch/lib/c10d/CMakeLists.txt

  Line 147 in c5ce995

  target_link_libraries(c10d PUBLIC ${C10D_LIBS})

• torch in turn links against torch_cuda_library (I'm ignoring CPU-only deps as we cannot put CUDAFuture there)
  

  pytorch/caffe2/CMakeLists.txt

  Line 1167 in c5ce995

  target_link_libraries(torch PUBLIC torch_cuda_library)

• which itself is an interface library for torch_cuda
  

  pytorch/caffe2/CMakeLists.txt

  Line 1150 in c5ce995

  caffe2_interface_library(torch_cuda torch_cuda_library)

• which in turn pulls in a lot of stuff (through sources and dependencies)
  

  pytorch/caffe2/CMakeLists.txt

  Line 763 in c5ce995

  cuda_add_library(torch_cuda ${Caffe2_GPU_SRCS} ${Caffe2_GPU_SRCS_W_SORT_BY_KEY})

  
  

  pytorch/caffe2/CMakeLists.txt

  Line 1186 in c5ce995

  target_link_libraries(torch_cuda PUBLIC c10_cuda torch::nvtoolsext)

  
  

  pytorch/caffe2/CMakeLists.txt

  Lines 1190 to 1193 in c5ce995

  	target_include_directories(
  	torch_cuda PRIVATE ${Caffe2_GPU_INCLUDE})
  	target_link_libraries(
  	torch_cuda PRIVATE ${Caffe2_CUDA_DEPENDENCY_LIBS})

  
  

  pytorch/caffe2/CMakeLists.txt

  Line 1198 in c5ce995

  target_link_libraries(torch_cuda PUBLIC torch_cpu_library ${Caffe2_PUBLIC_CUDA_DEPENDENCY_LIBS})

• which however either resolves to ATen
  

  pytorch/caffe2/CMakeLists.txt

  Lines 78 to 79 in c5ce995

  	list(APPEND Caffe2_GPU_SRCS ${ATen_CUDA_SRCS})
  	list(APPEND Caffe2_GPU_SRCS_W_SORT_BY_KEY ${ATen_CUDA_SRCS_W_SORT_BY_KEY})

  
  

  pytorch/caffe2/CMakeLists.txt

  Line 88 in c5ce995

  list(APPEND Caffe2_GPU_INCLUDE ${ATen_CUDA_INCLUDE})

  
  

  pytorch/caffe2/CMakeLists.txt

  Line 92 in c5ce995

  list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS ${ATen_CUDA_DEPENDENCY_LIBS})

• or to something inside Caffe2, which I didn't chase

• or some more fundamental (external) dependencies

[dr-ci]

💊 CI failures summary and remediations

As of commit d432680 (more details on the Dr. CI page):

---

• 1/1 failures introduced in this PR

---

🕵️ 1 new failure recognized by patterns

The following CI failures do not appear to be due to upstream breakages:

pytorch_linux_xenial_py3_6_gcc5_4_build (1/1)

Step: "Build" (full log | diagnosis details | 🔁 rerun)

Dec 09 15:30:37 sccache: error: couldn't connect to server

Dec 09 15:30:37 +++ eval 'extract_trap_cmd ' 
Dec 09 15:30:37 ++++ extract_trap_cmd 
Dec 09 15:30:37 ++++ printf '%s\n' '' 
Dec 09 15:30:37 +++ printf '%s\n' cleanup 
Dec 09 15:30:37 ++ trap -- ' 
Dec 09 15:30:37 cleanup' EXIT 
Dec 09 15:30:37 ++ [[ pytorch-linux-xenial-py3.6-gcc5.4-build != *pytorch-win-* ]] 
Dec 09 15:30:37 ++ which sccache 
Dec 09 15:30:37 ++ sccache --stop-server 
Dec 09 15:30:37 Stopping sccache server... 
Dec 09 15:30:37 sccache: error: couldn't connect to server 
Dec 09 15:30:37 sccache: caused by: Connection refused (os error 111) 
Dec 09 15:30:37 ++ true 
Dec 09 15:30:37 ++ rm /var/lib/jenkins/sccache_error.log 
Dec 09 15:30:37 rm: cannot remove '/var/lib/jenkins/sccache_error.log': No such file or directory 
Dec 09 15:30:37 ++ true 
Dec 09 15:30:37 ++ [[ pytorch-linux-xenial-py3.6-gcc5.4-build == *rocm* ]] 
Dec 09 15:30:37 ++ SCCACHE_ERROR_LOG=/var/lib/jenkins/sccache_error.log 
Dec 09 15:30:37 ++ SCCACHE_IDLE_TIMEOUT=1200 
Dec 09 15:30:37 ++ RUST_LOG=sccache::server=error 
Dec 09 15:30:37 ++ sccache --start-server 

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 27 times.

[mrshenli]

hmm, this means users cannot call wait on CUDAFuture before it is marked as completed?

[lw]

Not sure I follow: the postWaitHook is called by ivalue::Future once a wait has finished, and that only happens when the future is marked complete. Thus one can certainly start waiting before the future is complete, in which case one will block inside ivalue::Future (waiting on a condition variable). Since that only unblocks after the future is complete, we're guaranteed that postWaitHook will be called after postMarkCompletedHook. Is this what you were wondering about?

[mrshenli]

Are these events used anywhere? Will the markCompleted call below this line immediately override these cudaEvents_?

[lw]

These events will not be overridden when the future is marked complete, because FutureNCCL provides its own postMarkCompletedHook which does nothing (and which overrides the one from CUDAFuture, which is where we store the events).

However, these events will still be used by the postWaitHook, by the wrapCallback, etc.

It's somewhat hacky, but it works, and it's a temporary solution only until we finish merging Future and Work (at which point ProcessGroupNCCL should be able to use CUDAFuture unmodified).

[wanchaol]

looks good, left some minor questions.

[wanchaol]

if this is the only way we want to create an instance, can we delete the default constructor? also, where's the PR that have the createInstance for ivalue::Future?

[lw]

This is the only way I found to have at::ivalue::Future::then() be able to create an instance of at::cuda::CUDAFuture. In short, the a method of the superclass needs to create an instance of the subclass. It's not supposed to be used by external users, hence it's protected.

The superclass implementation of createInstance is here: https://github.com/pytorch/pytorch/pull/48505/files#diff-2f833078d12338d0ac920ab654b5c791cb1219729d6e7c97a08393b32a46d173R491

[wanchaol]

what's the difference of this between capture by reference? are the callbacks gonna destructed sometime but we still want them alive in the lambda?

[lw]

std::function AFAIK is a RAII wrapper so it might hold some state which gets destroyed once the function goes out of scope. If we captured a reference to the std::function that would thus risk being a dangling reference.

[facebook-github-bot]

This pull request has been merged in 030fa6c.