Allow PyTorch to be built without NCCL (#17295)

Summary:
With this patch you can use USE_DISTRIBUTED=OFF (possibly in combination with USE_NCCL=OFF (?))

The significance is partly because the NCCL doesn't build with CUDA 8.
This is written under the assumption that NCCL is required for distributed if not, the USE_DISTRIBUTED check in nccl.py should be replaced by a check for the USE_NCCL environment variable.

Fixes: #17274
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17295

Differential Revision: D14155080

Pulled By: ezyang

fbshipit-source-id: 0d133f7c5b4d118849f041bd4d4cbbd7ffc3c7b4

diff --git a/tools/build_pytorch_libs.py b/tools/build_pytorch_libs.py
index f331c09..73c55df 100644 (file)
--- a/tools/build_pytorch_libs.py
+++ b/tools/build_pytorch_libs.py
@@ -12,7 +12,7 @@ from subprocess import check_call, call, check_output
 from distutils.version import LooseVersion
 from .setup_helpers.cuda import USE_CUDA, CUDA_HOME
 from .setup_helpers.dist_check import USE_DISTRIBUTED, USE_GLOO_IBVERBS
-from .setup_helpers.nccl import USE_SYSTEM_NCCL, NCCL_INCLUDE_DIR, NCCL_ROOT_DIR, NCCL_SYSTEM_LIB
+from .setup_helpers.nccl import USE_SYSTEM_NCCL, NCCL_INCLUDE_DIR, NCCL_ROOT_DIR, NCCL_SYSTEM_LIB, USE_NCCL
 from .setup_helpers.rocm import ROCM_HOME, ROCM_VERSION, USE_ROCM
 from .setup_helpers.nnpack import USE_NNPACK
 from .setup_helpers.qnnpack import USE_QNNPACK
@@ -183,7 +183,8 @@ def run_cmake(version,
         USE_FFMPEG=check_env_flag('USE_FFMPEG'),
         USE_SYSTEM_EIGEN_INSTALL="OFF",
         USE_MKLDNN=USE_MKLDNN,
-        NCCL_EXTERNAL=USE_CUDA,
+        USE_NCCL=USE_NCCL,
+        NCCL_EXTERNAL=USE_NCCL,
         CMAKE_INSTALL_PREFIX=install_dir,
         CMAKE_C_FLAGS=cflags,
         CMAKE_CXX_FLAGS=cflags,

diff --git a/tools/setup_helpers/nccl.py b/tools/setup_helpers/nccl.py
index c1cc886..5fbbf77 100644 (file)
--- a/tools/setup_helpers/nccl.py
+++ b/tools/setup_helpers/nccl.py
@@ -8,8 +8,7 @@ from .env import IS_WINDOWS, IS_DARWIN, IS_CONDA, CONDA_DIR, check_negative_env_
 
 from .cuda import USE_CUDA, CUDA_HOME
 
-
-USE_NCCL = USE_CUDA and not IS_DARWIN and not IS_WINDOWS
+USE_NCCL = USE_CUDA and not check_negative_env_flag('USE_NCCL') and not IS_DARWIN and not IS_WINDOWS
 USE_SYSTEM_NCCL = False
 NCCL_LIB_DIR = None
 NCCL_SYSTEM_LIB = None