fix MPI & NCCL in cmake

InternLM · Sep 25, 2023 · 763dd23 · 763dd23
1 parent e980377
commit 763dd23
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 cmake_minimum_required(VERSION 3.11 FATAL_ERROR) # for PyTorch extensions, version should be greater than 3.13
+cmake_policy(SET CMP0074 NEW)
 project(TurboMind LANGUAGES CXX CUDA)
 
 find_package(CUDA 10.2 REQUIRED)
@@ -265,8 +266,8 @@ print(torch._C._GLIBCXX_USE_CXX11_ABI,end='');"
 endif()
 
 if (BUILD_MULTI_GPU)
-  list(APPEND COMMON_HEADER_DIRS ${MPI_INCLUDE_PATH})
-  list(APPEND COMMON_LIB_DIRS /usr/local/mpi/lib)
+  list(APPEND COMMON_HEADER_DIRS ${MPI_CXX_INCLUDE_DIRS})
+  list(APPEND COMMON_HEADER_DIRS ${NCCL_INCLUDE_DIRS})
 endif()
 
 if(USE_TRITONSERVER_DATATYPE)
@@ -347,7 +348,7 @@ add_library(transformer-shared SHARED
 
 if (BUILD_MULTI_GPU)
 target_link_libraries(transformer-shared PUBLIC
-  -lmpi
+  ${MPI_CXX_LIBRARIES}
   ${NCCL_LIBRARIES}
 )
 endif()

diff --git a/src/turbomind/utils/CMakeLists.txt b/src/turbomind/utils/CMakeLists.txt
@@ -63,7 +63,7 @@ add_library(mpi_utils STATIC mpi_utils.cc)
 set_property(TARGET mpi_utils PROPERTY POSITION_INDEPENDENT_CODE  ON)
 set_property(TARGET mpi_utils PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS  ON)
 if (BUILD_MULTI_GPU)
-    target_link_libraries(mpi_utils PUBLIC mpi logger)
+    target_link_libraries(mpi_utils PUBLIC ${MPI_CXX_LIBRARIES} logger)
 endif()
 
 add_library(nccl_utils STATIC nccl_utils.cc)