frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Turn on compile with NVSHMEM (#154538)

Browse Source 
Before:
`USE_NVSHMEM=1` need to be explicit set in build environment.

After:
`USE_NVSHMEM=1` is the default for CUDA/Rocm on Linux.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/154538
Approved by: https://github.com/ngimel
This commit is contained in:
Ke Wen
2025-05-29 11:44:25 -07:00
committed by PyTorch MergeBot
parent a1a268aff5
commit 3685b10170
3 changed files with 7 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
CMakeLists.txt
View File

@ -265,6 +265,8 @@ cmake_dependent_option(USE_RCCL "Use RCCL" ON USE_NCCL OFF)
cmake_dependent_option(USE_STATIC_NCCL "Use static NCCL" OFF "USE_NCCL" OFF)
cmake_dependent_option(USE_SYSTEM_NCCL "Use system-wide NCCL" OFF "USE_NCCL"
OFF)
cmake_dependent_option(USE_NVSHMEM "Use NVSHMEM" ON
"USE_CUDA OR USE_ROCM;UNIX;NOT APPLE" OFF)
option(USE_NNAPI "Use NNAPI" OFF)
option(USE_NNPACK "Use NNPACK" ON)
cmake_dependent_option(USE_NUMA "Use NUMA. Only available on Linux." ON "LINUX"

14
caffe2/CMakeLists.txt
View File

@ -985,22 +985,16 @@ elseif(USE_CUDA)
target_compile_definitions(torch_cuda PRIVATE USE_NCCL)
endif()
# Use env var for these for now for prototyping purposes
set(USE_NVSHMEM $ENV{USE_NVSHMEM} CACHE BOOL "Whether to build with NVSHMEM support")
# Compile with NVSHMEM
# Default value of `USE_NVSHMEM` is set in CMakeLists.txt under root, to ON.
# If user has specified NVSHMEM_HOME, we use it;
# Otherwise, NVSHMEM_HOME is auto detected in tools/setup_helpers/cmake.py
if($ENV{NVSHMEM_HOME})
set(NVSHMEM_HOME $ENV{NVSHMEM_HOME} CACHE PATH "Path to NVSHMEM build dir")
endif()
if(USE_NVSHMEM AND NOT DEFINED NVSHMEM_HOME)
message(WARNING "USE_NVSHMEM set to 1 but NVSHMEM_HOME not found. Please run `pip install nvidia-nvshmem-<version>`, or set NVSHMEM_HOME to the NVSHMEM build dir")
# Disable nvshmem if NVSHMEM_HOME is not found
set(USE_NVSHMEM FALSE CACHE BOOL "Whether to build with NVSHMEM support")
endif()
if(USE_NVSHMEM)
message("Building with NVSHMEM support: '${NVSHMEM_HOME}'")
if(USE_NVSHMEM AND DEFINED NVSHMEM_HOME)
message(STATUS "Building with NVSHMEM support: '${NVSHMEM_HOME}'")
set(NVSHMEM_INCLUDE_DIR "${NVSHMEM_HOME}/include")
set(NVSHMEM_LIB_DIR "${NVSHMEM_HOME}/lib")

1
cmake/Summary.cmake
View File

@ -172,6 +172,7 @@ function(caffe2_print_configuration_summary)
if(${USE_NCCL})
message(STATUS " USE_SYSTEM_NCCL : ${USE_SYSTEM_NCCL}")
endif()
message(STATUS " USE_NVSHMEM : ${USE_NVSHMEM}")
message(STATUS " USE_NNPACK : ${USE_NNPACK}")
message(STATUS " USE_NUMPY : ${USE_NUMPY}")
message(STATUS " USE_OBSERVERS : ${USE_OBSERVERS}")