zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-28 00:36:32 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Use nccl header only when nccl is not present (#2539)

Browse Source 
* use nccl header only when nccl is not present

* larger machine for cuda build
This commit is contained in:
Awni Hannun 2025-08-25 14:17:25 -07:00 committed by GitHub
parent 333ffea273
commit d2f540f4e0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 16 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.circleci/config.yml
View File

@ -394,7 +394,7 @@ jobs:
default: ""
machine:
image: ubuntu-2204:current
resource_class: large
resource_class: xlarge
steps:
- checkout
- run:

18
mlx/distributed/nccl/CMakeLists.txt
View File

@ -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()