Add ability to override cuda archs for conda (#1282)

pytorch · Jan 31, 2023 · 0a2f8ef · 0a2f8ef
1 parent ff4c588
commit 0a2f8ef
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/conda/pytorch-nightly/build.sh b/conda/pytorch-nightly/build.sh
@@ -52,31 +52,35 @@ if [[ -z "$USE_CUDA" || "$USE_CUDA" == 1 ]]; then
 fi
 if [[ -n "$build_with_cuda" ]]; then
     export TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
-    export TORCH_CUDA_ARCH_LIST="3.7+PTX;5.0"
+    TORCH_CUDA_ARCH_LIST="3.7+PTX;5.0"
     export USE_STATIC_CUDNN=1 # links cudnn statically (driven by tools/setup_helpers/cudnn.py)
 
     if [[ $CUDA_VERSION == 11.6* ]]; then
-        export TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;6.0;6.1;7.0;7.5;8.0;8.6"
+        TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;6.0;6.1;7.0;7.5;8.0;8.6"
         #for cuda 11.5 we use cudnn 8.3.2.44 https://docs.nvidia.com/deeplearning/cudnn/release-notes/rel_8.html
         #which does not have single static libcudnn_static.a deliverable to link with
         export USE_STATIC_CUDNN=0
         #for cuda 11.5 include all dynamic loading libraries
         DEPS_LIST=(/usr/local/cuda/lib64/libcudnn*.so.8 /usr/local/cuda-11.6/extras/CUPTI/lib64/libcupti.so.11.6)
     elif [[ $CUDA_VERSION == 11.7* ]]; then
-        export TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;6.0;6.1;7.0;7.5;8.0;8.6"
+        TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;6.0;6.1;7.0;7.5;8.0;8.6"
         #for cuda 11.7 we use cudnn 8.5
         #which does not have single static libcudnn_static.a deliverable to link with
         export USE_STATIC_CUDNN=0
         #for cuda 11.7 include all dynamic loading libraries
         DEPS_LIST=(/usr/local/cuda/lib64/libcudnn*.so.8 /usr/local/cuda-11.7/extras/CUPTI/lib64/libcupti.so.11.7)
     elif [[ $CUDA_VERSION == 11.8* ]]; then
-	export TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;6.0;6.1;7.0;7.5;8.0;8.6;9.0"
+	TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;6.0;6.1;7.0;7.5;8.0;8.6;9.0"
 	#for cuda 11.8 we use cudnn 8.7
 	#which does not have single static libcudnn_static.a deliverable to link with
 	export USE_STATIC_CUDNN=0
 	#for cuda 11.8 include all dynamic loading libraries
 	DEPS_LIST=(/usr/local/cuda/lib64/libcudnn*.so.8 /usr/local/cuda-11.8/extras/CUPTI/lib64/libcupti.so.11.8)
     fi
+    if [[ -n "$OVERRIDE_TORCH_CUDA_ARCH_LIST" ]]; then
+        TORCH_CUDA_ARCH_LIST="$OVERRIDE_TORCH_CUDA_ARCH_LIST"
+    fi
+    export TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST"
     export NCCL_ROOT_DIR=/usr/local/cuda
     export USE_STATIC_NCCL=1  # links nccl statically (driven by tools/setup_helpers/nccl.py, some of the NCCL cmake files such as FindNCCL.cmake and gloo/FindNCCL.cmake)
 

diff --git a/conda/pytorch-nightly/meta.yaml b/conda/pytorch-nightly/meta.yaml
@@ -86,6 +86,7 @@ build:
     - USE_PYTORCH_METAL_EXPORT # [osx]
     - USE_COREML_DELEGATE # [osx]
     - _GLIBCXX_USE_CXX11_ABI # [unix]
+    - OVERRIDE_TORCH_CUDA_ARCH_LIST
 
 test:
  imports: