diff --git a/apex/contrib/csrc/nccl_allocator/NCCLAllocator.cpp b/apex/contrib/csrc/nccl_allocator/NCCLAllocator.cpp index 1c4a70528..14cdaff86 100644 --- a/apex/contrib/csrc/nccl_allocator/NCCLAllocator.cpp +++ b/apex/contrib/csrc/nccl_allocator/NCCLAllocator.cpp @@ -324,6 +324,14 @@ c10::cuda::CUDACachingAllocator::SnapshotInfo NCCLAllocator::snapshot() { return result; } +c10::cuda::CUDACachingAllocator::ShareableHandle NCCLAllocator:: + shareIpcHandle(void* ptr) { + TORCH_CHECK( + false, + "CUDAPluggableAllocator does not yet support shareIPcHandle. " + "If you need it, please file an issue describing your use case."); +} + std::shared_ptr NCCLAllocator::getIpcDevPtr(std::string handle) { TORCH_CHECK( false, @@ -432,4 +440,10 @@ void NCCLAllocator::copy_data(void* dest, const void* src, std::size_t count) cudaMemcpy(dest, src, count, cudaMemcpyKind::cudaMemcpyDeviceToDevice)); } +void NCCLAllocator::enable(bool) {} + +bool NCCLAllocator::isEnabled() const { + return true; +} + } // namespace nccl_allocator::cuda diff --git a/apex/contrib/csrc/nccl_allocator/NCCLAllocator.h b/apex/contrib/csrc/nccl_allocator/NCCLAllocator.h index 3a4e76b83..5a71325d8 100644 --- a/apex/contrib/csrc/nccl_allocator/NCCLAllocator.h +++ b/apex/contrib/csrc/nccl_allocator/NCCLAllocator.h @@ -59,6 +59,8 @@ struct NCCLAllocator override; void releasePool(c10::DeviceIndex device, c10::cuda::MempoolId_t mempool_id) override; std::shared_ptr getIpcDevPtr(std::string handle) override; + c10::cuda::CUDACachingAllocator::ShareableHandle shareIpcHandle( + void*) override; void recordHistory( bool enabled, c10::cuda::CUDACachingAllocator::CreateContextFn context_recorder, @@ -85,6 +87,8 @@ struct NCCLAllocator bool p2p_enabled) override; std::string name() override; void copy_data(void* dest, const void* src, std::size_t count) const; + void enable(bool) override; + bool isEnabled() const override; protected: std::function init_fn_;