diff --git a/CMakeLists.txt b/CMakeLists.txt index c3fb8fba5..985c2bd62 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -226,6 +226,10 @@ if(MSCCLPP_USE_CUDA) else() set(GPU_LIBRARIES CUDA::cudart CUDA::cuda_driver) endif() + if(NOT TARGET CUDA::nvml) + message(FATAL_ERROR "CUDA NVML target CUDA::nvml is required for MSCCLPP CUDA builds.") + endif() + list(APPEND GPU_LIBRARIES CUDA::nvml) else() set(CMAKE_HIP_STANDARD 20) set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -Wall -Wextra") diff --git a/include/mscclpp/core.hpp b/include/mscclpp/core.hpp index 45b56bcc0..4c14f1eec 100644 --- a/include/mscclpp/core.hpp +++ b/include/mscclpp/core.hpp @@ -46,6 +46,10 @@ class Bootstrap { /// @return The total number of ranks per node. virtual int getNranksPerNode() const = 0; + /// Return the number of ranks in this rank's GPU IPC domain. + /// @return The number of ranks in the GPU IPC domain. + virtual int getNranksPerIpcDomain() const; + /// Send arbitrary data to another process. /// /// Data sent via `send(senderBuff, size, receiverRank, tag)` can be received via `recv(receiverBuff, size, @@ -144,6 +148,9 @@ class TcpBootstrap : public Bootstrap { /// Return the total number of ranks per node. int getNranksPerNode() const override; + /// Return the number of ranks in this rank's GPU IPC domain. + int getNranksPerIpcDomain() const override; + /// Send arbitrary data to another process. /// /// Data sent via `send(senderBuff, size, receiverRank, tag)` can be received via `recv(receiverBuff, size, diff --git a/python/csrc/core_py.cpp b/python/csrc/core_py.cpp index a94f9863a..7e9af6c1f 100644 --- a/python/csrc/core_py.cpp +++ b/python/csrc/core_py.cpp @@ -56,6 +56,7 @@ void register_core(nb::module_& m) { .def("get_rank", &Bootstrap::getRank) .def("get_n_ranks", &Bootstrap::getNranks) .def("get_n_ranks_per_node", &Bootstrap::getNranksPerNode) + .def("get_n_ranks_per_ipc_domain", &Bootstrap::getNranksPerIpcDomain) .def( "send", [](Bootstrap* self, uintptr_t ptr, size_t size, int peer, int tag) { diff --git a/python/mscclpp/_core/comm.py b/python/mscclpp/_core/comm.py index d42349ddb..9e8e6c2cb 100644 --- a/python/mscclpp/_core/comm.py +++ b/python/mscclpp/_core/comm.py @@ -73,6 +73,7 @@ def __init__( self.my_rank = self.bootstrap.get_rank() self.nranks = self.bootstrap.get_n_ranks() self.nranks_per_node = self.bootstrap.get_n_ranks_per_node() + self.nranks_per_ipc_domain = self.bootstrap.get_n_ranks_per_ipc_domain() def barrier(self): self.bootstrap.barrier() diff --git a/src/core/bootstrap/bootstrap.cc b/src/core/bootstrap/bootstrap.cc index b3032e502..2f18d70c3 100644 --- a/src/core/bootstrap/bootstrap.cc +++ b/src/core/bootstrap/bootstrap.cc @@ -50,6 +50,8 @@ MSCCLPP_API_CPP void Bootstrap::groupBarrier(const std::vector& ranks) { } } +MSCCLPP_API_CPP int Bootstrap::getNranksPerIpcDomain() const { return getNranksPerNode(); } + MSCCLPP_API_CPP void Bootstrap::send(const std::vector& data, int peer, int tag) { size_t size = data.size(); send((void*)&size, sizeof(size_t), peer, tag); @@ -83,6 +85,7 @@ class TcpBootstrap::Impl { int getRank(); int getNranks(); int getNranksPerNode(); + int getNranksPerIpcDomain(); void allGather(void* allData, int size); void broadcast(void* data, int size, int root); void send(void* data, int size, int peer, int tag); @@ -95,6 +98,7 @@ class TcpBootstrap::Impl { int rank_; int nRanks_; int nRanksPerNode_; + int nRanksPerIpcDomain_; bool netInitialized; std::unique_ptr listenSockRoot_; std::unique_ptr listenSock_; @@ -148,6 +152,7 @@ TcpBootstrap::Impl::Impl(int rank, int nRanks) : rank_(rank), nRanks_(nRanks), nRanksPerNode_(0), + nRanksPerIpcDomain_(0), netInitialized(false), peerCommAddresses_(nRanks, SocketAddress()), barrierArr_(nRanks, 0), @@ -451,6 +456,42 @@ int TcpBootstrap::Impl::getNranksPerNode() { return nRanksPerNode_; } +int TcpBootstrap::Impl::getNranksPerIpcDomain() { + if (nRanksPerIpcDomain_ > 0) return nRanksPerIpcDomain_; + std::vector ipcDomainHashes(nRanks_); + ipcDomainHashes[rank_] = getIpcDomainHash(); + allGather(ipcDomainHashes.data(), sizeof(uint64_t)); + + std::unordered_map ipcDomainCounts; + for (uint64_t ipcDomainHash : ipcDomainHashes) { + ipcDomainCounts[ipcDomainHash]++; + } + + const int nRanksPerIpcDomain = ipcDomainCounts[ipcDomainHashes[rank_]]; + const std::string invalidIpcDomainLayout = + "IPC domain ranks must have the same size and be arranged in consecutive rank order"; + if (nRanks_ % nRanksPerIpcDomain != 0) { + throw Error(invalidIpcDomainLayout + ": rank count is not divisible by IPC domain size", ErrorCode::InvalidUsage); + } + + for (const auto& entry : ipcDomainCounts) { + if (entry.second != nRanksPerIpcDomain) { + throw Error(invalidIpcDomainLayout + ": IPC domains have different sizes", ErrorCode::InvalidUsage); + } + } + for (int i = 0; i < nRanks_; ++i) { + const int ipcDomainFirstRank = i - i % nRanksPerIpcDomain; + if (ipcDomainHashes[i] != ipcDomainHashes[ipcDomainFirstRank]) { + throw Error(invalidIpcDomainLayout + ": ranks are not grouped by IPC domain", ErrorCode::InvalidUsage); + } + } + + INFO(MSCCLPP_INIT, "rank %d IPC domain fabric hash 0x%016llx nRanksPerIpcDomain %d", rank_, + static_cast(ipcDomainHashes[rank_]), nRanksPerIpcDomain); + nRanksPerIpcDomain_ = nRanksPerIpcDomain; + return nRanksPerIpcDomain_; +} + void TcpBootstrap::Impl::allGather(void* allData, int size) { char* data = static_cast(allData); int rank = rank_; @@ -592,6 +633,8 @@ MSCCLPP_API_CPP int TcpBootstrap::getNranks() const { return pimpl_->getNranks() MSCCLPP_API_CPP int TcpBootstrap::getNranksPerNode() const { return pimpl_->getNranksPerNode(); } +MSCCLPP_API_CPP int TcpBootstrap::getNranksPerIpcDomain() const { return pimpl_->getNranksPerIpcDomain(); } + MSCCLPP_API_CPP void TcpBootstrap::send(void* data, int size, int peer, int tag) { pimpl_->send(data, size, peer, tag); } diff --git a/src/core/include/utils_internal.hpp b/src/core/include/utils_internal.hpp index c5c67e26c..c6934194d 100644 --- a/src/core/include/utils_internal.hpp +++ b/src/core/include/utils_internal.hpp @@ -37,6 +37,7 @@ int64_t busIdToInt64(const std::string busId); uint64_t getHash(const char* string, int n); uint64_t getHostHash(); uint64_t getPidHash(); +uint64_t getIpcDomainHash(); void getRandomData(void* buffer, size_t bytes); struct netIf { diff --git a/src/core/utils_internal.cc b/src/core/utils_internal.cc index 8cc554301..d48c42254 100644 --- a/src/core/utils_internal.cc +++ b/src/core/utils_internal.cc @@ -6,6 +6,10 @@ #include #include +#if defined(MSCCLPP_USE_CUDA) +#include +#endif + #include #include #include @@ -175,6 +179,83 @@ uint64_t getPidHash(void) { return *pidHash; } +#if defined(MSCCLPP_USE_CUDA) && defined(NVML_GPU_FABRIC_UUID_LEN) +namespace { + +class NvmlState { + public: + NvmlState() : initialized_(nvmlInit_v2() == NVML_SUCCESS) {} + + NvmlState(const NvmlState&) = delete; + NvmlState& operator=(const NvmlState&) = delete; + NvmlState(NvmlState&&) = delete; + NvmlState& operator=(NvmlState&&) = delete; + + ~NvmlState() { + if (initialized_) { + (void)nvmlShutdown(); + } + } + + bool isInitialized() const { return initialized_; } + + private: + bool initialized_ = false; +}; + +template +uint64_t getFabricHash(const FabricInfo& fabricInfo) { + char hashData[NVML_GPU_FABRIC_UUID_LEN + sizeof(fabricInfo.cliqueId)]; + std::memcpy(hashData, fabricInfo.clusterUuid, NVML_GPU_FABRIC_UUID_LEN); + std::memcpy(hashData + NVML_GPU_FABRIC_UUID_LEN, &fabricInfo.cliqueId, sizeof(fabricInfo.cliqueId)); + return getHash(hashData, sizeof(hashData)); +} + +bool tryGetNvmlIpcDomainHash(uint64_t& ipcDomainHash) { + // Use the current CUDA device; callers must set the rank's device before querying. + int deviceId; + char pciBusId[] = "00000000:00:00.0"; + if (cudaGetDevice(&deviceId) != cudaSuccess || + cudaDeviceGetPCIBusId(pciBusId, sizeof(pciBusId), deviceId) != cudaSuccess) { + return false; + } + + static NvmlState nvml; + nvmlDevice_t nvmlDevice; + if (!nvml.isInitialized() || nvmlDeviceGetHandleByPciBusId_v2(pciBusId, &nvmlDevice) != NVML_SUCCESS) { + return false; + } + +#if defined(nvmlGpuFabricInfo_v2) + nvmlGpuFabricInfoV_t fabricInfo = {}; + fabricInfo.version = nvmlGpuFabricInfo_v2; + nvmlReturn_t result = nvmlDeviceGetGpuFabricInfoV(nvmlDevice, &fabricInfo); +#else + nvmlGpuFabricInfo_t fabricInfo = {}; + nvmlReturn_t result = nvmlDeviceGetGpuFabricInfo(nvmlDevice, &fabricInfo); +#endif + if (result != NVML_SUCCESS || fabricInfo.state != NVML_GPU_FABRIC_STATE_COMPLETED || + fabricInfo.status != NVML_SUCCESS) { + return false; + } + + ipcDomainHash = getFabricHash(fabricInfo); + return true; +} + +} // namespace +#endif + +uint64_t getIpcDomainHash(void) { +#if defined(MSCCLPP_USE_CUDA) && defined(NVML_GPU_FABRIC_UUID_LEN) + uint64_t ipcDomainHash; + if (tryGetNvmlIpcDomainHash(ipcDomainHash)) { + return ipcDomainHash; + } +#endif + return getHostHash(); +} + int parseStringList(const char* string, netIf* ifList, int maxList) { if (!string) return 0; diff --git a/test/mp_unit/bootstrap_tests.cc b/test/mp_unit/bootstrap_tests.cc index c28087a45..27ca1f5fb 100644 --- a/test/mp_unit/bootstrap_tests.cc +++ b/test/mp_unit/bootstrap_tests.cc @@ -42,10 +42,17 @@ void BootstrapTest::bootstrapTestSendRecv(std::shared_ptr bo } } +void BootstrapTest::bootstrapTestIpcDomain(std::shared_ptr bootstrap) { + int nRanksPerIpcDomain = bootstrap->getNranksPerIpcDomain(); + EXPECT_GT(nRanksPerIpcDomain, 0); + EXPECT_LE(nRanksPerIpcDomain, bootstrap->getNranks()); +} + void BootstrapTest::bootstrapTestAll(std::shared_ptr bootstrap) { bootstrapTestAllGather(bootstrap); bootstrapTestBarrier(bootstrap); bootstrapTestSendRecv(bootstrap); + bootstrapTestIpcDomain(bootstrap); } TEST(BootstrapTest, WithId) { @@ -127,6 +134,7 @@ class MPIBootstrap : public mscclpp::Bootstrap { MPI_Comm_size(shmcomm, &shmrank); return shmrank; } + int getNranksPerIpcDomain() const override { return getNranksPerNode(); } void allGather(void* sendbuf, int size) override { MPI_Allgather(MPI_IN_PLACE, 0, MPI_BYTE, sendbuf, size, MPI_BYTE, MPI_COMM_WORLD); } diff --git a/test/mp_unit/mp_unit_tests.hpp b/test/mp_unit/mp_unit_tests.hpp index eb8b54859..d01857f3f 100644 --- a/test/mp_unit/mp_unit_tests.hpp +++ b/test/mp_unit/mp_unit_tests.hpp @@ -56,6 +56,8 @@ class BootstrapTest : public MultiProcessTest { void bootstrapTestSendRecv(std::shared_ptr bootstrap); + void bootstrapTestIpcDomain(std::shared_ptr bootstrap); + void bootstrapTestAll(std::shared_ptr bootstrap); // Each test case should finish within 30 seconds.