diff --git a/mlx/distributed/nccl/CMakeLists.txt b/mlx/distributed/nccl/CMakeLists.txt index 389a86c66..4c59c56e2 100644 --- a/mlx/distributed/nccl/CMakeLists.txt +++ b/mlx/distributed/nccl/CMakeLists.txt @@ -9,11 +9,17 @@ if(MLX_BUILD_CUDA) STATUS "NCCL not found, using stubs. To run distributed with NCCL backend, install NCCL." ) - file( - DOWNLOAD - "https://raw.githubusercontent.com/NVIDIA/nccl/refs/tags/v2.27.5-1/src/nccl.h.in" - "${CMAKE_CURRENT_BINARY_DIR}/nccl.h") - target_include_directories(mlx PRIVATE "${CMAKE_CURRENT_BINARY_DIR}") + + include(ExternalProject) + ExternalProject_Add( + nccl_stub + SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/nccl_stub" + BUILD_COMMAND ${CMAKE_COMMAND} --build . + INSTALL_COMMAND "") + set(NCCL_PATH + "${CMAKE_CURRENT_BINARY_DIR}/nccl_stub-prefix/src/nccl_stub-build/") + target_link_libraries(mlx PRIVATE ${NCCL_PATH}/libnccl.so) + target_include_directories(mlx PRIVATE ${NCCL_PATH}) endif() else() target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp) diff --git a/mlx/distributed/nccl/nccl_stub/CMakeLists.txt b/mlx/distributed/nccl/nccl_stub/CMakeLists.txt new file mode 100644 index 000000000..54d270347 --- /dev/null +++ b/mlx/distributed/nccl/nccl_stub/CMakeLists.txt @@ -0,0 +1,13 @@ +cmake_minimum_required(VERSION 3.25) + +project(nccl LANGUAGES C CXX) + +file( + DOWNLOAD + "https://raw.githubusercontent.com/NVIDIA/nccl/refs/tags/v2.27.5-1/src/nccl.h.in" + "${CMAKE_CURRENT_BINARY_DIR}/nccl.h") + +add_library(nccl SHARED nccl_stubs.cpp) +find_package(CUDAToolkit REQUIRED) +target_include_directories(nccl PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) +target_include_directories(nccl PRIVATE ${CMAKE_CURRENT_BINARY_DIR}) diff --git a/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp b/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp new file mode 100644 index 000000000..171340f67 --- /dev/null +++ b/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp @@ -0,0 +1,7 @@ +#include +#include +#include + +ncclResult_t ncclGetUniqueId(ncclUniqueId*) { + return ncclSuccess; +}