diff --git a/docs/troubleshooting.rst b/docs/troubleshooting.rst index cd5bdd1838..abc136f5b7 100644 --- a/docs/troubleshooting.rst +++ b/docs/troubleshooting.rst @@ -69,6 +69,10 @@ UV and virtual environments export CUDNN_HOME=$CUDNN_PATH export LD_LIBRARY_PATH=$CUDNN_PATH/lib:$LD_LIBRARY_PATH + At runtime, Transformer Engine searches for each CUDA library in ``_HOME``, + ``_PATH``, ``CUDA_HOME``, ``CUDA_PATH``, ``/usr/local/cuda``, and ``LD_LIBRARY_PATH``. + It then searches the Python environment, and finally the dynamic linker's default locations. + 3. **Building wheels** Use ``uv build --wheel --no-build-isolation -v`` when building the wheel and diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 2627de023a..d73a0891e9 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -8,8 +8,13 @@ import pytest import os import copy +import ctypes +import subprocess +import sys +import textwrap import transformer_engine +import transformer_engine.common as te_common import transformer_engine.pytorch as te from transformer_engine.pytorch.quantization import FP8GlobalStateManager from transformer_engine.pytorch.utils import ( @@ -1651,3 +1656,136 @@ def check_weights(): with autocast(enabled=with_quantization, recipe=quantization_recipe): y = module(x, **kwargs) check_weights() + + +_SEARCH_LIB = "tesearchordertest" + + +def _run_cuda_library_search(monkeypatch, tmp_path, ld_library_path, linker_finds_library=True): + """Run the CUDA library search with a fake loader. Return its result and the loaded names.""" + loaded = [] + + def fake_cdll(name, mode=0): + missing = not os.path.exists(name) if "/" in name else not linker_finds_library + if missing: + raise OSError(f"{name}: cannot open shared object file") + loaded.append(name) + return object() + + monkeypatch.setattr(ctypes, "CDLL", fake_cdll) + monkeypatch.setattr(te_common.sysconfig, "get_path", lambda _: str(tmp_path / "python")) + for variable in ("CUDA_HOME", "CUDA_PATH", f"{_SEARCH_LIB.upper()}_HOME"): + monkeypatch.delenv(variable, raising=False) + monkeypatch.setenv("LD_LIBRARY_PATH", ld_library_path) + monkeypatch.setenv(f"{_SEARCH_LIB.upper()}_PATH", str(tmp_path / "lib_path")) + + te_common._load_cuda_library_from_system.cache_clear() + te_common._load_cuda_library_from_python.cache_clear() + te_common._load_cuda_library_from_linker_default.cache_clear() + te_common._load_cuda_library.cache_clear() + + found_in_system, _ = te_common._load_cuda_library(_SEARCH_LIB) + return found_in_system, loaded + + +def _create_libraries(directory, *suffixes): + directory.mkdir(parents=True) + for suffix in suffixes: + (directory / f"lib{_SEARCH_LIB}.so{suffix}").touch() + + +@pytest.mark.parametrize( + "available, expected", + [ + (("linker",), "linker"), + (("linker", "python"), "python"), + (("linker", "python", "ld_library_path"), "ld_library_path"), + (("linker", "python", "ld_library_path", "lib_path"), "lib_path"), + ], +) +def test_cuda_library_search_order(monkeypatch, tmp_path, available, expected): + """Selected locations outrank the Python environment, which outranks a host-wide library.""" + locations = { + "python": tmp_path / "python" / "nvidia" / _SEARCH_LIB / "lib", + "ld_library_path": tmp_path / "ld_library_path", + "lib_path": tmp_path / "lib_path" / "lib", + } + for name, location in locations.items(): + if name in available: + _create_libraries(location, ".1") + + found_in_system, loaded = _run_cuda_library_search( + monkeypatch, + tmp_path, + str(locations["ld_library_path"]), + linker_finds_library="linker" in available, + ) + + assert found_in_system == (expected != "python") + if expected == "linker": + assert loaded == [f"lib{_SEARCH_LIB}.so"] + else: + assert loaded == [str(locations[expected] / f"lib{_SEARCH_LIB}.so.1")] + + +@pytest.mark.parametrize( + "suffixes, expected_suffix", + [(("", ".1", ".2"), ""), ((".1", ".2"), ".2")], +) +def test_cuda_library_ld_library_path_name_choice(monkeypatch, tmp_path, suffixes, expected_suffix): + """An LD_LIBRARY_PATH entry follows the unversioned name, then the newest versioned name.""" + directory = tmp_path / "ld_library_path" + _create_libraries(directory, *suffixes) + + _, loaded = _run_cuda_library_search(monkeypatch, tmp_path, str(directory)) + + assert loaded == [str(directory / f"lib{_SEARCH_LIB}.so{expected_suffix}")] + + +@pytest.mark.parametrize("entry", ["$ORIGIN/lib", "${ORIGIN}/lib"]) +def test_cuda_library_ld_library_path_expands_origin(monkeypatch, tmp_path, entry): + """A $ORIGIN entry that has the library outranks the Python environment.""" + _create_libraries(tmp_path / "python" / "nvidia" / _SEARCH_LIB / "lib", ".1") + _create_libraries(tmp_path / "bin" / "lib", "") + monkeypatch.setattr(sys, "executable", str(tmp_path / "bin" / "python")) + + found_in_system, loaded = _run_cuda_library_search(monkeypatch, tmp_path, entry) + + assert found_in_system + assert loaded == [str(tmp_path / "bin" / "lib" / f"lib{_SEARCH_LIB}.so")] + + +def test_cuda_library_token_entry_without_library_keeps_python_priority(monkeypatch, tmp_path): + """A token entry that lacks the library must not let a host-wide library win.""" + location = tmp_path / "python" / "nvidia" / _SEARCH_LIB / "lib" + _create_libraries(location, ".1") + + found_in_system, loaded = _run_cuda_library_search( + monkeypatch, tmp_path, "$ORIGIN/missing:$LIB/missing" + ) + + assert not found_in_system + assert loaded == [str(location / f"lib{_SEARCH_LIB}.so.1")] + + +def test_import_does_not_require_curand_library(): + """TE only uses cuRAND's device headers, and torch no longer always installs its library.""" + probe = textwrap.dedent( + """ + import ctypes + + cdll = ctypes.CDLL + + + def cdll_without_curand(name, *args, **kwargs): + if name is not None and "curand" in str(name): + raise OSError(f"{name}: cannot open shared object file") + return cdll(name, *args, **kwargs) + + + ctypes.CDLL = cdll_without_curand + + import transformer_engine.pytorch + """ + ) + subprocess.run([sys.executable, "-c", probe], check=True) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 6237c6732b..cbd6a373cf 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -345,12 +345,11 @@ def _nvidia_cudart_include_dir(cuda_major_version: int) -> str: @functools.lru_cache(maxsize=None) -def _load_cuda_library_from_python(lib_name: str, strict: bool = False): +def _load_cuda_library_from_python(lib_name: str): """ Attempts to load shared object file installed via python packages. `lib_name` : Name of package as found in the `nvidia` dir in python environment. - `strict` : If set to `True`, throw an error if lib is not found. """ ext = _get_sys_extension() @@ -382,20 +381,25 @@ def _load_cuda_library_from_python(lib_name: str, strict: bool = False): for so_path in so_paths: ctypes_handles.append(ctypes.CDLL(so_path, mode=ctypes.RTLD_GLOBAL)) - if strict and not path_found: - raise RuntimeError(f"{lib_name} shared object not found.") - return path_found, ctypes_handles +def _newest_cuda_library(pattern: str) -> Optional[str]: + """Return the highest-sorted non-stub library that matches a glob pattern.""" + libs = [lib for lib in glob.glob(pattern, recursive=True) if "stub" not in lib] + return max(libs, key=os.path.basename, default=None) + + @functools.lru_cache(maxsize=None) def _load_cuda_library_from_system(lib_name: str): """ - Attempts to load shared object file installed via system/cuda-toolkit. + Attempts to load shared object file from a location selected by the environment or toolkit. `lib_name`: Name of library to load without extension or `lib` prefix. """ + file_name = f"lib{lib_name}{_get_sys_extension()}" + # Where to look for the shared lib in decreasing order of preference. paths = ( os.environ.get(f"{lib_name.upper()}_HOME"), @@ -408,16 +412,41 @@ def _load_cuda_library_from_system(lib_name: str): for path in paths: if path is None: continue - libs = glob.glob(f"{path}/**/lib{lib_name}{_get_sys_extension()}*", recursive=True) - libs = [lib for lib in libs if "stub" not in lib] - libs.sort(reverse=True, key=os.path.basename) - if libs: - return True, ctypes.CDLL(libs[0], mode=ctypes.RTLD_GLOBAL) + lib = _newest_cuda_library(f"{path}/**/{file_name}*") + if lib: + return True, ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL) + + origin = os.path.dirname(os.path.realpath(sys.executable)) + for path in os.environ.get("LD_LIBRARY_PATH", "").split(os.pathsep): + if not path: + continue + # In this variable, $ORIGIN is the executable's directory. dlopen would use the caller's. + path = path.replace("${ORIGIN}", origin).replace("$ORIGIN", origin) + exact = os.path.join(path, file_name) + # dlopen expands the other tokens, such as $LIB, but a file search cannot. + if "$" in path: + try: + return True, ctypes.CDLL(exact, mode=ctypes.RTLD_GLOBAL) + except OSError: + continue + # Follow the unversioned name as the linker does. Python packages only have versioned names. + lib = exact if os.path.exists(exact) else _newest_cuda_library(f"{exact}*") + if lib: + return True, ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL) + + return False, None + + +@functools.lru_cache(maxsize=None) +def _load_cuda_library_from_linker_default(lib_name: str): + """ + Attempts to load shared object file through the dynamic linker's default search. + + `lib_name`: Name of library to load without extension or `lib` prefix. + """ - # Search in LD_LIBRARY_PATH. try: - _lib_handle = ctypes.CDLL(f"lib{lib_name}{_get_sys_extension()}", mode=ctypes.RTLD_GLOBAL) - return True, _lib_handle + return True, ctypes.CDLL(f"lib{lib_name}{_get_sys_extension()}", mode=ctypes.RTLD_GLOBAL) except OSError: return False, None @@ -426,11 +455,11 @@ def _load_cuda_library_from_system(lib_name: str): def _load_cuda_library(lib_name: str): """ Load given shared library. - Prioritize loading from system/toolkit - before checking python packages. + Prioritize locations selected by the environment or toolkit, then python packages, + then the dynamic linker's default search. """ - # Attempt to locate library in system. + # Attempt to locate library in a selected location. found, handle = _load_cuda_library_from_system(lib_name) if found: return True, handle @@ -440,6 +469,12 @@ def _load_cuda_library(lib_name: str): if found: return False, handle + # A host-wide library comes last: pip-installed frameworks depend on the version in their own + # Python environment, and loading another copy globally would replace it in this process. + found, handle = _load_cuda_library_from_linker_default(lib_name) + if found: + return True, handle + raise RuntimeError(f"{lib_name} shared object not found.") @@ -503,16 +538,14 @@ def _load_core_library(): # and an error is thrown if not found. _, _CUDNN_LIB_CTYPES = _load_cuda_library("cudnn") system_nvrtc, _NVRTC_LIB_CTYPES = _load_cuda_library("nvrtc") - system_curand, _CURAND_LIB_CTYPES = _load_cuda_library("curand") # This additional step is necessary to be able to install TE wheels # and import TE (without any guards) in an environment where the cuda # toolkit might be absent without being guarded - load_libs_for_no_ctk = not system_nvrtc and not system_curand + load_libs_for_no_ctk = not system_nvrtc if load_libs_for_no_ctk: - _CUBLAS_LIB_CTYPES = _load_cuda_library_from_python("cublas", strict=True) - _CUDART_LIB_CTYPES = _load_cuda_library_from_python("cudart", strict=True) - _CUDNN_ALL_LIB_CTYPES = _load_cuda_library_from_python("cudnn", strict=True) + _CUBLAS_LIB_CTYPES = _load_cuda_library_from_python("cublas") + _CUDART_LIB_CTYPES = _load_cuda_library_from_python("cudart") # Prepare CuTeDSL backend for on-demand compilation via TVM-FFI if user enables it. if os.environ.get("NVTE_ENABLE_CUTEDSL_BACKEND", "0") != "0":