diff --git a/cmake/FindNCCL.cmake b/cmake/FindNCCL.cmake new file mode 100644 index 000000000..7f8791476 --- /dev/null +++ b/cmake/FindNCCL.cmake @@ -0,0 +1,64 @@ +# Find the nccl libraries +# +# The following variables are optionally searched for defaults NCCL_ROOT_DIR: +# Base directory where all NCCL components are found NCCL_INCLUDE_DIR: Directory +# where NCCL header is found NCCL_LIB_DIR: Directory where NCCL library is found +# +# The following are set after configuration is done: NCCL_FOUND +# NCCL_INCLUDE_DIRS NCCL_LIBRARIES +# +# The path hints include CUDA_TOOLKIT_ROOT_DIR seeing as some folks install NCCL +# in the same location as the CUDA toolkit. See +# https://github.com/caffe2/caffe2/issues/1601 + +set(NCCL_ROOT_DIR + $ENV{NCCL_ROOT_DIR} + CACHE PATH "Folder contains NVIDIA NCCL") + +find_path( + NCCL_INCLUDE_DIRS + NAMES nccl.h + HINTS ${NCCL_INCLUDE_DIR} ${NCCL_ROOT_DIR} ${NCCL_ROOT_DIR}/include + ${CUDA_TOOLKIT_ROOT_DIR}/include) + +if($ENV{USE_STATIC_NCCL}) + message( + STATUS "USE_STATIC_NCCL detected. Linking against static NCCL library") + set(NCCL_LIBNAME "libnccl_static.a") +else() + set(NCCL_LIBNAME "nccl") +endif() + +find_library( + NCCL_LIBRARIES + NAMES ${NCCL_LIBNAME} + HINTS ${NCCL_LIB_DIR} + ${NCCL_ROOT_DIR} + ${NCCL_ROOT_DIR}/lib + ${NCCL_ROOT_DIR}/lib/x86_64-linux-gnu + ${NCCL_ROOT_DIR}/lib64 + ${CUDA_TOOLKIT_ROOT_DIR}/lib + ${CUDA_TOOLKIT_ROOT_DIR}/lib64) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(NCCL DEFAULT_MSG NCCL_INCLUDE_DIRS + NCCL_LIBRARIES) + +if(NCCL_FOUND) + set(NCCL_HEADER_FILE "${NCCL_INCLUDE_DIRS}/nccl.h") + message( + STATUS "Determining NCCL version from the header file: ${NCCL_HEADER_FILE}") + file( + STRINGS ${NCCL_HEADER_FILE} NCCL_MAJOR_VERSION_DEFINED + REGEX "^[ \t]*#define[ \t]+NCCL_MAJOR[ \t]+[0-9]+.*$" + LIMIT_COUNT 1) + if(NCCL_MAJOR_VERSION_DEFINED) + string(REGEX REPLACE "^[ \t]*#define[ \t]+NCCL_MAJOR[ \t]+" "" + NCCL_MAJOR_VERSION ${NCCL_MAJOR_VERSION_DEFINED}) + message(STATUS "NCCL_MAJOR_VERSION: ${NCCL_MAJOR_VERSION}") + endif() + message( + STATUS + "Found NCCL (include: ${NCCL_INCLUDE_DIRS}, library: ${NCCL_LIBRARIES})") + mark_as_advanced(NCCL_ROOT_DIR NCCL_INCLUDE_DIRS NCCL_LIBRARIES) +endif() diff --git a/mlx/backend/cuda/primitives.cu b/mlx/backend/cuda/primitives.cu index e32befc9c..ec1dde600 100644 --- a/mlx/backend/cuda/primitives.cu +++ b/mlx/backend/cuda/primitives.cu @@ -54,6 +54,28 @@ bool fast::ScaledDotProductAttention::use_fallback( return true; } +namespace distributed { +void AllReduce::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + // Here I assume for now that in is donatable and contiguous. + // TODO + + auto& input = inputs[0]; + auto& output = outputs[0]; + + output.copy_shared_buffer(input); + auto& s = stream(); + switch (reduce_type_) { + case Sum: + distributed::detail::all_sum(group(), input, output, s); + break; + default: + throw std::runtime_error("Only all reduce sum is supported for now"); + } +} +} // namespace distributed + #define NO_GPU_MULTI(func) \ void func::eval_gpu( \ const std::vector& inputs, std::vector& outputs) { \ @@ -97,7 +119,6 @@ NO_GPU_MULTI(CustomKernel) } // namespace fast namespace distributed { -NO_GPU_MULTI(AllReduce) NO_GPU_MULTI(AllGather) NO_GPU_MULTI(Send) NO_GPU_MULTI(Recv) diff --git a/mlx/distributed/CMakeLists.txt b/mlx/distributed/CMakeLists.txt index 8e16bd40d..b7762f6a7 100644 --- a/mlx/distributed/CMakeLists.txt +++ b/mlx/distributed/CMakeLists.txt @@ -6,3 +6,4 @@ target_sources( add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/mpi) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/ring) +add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/nccl) diff --git a/mlx/distributed/distributed.cpp b/mlx/distributed/distributed.cpp index cc01e6090..f791ee29e 100644 --- a/mlx/distributed/distributed.cpp +++ b/mlx/distributed/distributed.cpp @@ -5,6 +5,7 @@ #include "mlx/distributed/distributed.h" #include "mlx/distributed/distributed_impl.h" #include "mlx/distributed/mpi/mpi.h" +#include "mlx/distributed/nccl/nccl.h" #include "mlx/distributed/ring/ring.h" namespace mlx::core::distributed { @@ -111,6 +112,8 @@ Group init(bool strict /* = false */, const std::string& bk /* = "any" */) { group = mpi::init(strict); } else if (bk == "ring") { group = ring::init(strict); + } else if (bk == "nccl") { + group = nccl::init(strict); } else if (bk == "any") { group = ring::init(false); bk_ = "ring"; diff --git a/mlx/distributed/nccl/CMakeLists.txt b/mlx/distributed/nccl/CMakeLists.txt new file mode 100644 index 000000000..2f764c6ac --- /dev/null +++ b/mlx/distributed/nccl/CMakeLists.txt @@ -0,0 +1,8 @@ +if(MLX_BUILD_CUDA) + target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nccl.cpp) + find_package(NCCL REQUIRED) + target_link_libraries(mlx PRIVATE ${NCCL_LIBRARIES}) + target_include_directories(mlx PRIVATE ${NCCL_INCLUDE_DIRS}) +else() + target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp) +endif() diff --git a/mlx/distributed/nccl/nccl.cpp b/mlx/distributed/nccl/nccl.cpp new file mode 100644 index 000000000..02b1fc20c --- /dev/null +++ b/mlx/distributed/nccl/nccl.cpp @@ -0,0 +1,389 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mlx/backend/cuda/device.h" +#include "mlx/distributed/distributed.h" +#include "mlx/distributed/distributed_impl.h" + +namespace mlx::core::distributed::nccl { + +#define CHECK_CUDA(cmd) \ + do { \ + cudaError_t e = cmd; \ + if (e != cudaSuccess) { \ + fprintf( \ + stderr, \ + "CUDA error %s:%d '%s'\n", \ + __FILE__, \ + __LINE__, \ + cudaGetErrorString(e)); \ + exit(1); \ + } \ + } while (0) + +#define CHECK_NCCL(cmd) \ + do { \ + ncclResult_t r = cmd; \ + if (r != ncclSuccess) { \ + fprintf( \ + stderr, \ + "NCCL error %s:%d '%s'\n", \ + __FILE__, \ + __LINE__, \ + ncclGetErrorString(r)); \ + exit(1); \ + } \ + } while (0) + +namespace detail { + +inline void sendAll(int sock, const void* buf, size_t len) { + const char* ptr = reinterpret_cast(buf); + while (len > 0) { + ssize_t sent = send(sock, ptr, len, 0); + if (sent <= 0) { + perror("send"); + exit(1); + } + ptr += sent; + len -= sent; + } +} + +inline void recvAll(int sock, void* buf, size_t len) { + char* ptr = reinterpret_cast(buf); + while (len > 0) { + ssize_t rec = recv(sock, ptr, len, 0); + if (rec <= 0) { + perror("recv"); + exit(1); + } + ptr += rec; + len -= rec; + } +} + +inline void bootstrap_unique_id( + ncclUniqueId& id, + int rank, + int size, + const std::string& initMethod) { + // Parse the init method to extract the host and port + if (initMethod.rfind("tcp://", 0) != 0) + throw; + auto hostport = initMethod.substr(6); + auto colon = hostport.find(':'); + std::string host = hostport.substr(0, colon); + int port = std::stoi(hostport.substr(colon + 1)); + + if (rank == 0) { + // create a unique id on the rank 0 + CHECK_NCCL(ncclGetUniqueId(&id)); + + // create a socket to send the unique id to all other ranks + int sock = socket(AF_INET, SOCK_STREAM, 0); + + if (sock < 0) { + std::ostringstream msg; + msg << "[nccl] Couldn't create socket (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } + + sockaddr_in serv = {}; + serv.sin_family = AF_INET; + serv.sin_addr.s_addr = htonl(INADDR_ANY); + serv.sin_port = htons(port); + + int reuse = 1; + // Without this, if rank-0 crashes or restarts process quickly, + // the OS might refuse to let binding to the same port, so reuse + + if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) < 0) { + std::ostringstream msg; + msg << "[nccl] setsockopt() failed: " << strerror(errno); + throw std::runtime_error(msg.str()); + } + + if (bind(sock, reinterpret_cast(&serv), sizeof(serv)) < 0) { + std::ostringstream msg; + msg << "[nccl] bind() failed: " << strerror(errno); + throw std::runtime_error(msg.str()); + } + if (listen(sock, size - 1) < 0) { + std::ostringstream msg; + msg << "[nccl] listen() failed: " << strerror(errno); + throw std::runtime_error(msg.str()); + } + + for (int peer = 1; peer < size; ++peer) { + int conn = accept(sock, nullptr, nullptr); + if (conn < 0) { + std::ostringstream msg; + msg << "[nccl] accept() failed: " << strerror(errno); + throw std::runtime_error(msg.str()); + } + sendAll(conn, &id, sizeof(id)); + close(conn); + } + close(sock); + + } else { + // Here just wanted to make show that rank 0 has enough time to bind + // so we will retry to connect until max attempts + + int sock = socket(AF_INET, SOCK_STREAM, 0); + if (sock < 0) { + std::ostringstream msg; + msg << "[nccl] socket() failed: " << strerror(errno); + throw std::runtime_error(msg.str()); + } + + hostent* he = gethostbyname(host.c_str()); + if (!he) { + throw std::runtime_error("[nccl] lookup failed for host: " + host); + } + sockaddr_in serv = {}; + serv.sin_family = AF_INET; + memcpy(&serv.sin_addr, he->h_addr_list[0], he->h_length); + serv.sin_port = htons(port); + + const int max_retries = 30; + int attempt = 0; + bool connected = false; + + for (attempt = 0; attempt < max_retries; ++attempt) { + if (connect(sock, reinterpret_cast(&serv), sizeof(serv)) == + 0) { + connected = true; + std::cout << "[Rank " << rank << "] Connected successfully on attempt " + << attempt + 1 << std::endl; + break; + } + if (errno != ECONNREFUSED) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + } + + if (!connected) { + std::ostringstream msg; + msg << "[Rank " << rank << "] connect() failed after " << attempt + << " retries: " << strerror(errno); + close(sock); + throw std::runtime_error(msg.str()); + } + recvAll(sock, &id, sizeof(id)); + close(sock); + } +} + +template +struct type_identity { + using type = T; +}; + +template +void dispatch_dtype(const array& arr, F&& f) { + switch (arr.dtype()) { + case bool_: + throw std::invalid_argument("[nccl] Boolean arrays not supported"); + case int8: + f(type_identity{}, ncclChar); + break; + case uint8: + f(type_identity{}, ncclUint8); + break; + case int32: + f(type_identity{}, ncclInt); + break; + case uint32: + f(type_identity{}, ncclUint32); + break; + case int64: + f(type_identity{}, ncclInt64); + break; + case uint64: + f(type_identity{}, ncclUint64); + break; + case float16: + f(type_identity{}, ncclHalf); + break; + case bfloat16: + f(type_identity{}, ncclBfloat16); + break; + case float32: + f(type_identity{}, ncclFloat); + break; + case float64: + f(type_identity{}, ncclDouble); + break; + default: + throw std::invalid_argument("[nccl] Unknown or unsupported dtype"); + } +} + +} // namespace detail + +using GroupImpl = mlx::core::distributed::detail::GroupImpl; +class NCCLGroup : public GroupImpl { + public: + NCCLGroup(int worldRank, int worldSize, const std::string initMethod) + : rank_(worldRank), + size_(worldSize), + comm_(nullptr), + initMethod_(initMethod) { + if (initialized_) + return; + int ndev; + CHECK_CUDA(cudaGetDeviceCount(&ndev)); + CHECK_CUDA(cudaSetDevice(rank_ % ndev)); + detail::bootstrap_unique_id(uniqueId_, rank_, size_, initMethod_); + CHECK_NCCL(ncclCommInitRank(&comm_, size_, uniqueId_, rank_)); + initialized_ = true; + } + + ~NCCLGroup() { + ncclCommDestroy(comm_); + ncclGroupEnd(); + initialized_ = false; + } + + int rank() override { + return rank_; + } + + int size() override { + return size_; + } + + void all_sum(const array& input, array& output, Stream stream) override { + detail::dispatch_dtype(input, [&](auto type_tag, ncclDataType_t dt) { + using T = typename decltype(type_tag)::type; + detail::all_reduce_impl(input, output, stream, dt, ncclSum); + }); + } + + virtual std::shared_ptr split(int color, int key = -1) override { + throw std::runtime_error("[nccl] Group split not supported."); + } + + void all_gather(const array& input, array& output, Stream stream) override { + detail::dispatch_dtype(input, [&](auto type_tag, ncclDataType_t dt) { + using T = typename decltype(type_tag)::type; + CHECK_NCCL(ncclAllGather( + input.data(), + output.data(), + input.size(), + dt, + comm_, + cu::get_stream(stream).last_cuda_stream())); + }); + } + + void send(const array& input, int dst, Stream stream) override { + detail::dispatch_dtype(input, [&](auto type_tag, ncclDataType_t dt) { + using T = typename decltype(type_tag)::type; + CHECK_NCCL(ncclSend( + input.data(), + input.size(), + dt, + dst, + comm_, + cu::get_stream(stream).last_cuda_stream())); + }); + } + + void recv(array& output, int src, Stream stream) override { + detail::dispatch_dtype(output, [&](auto type_tag, ncclDataType_t dt) { + using T = typename decltype(type_tag)::type; + CHECK_NCCL(ncclRecv( + output.data(), + output.size(), + dt, + src, + comm_, + cu::get_stream(stream).last_cuda_stream())); + }); + } + + void all_max(const array& input, array& output, Stream stream) override { + detail::dispatch_dtype(input, [&](auto type_tag, ncclDataType_t dt) { + using T = typename decltype(type_tag)::type; + all_reduce_impl(input, output, stream, dt, ncclMax); + }); + } + + void all_min(const array& input, array& output, Stream stream) override { + detail::dispatch_dtype(input, [&](auto type_tag, ncclDataType_t dt) { + using T = typename decltype(type_tag)::type; + all_reduce_impl(input, output, stream, dt, ncclMin); + }); + } + + template + void all_reduce_impl( + const array& input, + array& output, + Stream stream, + ncclDataType_t dt, + ncclRedOp_t op) { + CHECK_NCCL(ncclAllReduce( + input.data(), + output.data(), + input.size(), + dt, + op, + comm_, + cu::get_stream(stream).last_cuda_stream())); + } + + int rank_, size_; + std::string initMethod_; + ncclUniqueId uniqueId_; + ncclComm_t comm_; + bool initialized_ = false; +}; + +bool is_available() { + return true; +} + +namespace detail { +static std::string get_env_var_or_throw(const char* env_var_name) { + const char* value = std::getenv(env_var_name); + if (value == nullptr) { + std::ostringstream msg; + msg << "[nccl] Required environment variable '" << env_var_name + << "' is not set. " + << "Please set it before initializing the distributed backend."; + throw std::runtime_error(msg.str()); + } + return std::string(value); +} +} // namespace detail + +std::shared_ptr init(bool strict /* = false */) { + std::string host = detail::get_env_var_or_throw("NCCL_HOST_IP"); + std::string port = detail::get_env_var_or_throw("NCCL_PORT"); + std::string rank_str = detail::get_env_var_or_throw("MLX_RANK"); + std::string n_nodes_str = detail::get_env_var_or_throw("MLX_WORLD_SIZE"); + + int rank = std::stoi(rank_str); + int n_nodes = std::stoi(n_nodes_str); + std::string init_method = "tcp://" + host + ":" + port; + + return std::make_shared(rank, n_nodes, init_method); +} +} // namespace mlx::core::distributed::nccl diff --git a/mlx/distributed/nccl/nccl.h b/mlx/distributed/nccl/nccl.h new file mode 100644 index 000000000..5370d2daf --- /dev/null +++ b/mlx/distributed/nccl/nccl.h @@ -0,0 +1,12 @@ +// Copyright © 2024 Apple Inc. + +#include "mlx/distributed/distributed.h" + +namespace mlx::core::distributed::nccl { + +using GroupImpl = mlx::core::distributed::detail::GroupImpl; + +bool is_available(); +std::shared_ptr init(bool strict = false); + +} // namespace mlx::core::distributed::nccl diff --git a/mlx/distributed/nccl/no_nccl.cpp b/mlx/distributed/nccl/no_nccl.cpp new file mode 100644 index 000000000..1be256a11 --- /dev/null +++ b/mlx/distributed/nccl/no_nccl.cpp @@ -0,0 +1,20 @@ +// Copyright © 2024 Apple Inc. + +#include "mlx/distributed/nccl/nccl.h" + +namespace mlx::core::distributed::nccl { + +using GroupImpl = mlx::core::distributed::detail::GroupImpl; + +bool is_available() { + return false; +} + +std::shared_ptr init(bool strict /* = false */) { + if (strict) { + throw std::runtime_error("Cannot initialize nccl distributed backend."); + } + return nullptr; +} + +} // namespace mlx::core::distributed::nccl diff --git a/mlx/distributed/ops.cpp b/mlx/distributed/ops.cpp index 0a5114805..9c251a944 100644 --- a/mlx/distributed/ops.cpp +++ b/mlx/distributed/ops.cpp @@ -31,8 +31,7 @@ array all_sum( return array( x.shape(), x.dtype(), - std::make_shared( - to_stream(s, Device::cpu), group, AllReduce::Sum), + std::make_shared(to_stream(s), group, AllReduce::Sum), {x}); }