diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index 09f590c9bc83c..9550c7f75c253 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -953,6 +953,12 @@ if(USE_ROCM) target_compile_definitions(torch_hip PRIVATE USE_NCCL) endif() + # IntraNodeComm loads libamd_smi at runtime via dlopen (avoids linking amd_smi + # which would cause bus errors if Python amdsmi loads a different copy). + if(USE_DISTRIBUTED) + target_link_libraries(torch_hip PRIVATE ${CMAKE_DL_LIBS}) + endif() + if(USE_PRECOMPILED_HEADERS) target_precompile_headers(torch_hip PRIVATE "$<$:ATen/core/ATen_pch.h>") diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index a7a2840ef0649..93fd2e5674db5 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -1090,13 +1090,6 @@ if(USE_ROCM) ) endif() - # ROCM-SMI needed to support symmetric memory - if(USE_DISTRIBUTED AND UNIX) - list(APPEND Caffe2_PUBLIC_HIP_DEPENDENCY_LIBS - rocm_smi64 - ) - endif() - # ---[ Kernel asserts # Kernel asserts is disabled for ROCm by default. # It can be turned on by turning on the env USE_ROCM_KERNEL_ASSERT to the build system. diff --git a/cmake/public/LoadHIP.cmake b/cmake/public/LoadHIP.cmake index fe6fd4cdfbb09..7ecaff5109f42 100644 --- a/cmake/public/LoadHIP.cmake +++ b/cmake/public/LoadHIP.cmake @@ -194,7 +194,6 @@ endif() if(UNIX) find_package_and_print_version(rccl) find_package_and_print_version(hsa-runtime64 REQUIRED) - find_package_and_print_version(rocm_smi REQUIRED) endif() # Optional components. diff --git a/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp index c4bf7ffacfe66..50a36bb77c5e7 100644 --- a/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp +++ b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp @@ -891,7 +891,7 @@ ProcessGroupNCCL::ProcessGroupNCCL( options_(std::move(options)), terminateProcessGroup_(false), local_id_(process_group_id++), - intraNodeComm_(initIntraNodeComm()) { + intraNodeComm_(nullptr) { TORCH_CHECK_WITH( ValueError, at::cuda::getNumGPUs() != 0, @@ -1183,7 +1183,10 @@ c10::intrusive_ptr ProcessGroupNCCL:: return nullptr; } auto prefixStore = c10::make_intrusive("IntraNodeComm", store_); - auto comm = c10::make_intrusive(prefixStore, rank_, size_); + const std::string groupName = + options_->group_name.empty() ? "0" : options_->group_name; + auto comm = c10::make_intrusive( + prefixStore, rank_, size_, std::nullopt, groupName); if (comm->rendezvous()) { return comm; } else { @@ -4468,12 +4471,17 @@ c10::intrusive_ptr ProcessGroupNCCL::allreduce( } check_gpu_single_tensor(tensor); - if (intraNodeComm_ != nullptr && opts.reduceOp == ReduceOp::SUM) { + if (opts.reduceOp == ReduceOp::SUM) { using namespace intra_node_comm; - auto algo = intraNodeComm_->selectAllReduceAlgo(tensor); - if (algo != intra_node_comm::AllReduceAlgo::NONE) { - intraNodeComm_->allReduce(tensor, algo); - return c10::make_intrusive(); + if (intraNodeComm_ == nullptr && IntraNodeComm::isEnabled()) { + intraNodeComm_ = initIntraNodeComm(); + } + if (intraNodeComm_ != nullptr) { + auto algo = intraNodeComm_->selectAllReduceAlgo(tensor); + if (algo != intra_node_comm::AllReduceAlgo::NONE) { + intraNodeComm_->allReduce(tensor, algo); + return c10::make_intrusive(); + } } } TORCH_CHECK( diff --git a/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cpp b/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cpp index f62577e701847..a84c0dcebebfe 100644 --- a/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cpp +++ b/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cpp @@ -3,7 +3,10 @@ #include #if defined(USE_ROCM) -#include +#include +#include +#include +#include #endif namespace c10d::intra_node_comm { @@ -20,7 +23,7 @@ static int intraNodeCommIdx = 0; * Query the nvlink connection among devices. */ static NvlMesh getNvlMesh(const std::vector& rankToDeviceIdx) { -#if !defined(USE_RCOM) +#if !defined(USE_ROCM) auto connectivity = detect_dma_connectivity(c10::DeviceType::CUDA, "nvlink"); NvlMesh nvlMesh = {}; for (size_t srcRank = 0; srcRank < kMaxDevices; ++srcRank) { @@ -35,23 +38,126 @@ static NvlMesh getNvlMesh(const std::vector& rankToDeviceIdx) { } return nvlMesh; #else + // Load libamd_smi at runtime to avoid linking it into torch_hip (double-load + // with Python amdsmi causes bus errors). Types/constants from amdsmi.h only. + struct AmdsmiApi { + amdsmi_status_t (*init)(uint64_t); + amdsmi_status_t (*get_socket_handles)(uint32_t*, amdsmi_socket_handle*); + amdsmi_status_t (*get_processor_handles)( + amdsmi_socket_handle, + uint32_t*, + amdsmi_processor_handle*); + amdsmi_status_t (*is_P2P_accessible)( + amdsmi_processor_handle, + amdsmi_processor_handle, + bool*); + }; + static void* amdsmi_handle = nullptr; + static AmdsmiApi amdsmi = {}; + static bool amdsmi_resolved = false; + + if (!amdsmi_resolved) { + amdsmi_resolved = true; + const char* rocm = std::getenv("ROCM_PATH"); + std::string path = + rocm ? std::string(rocm) + "/lib/libamd_smi.so" : "libamd_smi.so"; + amdsmi_handle = dlopen(path.c_str(), RTLD_NOW | RTLD_LOCAL); + if (!amdsmi_handle) { + amdsmi_handle = dlopen("libamd_smi.so", RTLD_NOW | RTLD_LOCAL); + } + if (!amdsmi_handle) { + LOG(ERROR) << "IntraNodeComm:: getNvlMesh: dlopen libamd_smi.so failed: " + << dlerror(); + return {}; + } + amdsmi.init = reinterpret_cast( + dlsym(amdsmi_handle, "amdsmi_init")); + amdsmi.get_socket_handles = + reinterpret_cast( + dlsym(amdsmi_handle, "amdsmi_get_socket_handles")); + amdsmi.get_processor_handles = + reinterpret_cast( + dlsym(amdsmi_handle, "amdsmi_get_processor_handles")); + amdsmi.is_P2P_accessible = + reinterpret_cast( + dlsym(amdsmi_handle, "amdsmi_is_P2P_accessible")); + if (!amdsmi.init || !amdsmi.get_socket_handles || + !amdsmi.get_processor_handles || !amdsmi.is_P2P_accessible) { + LOG(ERROR) << "IntraNodeComm:: getNvlMesh: dlsym amdsmi failed"; + return {}; + } + } + NvlMesh nvlMesh = {}; const auto worldSize = rankToDeviceIdx.size(); - // For each device, loop over devices connected to it + + uint32_t socket_count = 0; + amdsmi_status_t ret = amdsmi.get_socket_handles(&socket_count, nullptr); + if (ret == AMDSMI_STATUS_NOT_INIT) { + ret = amdsmi.init(AMDSMI_INIT_AMD_GPUS); + if (ret != AMDSMI_STATUS_SUCCESS) { + LOG(ERROR) << "IntraNodeComm:: getNvlMesh: amdsmi_init failed, ret=" + << static_cast(ret); + return {}; + } + socket_count = 0; + ret = amdsmi.get_socket_handles(&socket_count, nullptr); + } + if (ret != AMDSMI_STATUS_SUCCESS) { + LOG(ERROR) + << "IntraNodeComm:: getNvlMesh: amdsmi_get_socket_handles failed, ret=" + << static_cast(ret); + return {}; + } + + std::vector socket_handles(socket_count); + ret = amdsmi.get_socket_handles(&socket_count, &socket_handles[0]); + if (ret != AMDSMI_STATUS_SUCCESS) { + LOG(ERROR) + << "IntraNodeComm:: getNvlMesh: amdsmi_get_socket_handles (buffer) failed, ret=" + << static_cast(ret); + return {}; + } + + std::vector processor_handles; + for (size_t i = 0; i < socket_count; ++i) { + uint32_t device_count = 0; + ret = + amdsmi.get_processor_handles(socket_handles[i], &device_count, nullptr); + if (ret != AMDSMI_STATUS_SUCCESS) { + LOG(ERROR) + << "IntraNodeComm:: getNvlMesh: amdsmi_get_processor_handles (count) failed, ret=" + << static_cast(ret); + return {}; + } + std::vector _processor_handles(device_count); + ret = amdsmi.get_processor_handles( + socket_handles[i], &device_count, &_processor_handles[0]); + if (ret != AMDSMI_STATUS_SUCCESS) { + LOG(ERROR) + << "IntraNodeComm:: getNvlMesh: amdsmi_get_processor_handles (buffer) failed, ret=" + << static_cast(ret); + return {}; + } + processor_handles.insert( + processor_handles.end(), + _processor_handles.begin(), + _processor_handles.end()); + } + for (size_t idx = 0; idx < worldSize; ++idx) { for (size_t link = 0; link < kMaxDevices; ++link) { if (idx == link) continue; - bool conn = false; - auto ret = rsmi_is_P2P_accessible(idx, link, &conn); - if (ret != RSMI_STATUS_SUCCESS) { + ret = amdsmi.is_P2P_accessible( + processor_handles[idx], processor_handles[link], &conn); + if (ret != AMDSMI_STATUS_SUCCESS) { LOG(ERROR) - << "IntraNodeComm: getNvlMesh: rsmi_is_P2P_accessible returned error ret=" - << ret; + << "IntraNodeComm: getNvlMesh: amdsmi_is_P2P_accessible failed, ret=" + << static_cast(ret); return {}; } - if (conn) { nvlMesh[idx][link] += 1; } @@ -88,11 +194,13 @@ IntraNodeComm::IntraNodeComm( c10::intrusive_ptr store, size_t rank, size_t worldSize, - std::optional bufferSize) + std::optional bufferSize, + std::string groupName) : store_(std::move(store)), rank_(rank), worldSize_(worldSize), - bufferSize_(bufferSize.has_value() ? *bufferSize : kDefaultBufferSize) {} + bufferSize_(bufferSize.has_value() ? *bufferSize : kDefaultBufferSize), + groupName_(std::move(groupName)) {} IntraNodeComm::~IntraNodeComm() { if (!isInitialized_) { @@ -171,14 +279,6 @@ bool IntraNodeComm::rendezvous() { gethostname(devInfo.hostname, sizeof(devInfo.hostname)); devInfo.deviceIdx = deviceIdx_; -#if defined(USE_ROCM) - auto ret = rsmi_init(0); - if (ret != RSMI_STATUS_SUCCESS) { - LOG(ERROR) << "IntraNodeComm:: rendezvous failed in rsmi_init, ret=" << ret; - return false; - } -#endif - auto peerDevInfos = storeAllGather(store_, "handshake-0", rank_, worldSize_, devInfo); @@ -214,11 +314,13 @@ bool IntraNodeComm::rendezvous() { return false; } - auto groupName = "IntraNodeComm" + std::to_string(intraNodeCommIdx++); + const std::string name = groupName_.empty() + ? "IntraNodeComm" + std::to_string(intraNodeCommIdx++) + : groupName_; set_group_info( - groupName, static_cast(rank_), static_cast(worldSize_), store_); + name, static_cast(rank_), static_cast(worldSize_), store_); auto allocator = get_allocator(c10::DeviceType::CUDA); - symmetricMemoryPtr_ = allocator->alloc(bufferSize_, deviceIdx_, groupName); + symmetricMemoryPtr_ = allocator->alloc(bufferSize_, deviceIdx_, name); symmetricMemory_ = allocator->rendezvous(symmetricMemoryPtr_, std::nullopt); isInitialized_ = true; return true; diff --git a/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cu b/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cu index 6a6a6520e36ba..1fc655fce23be 100644 --- a/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cu +++ b/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cu @@ -23,7 +23,7 @@ static void checkInput(const at::Tensor& input, int deviceIdx) { } bool isIntraNodeCommSupported() { -#if defined(USE_ROCM) || (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)) +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800) return false; #else return true; diff --git a/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp b/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp index 7b5e8ff999c5d..8b2a425c47ce2 100644 --- a/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp +++ b/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp @@ -33,7 +33,8 @@ class TORCH_API IntraNodeComm : public c10::intrusive_ptr_target { c10::intrusive_ptr store, size_t rank, size_t worldSize, - std::optional bufferSize = std::nullopt); + std::optional bufferSize = std::nullopt, + std::string groupName = ""); ~IntraNodeComm() override; @@ -67,6 +68,7 @@ class TORCH_API IntraNodeComm : public c10::intrusive_ptr_target { size_t rank_; size_t worldSize_; size_t bufferSize_; + std::string groupName_; /** * Members initialized after rendezvous