@@ -142,22 +142,14 @@ DEPS_SONAME=(
142
142
" libcublasLt.so.11"
143
143
" libgomp.so.1"
144
144
)
145
- elif [[ $CUDA_VERSION == " 11.7" ]]; then
145
+ elif [[ $CUDA_VERSION == " 11.7" || $CUDA_VERSION == " 11.8 " ]]; then
146
146
export USE_STATIC_CUDNN=0
147
147
# Try parallelizing nvcc as well
148
148
export TORCH_NVCC_FLAGS=" -Xfatbin -compress-all --threads 2"
149
149
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"
154
150
" $LIBGOMP_PATH "
155
151
)
156
152
DEPS_SONAME=(
157
- " libcudart.so.11.0"
158
- " libnvToolsExt.so.1"
159
- " libnvrtc.so.11.2"
160
- " libnvrtc-builtins.so.11.7"
161
153
" libgomp.so.1"
162
154
)
163
155
@@ -173,6 +165,10 @@ elif [[ $CUDA_VERSION == "11.7" ]]; then
173
165
" /usr/local/cuda/lib64/libcudnn.so.8"
174
166
" /usr/local/cuda/lib64/libcublas.so.11"
175
167
" /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"
176
172
)
177
173
DEPS_SONAME+=(
178
174
" libcudnn_adv_infer.so.8"
@@ -186,69 +182,31 @@ elif [[ $CUDA_VERSION == "11.7" ]]; then
186
182
" libcublasLt.so.11"
187
183
)
188
184
else
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."
185
+ echo " Using nvidia libs from pypi."
244
186
CUDA_RPATHS=(
245
187
' $ORIGIN/../../nvidia/cublas/lib'
188
+ ' $ORIGIN/../../nvidia/cuda_cupti/lib'
189
+ ' $ORIGIN/../../nvidia/cuda_nvrtc/lib'
190
+ ' $ORIGIN/../../nvidia/cuda_runtime/lib'
246
191
' $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'
247
198
)
248
199
CUDA_RPATHS=$( IFS=: ; echo " ${CUDA_RPATHS[*]} " )
249
200
export C_SO_RPATH=$CUDA_RPATHS ' :$ORIGIN:$ORIGIN/lib'
250
201
export LIB_SO_RPATH=$CUDA_RPATHS ' :$ORIGIN'
251
202
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/"
252
210
fi
253
211
else
254
212
echo " Unknown cuda version $CUDA_VERSION "
0 commit comments