diff --git a/.circleci/config.yml b/.circleci/config.yml index 502640d3e..14bac4437 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -394,7 +394,7 @@ jobs: default: "" machine: image: ubuntu-2204:current - resource_class: large + resource_class: xlarge steps: - checkout - run: diff --git a/mlx/distributed/nccl/CMakeLists.txt b/mlx/distributed/nccl/CMakeLists.txt index 2f764c6ac..389a86c66 100644 --- a/mlx/distributed/nccl/CMakeLists.txt +++ b/mlx/distributed/nccl/CMakeLists.txt @@ -1,8 +1,20 @@ 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}) + find_package(NCCL) + if(NCCL_FOUND) + target_link_libraries(mlx PRIVATE ${NCCL_LIBRARIES}) + target_include_directories(mlx PRIVATE ${NCCL_INCLUDE_DIRS}) + else() + message( + 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}") + endif() else() target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp) endif()