Use dlopen to load NCCL. (#9796)

This PR adds optional support for loading nccl with `dlopen` as an alternative of compile time linking. This is to address the size bloat issue with the PyPI binary release.
- Add CMake option to load `nccl` at runtime.
- Add an NCCL stub.

After this, `nccl` will be fetched from PyPI when using pip to install XGBoost, either by a user or by `pyproject.toml`. Others who want to link the nccl at compile time can continue to do so without any change.

At the moment, this is Linux only since we only support MNMG on Linux.
This commit is contained in:
Jiaming Yuan
2023-11-22 19:27:31 +08:00
committed by GitHub
parent fedd9674c8
commit 0715ab3c10
45 changed files with 658 additions and 268 deletions

View File

@@ -171,17 +171,24 @@ function(xgboost_set_cuda_flags target)
endif()
endfunction()
macro(xgboost_link_nccl target)
function(xgboost_link_nccl target)
set(xgboost_nccl_flags -DXGBOOST_USE_NCCL=1)
if(USE_DLOPEN_NCCL)
list(APPEND xgboost_nccl_flags -DXGBOOST_USE_DLOPEN_NCCL=1)
endif()
if(BUILD_STATIC_LIB)
target_include_directories(${target} PUBLIC ${NCCL_INCLUDE_DIR})
target_compile_definitions(${target} PUBLIC -DXGBOOST_USE_NCCL=1)
target_compile_definitions(${target} PUBLIC ${xgboost_nccl_flags})
target_link_libraries(${target} PUBLIC ${NCCL_LIBRARY})
else()
target_include_directories(${target} PRIVATE ${NCCL_INCLUDE_DIR})
target_compile_definitions(${target} PRIVATE -DXGBOOST_USE_NCCL=1)
target_link_libraries(${target} PRIVATE ${NCCL_LIBRARY})
target_compile_definitions(${target} PRIVATE ${xgboost_nccl_flags})
if(NOT USE_DLOPEN_NCCL)
target_link_libraries(${target} PRIVATE ${NCCL_LIBRARY})
endif()
endif()
endmacro()
endfunction()
# compile options
macro(xgboost_target_properties target)

View File

@@ -54,17 +54,24 @@ find_path(NCCL_INCLUDE_DIR
NAMES nccl.h
HINTS ${NCCL_ROOT}/include $ENV{NCCL_ROOT}/include)
find_library(NCCL_LIBRARY
NAMES ${NCCL_LIB_NAME}
HINTS ${NCCL_ROOT}/lib $ENV{NCCL_ROOT}/lib/)
if(USE_DLOPEN_NCCL)
include(FindPackageHandleStandardArgs)
find_package_handle_standard_args(Nccl DEFAULT_MSG NCCL_INCLUDE_DIR)
message(STATUS "Using nccl library: ${NCCL_LIBRARY}")
mark_as_advanced(NCCL_INCLUDE_DIR)
else()
find_library(NCCL_LIBRARY
NAMES ${NCCL_LIB_NAME}
HINTS ${NCCL_ROOT}/lib $ENV{NCCL_ROOT}/lib/)
include(FindPackageHandleStandardArgs)
find_package_handle_standard_args(Nccl DEFAULT_MSG
NCCL_INCLUDE_DIR NCCL_LIBRARY)
message(STATUS "Using nccl library: ${NCCL_LIBRARY}")
mark_as_advanced(
NCCL_INCLUDE_DIR
NCCL_LIBRARY
)
include(FindPackageHandleStandardArgs)
find_package_handle_standard_args(Nccl DEFAULT_MSG
NCCL_INCLUDE_DIR NCCL_LIBRARY)
mark_as_advanced(
NCCL_INCLUDE_DIR
NCCL_LIBRARY
)
endif()