@@ -142,14 +142,22 @@ DEPS_SONAME=(
142142 " libcublasLt.so.11"
143143 " libgomp.so.1"
144144)
145- elif [[ $CUDA_VERSION == " 11.7" || $CUDA_VERSION == " 11.8 " ]]; then
145+ elif [[ $CUDA_VERSION == " 11.7" ]]; then
146146 export USE_STATIC_CUDNN=0
147147 # Try parallelizing nvcc as well
148148 export TORCH_NVCC_FLAGS=" -Xfatbin -compress-all --threads 2"
149149 DEPS_LIST=(
150+ " /usr/local/cuda/lib64/libcudart.so.11.0"
151+ " /usr/local/cuda/lib64/libnvToolsExt.so.1"
152+ " /usr/local/cuda/lib64/libnvrtc.so.11.2" # this is not a mistake for 11.7, it links to 11.7.50
153+ " /usr/local/cuda/lib64/libnvrtc-builtins.so.11.7"
150154 " $LIBGOMP_PATH "
151155 )
152156 DEPS_SONAME=(
157+ " libcudart.so.11.0"
158+ " libnvToolsExt.so.1"
159+ " libnvrtc.so.11.2"
160+ " libnvrtc-builtins.so.11.7"
153161 " libgomp.so.1"
154162 )
155163
@@ -165,10 +173,6 @@ elif [[ $CUDA_VERSION == "11.7" || $CUDA_VERSION == "11.8" ]]; then
165173 " /usr/local/cuda/lib64/libcudnn.so.8"
166174 " /usr/local/cuda/lib64/libcublas.so.11"
167175 " /usr/local/cuda/lib64/libcublasLt.so.11"
168- " /usr/local/cuda/lib64/libnvrtc.so.11.2" # this is not a mistake for 11.7, it links to 11.7.50
169- " /usr/local/cuda/lib64/libnvrtc-builtins.so.11.7"
170- " /usr/local/cuda/lib64/libcudart.so.11.0"
171- " /usr/local/cuda/lib64/libnvToolsExt.so.1"
172176 )
173177 DEPS_SONAME+=(
174178 " libcudnn_adv_infer.so.8"
@@ -182,31 +186,69 @@ elif [[ $CUDA_VERSION == "11.7" || $CUDA_VERSION == "11.8" ]]; then
182186 " libcublasLt.so.11"
183187 )
184188 else
185- echo " Using nvidia libs from pypi."
189+ echo " Using cudnn and cublas from pypi."
190+ CUDA_RPATHS=(
191+ ' $ORIGIN/../../nvidia/cublas/lib'
192+ ' $ORIGIN/../../nvidia/cudnn/lib'
193+ )
194+ CUDA_RPATHS=$( IFS=: ; echo " ${CUDA_RPATHS[*]} " )
195+ export C_SO_RPATH=$CUDA_RPATHS ' :$ORIGIN:$ORIGIN/lib'
196+ export LIB_SO_RPATH=$CUDA_RPATHS ' :$ORIGIN'
197+ export FORCE_RPATH=" --force-rpath"
198+ fi
199+ elif [[ $CUDA_VERSION == " 11.8" ]]; then
200+ export USE_STATIC_CUDNN=0
201+ # Try parallelizing nvcc as well
202+ export TORCH_NVCC_FLAGS=" -Xfatbin -compress-all --threads 2"
203+ DEPS_LIST=(
204+ " /usr/local/cuda/lib64/libcudart.so.11.0"
205+ " /usr/local/cuda/lib64/libnvToolsExt.so.1"
206+ " /usr/local/cuda/lib64/libnvrtc.so.11.2" # this is not a mistake for 11.8, it links to 11.8.89
207+ " /usr/local/cuda/lib64/libnvrtc-builtins.so.11.8"
208+ " $LIBGOMP_PATH "
209+ )
210+ DEPS_SONAME=(
211+ " libcudart.so.11.0"
212+ " libnvToolsExt.so.1"
213+ " libnvrtc.so.11.2"
214+ " libnvrtc-builtins.so.11.8"
215+ " libgomp.so.1"
216+ )
217+
218+ if [[ -z " $PYTORCH_EXTRA_INSTALL_REQUIREMENTS " ]]; then
219+ echo " Bundling with cudnn and cublas."
220+ DEPS_LIST+=(
221+ " /usr/local/cuda/lib64/libcudnn_adv_infer.so.8"
222+ " /usr/local/cuda/lib64/libcudnn_adv_train.so.8"
223+ " /usr/local/cuda/lib64/libcudnn_cnn_infer.so.8"
224+ " /usr/local/cuda/lib64/libcudnn_cnn_train.so.8"
225+ " /usr/local/cuda/lib64/libcudnn_ops_infer.so.8"
226+ " /usr/local/cuda/lib64/libcudnn_ops_train.so.8"
227+ " /usr/local/cuda/lib64/libcudnn.so.8"
228+ " /usr/local/cuda/lib64/libcublas.so.11"
229+ " /usr/local/cuda/lib64/libcublasLt.so.11"
230+ )
231+ DEPS_SONAME+=(
232+ " libcudnn_adv_infer.so.8"
233+ " libcudnn_adv_train.so.8"
234+ " libcudnn_cnn_infer.so.8"
235+ " libcudnn_cnn_train.so.8"
236+ " libcudnn_ops_infer.so.8"
237+ " libcudnn_ops_train.so.8"
238+ " libcudnn.so.8"
239+ " libcublas.so.11"
240+ " libcublasLt.so.11"
241+ )
242+ else
243+ echo " Using cudnn and cublas from pypi."
186244 CUDA_RPATHS=(
187245 ' $ORIGIN/../../nvidia/cublas/lib'
188- ' $ORIGIN/../../nvidia/cuda_cupti/lib'
189- ' $ORIGIN/../../nvidia/cuda_nvrtc/lib'
190- ' $ORIGIN/../../nvidia/cuda_runtime/lib'
191246 ' $ORIGIN/../../nvidia/cudnn/lib'
192- ' $ORIGIN/../../nvidia/cufft/lib'
193- ' $ORIGIN/../../nvidia/curand/lib'
194- ' $ORIGIN/../../nvidia/cusolver/lib'
195- ' $ORIGIN/../../nvidia/cusparse/lib'
196- ' $ORIGIN/../../nvidia/nccl/lib'
197- ' $ORIGIN/../../nvidia/nvtx/lib'
198247 )
199248 CUDA_RPATHS=$( IFS=: ; echo " ${CUDA_RPATHS[*]} " )
200249 export C_SO_RPATH=$CUDA_RPATHS ' :$ORIGIN:$ORIGIN/lib'
201250 export LIB_SO_RPATH=$CUDA_RPATHS ' :$ORIGIN'
202251 export FORCE_RPATH=" --force-rpath"
203- export USE_STATIC_NCCL=0
204- export USE_SYSTEM_NCCL=1
205- export ATEN_STATIC_CUDA=0
206- export USE_CUDA_STATIC_LINK=0
207- export USE_CUPTI_SO=1
208- export NCCL_INCLUDE_DIR=" /usr/local/cuda/include/"
209- export NCCL_LIB_DIR=" /usr/local/cuda/lib64/"
210252 fi
211253else
212254 echo " Unknown cuda version $CUDA_VERSION "
0 commit comments