Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Enable TensorPipe's CUDA GDR channel #50763

Closed
wants to merge 3 commits into from
Closed

Enable TensorPipe's CUDA GDR channel #50763

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4621419
Enable TensorPipe's CUDA GDR channel
lw Jan 19, 2021
c494a22
Update on "Enable TensorPipe's CUDA GDR channel"
lw Jan 25, 2021
22a91c3
Update on "Enable TensorPipe's CUDA GDR channel"
lw Jan 28, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
28 changes: 28 additions & 0 deletions torch/csrc/distributed/rpc/tensorpipe_agent.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,10 @@ constexpr int64_t kBasicChannelPriority = 0;
constexpr int64_t kCudaIpcChannelPriority = 300;
#endif

#if TENSORPIPE_HAS_CUDA_GDR_CHANNEL && defined(USE_CUDA_NOT_ROCM)
constexpr int64_t kCudaGdrChannelPriority = 200;
#endif

#ifdef USE_CUDA_NOT_ROCM
constexpr int64_t kCudaXthChannelPriority = 400;
constexpr int64_t kCudaBasicChannelPriority = 100;
Expand Down Expand Up @@ -250,6 +254,30 @@ C10_REGISTER_CREATOR(

#endif

#if TENSORPIPE_HAS_CUDA_GDR_CHANNEL && defined(USE_CUDA_NOT_ROCM)

std::unique_ptr<CudaChannelRegistration> makeCudaGdrChannel() {
auto context = std::make_shared<tensorpipe::channel::cuda_gdr::Context>();
return std::make_unique<CudaChannelRegistration>(
CudaChannelRegistration{std::move(context), kCudaGdrChannelPriority});
}

// The cuda_gdr channel sends CUDA memory over InfiniBand using GPUDirect RDMA.
// It directly registers the user-provided tensor with libibverbs, an operation
// which is expensive the first time, but it then caches the registration in
// order to amortize the cost and get low latency for subsequent transfers. A
// ready-to-send/ready-to-receive handshake is still needed before the transfer
// in order to ensure readiness and to agree on the device indices and thus the
// queue pair to use. It automatically pairs each GPU to the "closest" NIC if
// there are multiple of them (closest = longest prefix match in PCI tree).
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
C10_REGISTER_CREATOR(
TensorPipeCudaChannelRegistry,
cuda_gdr,
makeCudaGdrChannel);

#endif

#ifdef USE_CUDA_NOT_ROCM

std::unique_ptr<CudaChannelRegistration> makeCudaXthChannel() {
Expand Down