Skip to content

Commit 141fc72

Browse files
atalmanpytorchmergebot
authored andcommitted
[CD] CUDA 13.0 fix preload logic to include nvidia/cu13/lib/ (pytorch#163661)
Preload logic no longer works with CUDA 13.0 See the installation path: ``` ls /home/ubuntu/.venv/lib/python3.10/site-packages/nvidia/cu13/lib/ libcheckpoint.so libcudadevrt.a libcufft.so.12 libcufile_rdma.so.1 libcusolver.so.12 libnvJitLink.so.13 libnvperf_target.so libnvrtc.alt.so.13 libpcsamplingutil.so libcublas.so.13 libcudart.so.13 libcufftw.so.12 libcupti.so.13 libcusolverMg.so.12 libnvblas.so.13 libnvrtc-builtins.alt.so.13.0 libnvrtc.so.13 libcublasLt.so.13 libcudart_static.a libcufile.so.0 libcurand.so.10 libcusparse.so.12 libnvperf_host.so libnvrtc-builtins.so.13.0 libnvtx3interop.so.1 ls /home/ubuntu/.venv/lib/python3.10/site-packages/nvidia/ cu13 cudnn cusparselt nccl nvshmem ``` Test using script from : pytorch#162367 ``` Kernel test passed! ``` Pull Request resolved: pytorch#163661 Approved by: https://github.com/nWEIdia, https://github.com/tinglvv, https://github.com/Camyll
1 parent b66aa1a commit 141fc72

File tree

1 file changed

+15
-4
lines changed

1 file changed

+15
-4
lines changed

torch/__init__.py

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -283,10 +283,20 @@ def _load_dll_libraries() -> None:
283283

284284

285285
def _get_cuda_dep_paths(path: str, lib_folder: str, lib_name: str) -> list[str]:
286-
# Libraries can either be in path/nvidia/lib_folder/lib or path/lib_folder/lib
286+
# Libraries can either be in
287+
# path/nvidia/lib_folder/lib or
288+
# path/nvidia/cuXX/lib (since CUDA 13.0) or
289+
# path/lib_folder/lib
290+
from torch.version import cuda as cuda_version
291+
287292
nvidia_lib_paths = glob.glob(
288293
os.path.join(path, "nvidia", lib_folder, "lib", lib_name)
289294
)
295+
if cuda_version is not None:
296+
maj_cuda_version = cuda_version.split(".")[0]
297+
nvidia_lib_paths += glob.glob(
298+
os.path.join(path, "nvidia", f"cu{maj_cuda_version}", "lib", lib_name)
299+
)
290300
lib_paths = glob.glob(os.path.join(path, lib_folder, "lib", lib_name))
291301

292302
return nvidia_lib_paths + lib_paths
@@ -330,12 +340,13 @@ def _load_global_deps() -> None:
330340
try:
331341
with open("/proc/self/maps") as f:
332342
_maps = f.read()
333-
# libtorch_global_deps.so always depends in cudart, check if its installed via wheel
334-
if "nvidia/cuda_runtime/lib/libcudart.so" not in _maps:
343+
344+
# libtorch_global_deps.so always depends in cudart, check if its installed and loaded
345+
if "libcudart.so" not in _maps:
335346
return
336347
# If all above-mentioned conditions are met, preload nvrtc and nvjitlink
337-
# Please note that order are important for CUDA-11.8 , as nvjitlink does not exist there
338348
_preload_cuda_deps("cuda_nvrtc", "libnvrtc.so.*[0-9]")
349+
_preload_cuda_deps("cuda_nvrtc", "libnvrtc-builtins.so.*[0-9]")
339350
_preload_cuda_deps("nvjitlink", "libnvJitLink.so.*[0-9]")
340351
except Exception:
341352
pass

0 commit comments

Comments
 (0)