[bvolpato]

Description

Fix two ways that the runtime CUDA library search breaks environments where a pip-installed framework supplies the CUDA libraries.

A host-wide cuDNN replaces the cuDNN that the framework needs. The CDLL("libcudnn.so") fallback is described as an LD_LIBRARY_PATH search, but it also searches the dynamic linker's default locations, and it runs before the Python environment. PyPI torch pins its own nvidia-cudnn-cu13, so on a host with an older system cuDNN it fails with:

RuntimeError: cuDNN version incompatibility: PyTorch was compiled against (9, 26, 0) but found runtime version (9, 20, 0).

With JAX, the same search maps two copies of cuDNN into one process.

The import also requires libcurand. PyTorch 2.15 removed curand from its cuda-toolkit extras, and JAX does not install it, so on a host without the CUDA toolkit the import fails with curand shared object not found. Transformer Engine only uses cuRAND's device-side headers. The core library does not link libcurand and calls none of its host functions.

Related to #2373 and #2297, which made LD_LIBRARY_PATH take priority over Python packages. This PR keeps that priority.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Search for each library in this order: <LIBRARY>_HOME, <LIBRARY>_PATH, CUDA_HOME, CUDA_PATH, /usr/local/cuda, LD_LIBRARY_PATH, the Python environment, and then the dynamic linker's default locations.
• Read LD_LIBRARY_PATH directly, so that it also finds versioned libraries such as libcudnn.so.9. The pip packages have no unversioned symlink, so the previous lookup could not find them there. An entry's unversioned name still takes priority when it exists. $ORIGIN expands to the executable's directory, and an entry with another token, such as $LIB, goes to dlopen as a path. An entry that does not have the library does not affect the search.
• Stop loading cuRAND, and decide whether the toolkit is present from NVRTC alone. When NVRTC comes from a Python package, load cuBLAS and cudart from Python packages only if they are there, and leave the rest to the dynamic linker.
• Document the search order in the troubleshooting guide.

PyTorch preloads the cuDNN from its own Python package, so when the environment selects the system cuDNN, PyTorch's copy is mapped too and PyTorch rejects the older version. That behavior is not new.

One behavior changes on purpose. If a framework uses a system cuDNN and an unrelated pip cuDNN package is also installed, Transformer Engine now loads the pip package unless the environment selects a location. CUDNN_PATH or LD_LIBRARY_PATH selects the system copy.

Before and after

These runs used one NVIDIA B200 with system cuDNN 9.20 and PyTorch 2.15.0+cu132, which installs cuDNN 9.26.

Environment	Before	After
Default, importing either PyTorch or Transformer Engine first	cuDNN version incompatibility	Loads pip cuDNN 9.26
LD_LIBRARY_PATH set to pip cuDNN	cuDNN version incompatibility	Loads pip cuDNN 9.26
CUDNN_PATH set to pip cuDNN	Loads pip cuDNN 9.26	Same
LD_LIBRARY_PATH or CUDNN_PATH set to system cuDNN	Loads system cuDNN 9.20, which PyTorch rejects	Same
No pip cuDNN, importing only transformer_engine.common	Loads system cuDNN 9.20	Same
JAX 0.12, default	Maps both pip and system cuDNN	Maps only pip cuDNN
No CUDA toolkit and no system cuDNN	curand shared object not found	Imports and runs from pip libraries
NVRTC from pip, with cuBLAS and cudart only from the system	curand shared object not found	Imports with the system cuBLAS and cudart
LD_LIBRARY_PATH entry with $ORIGIN or $LIB that selects system cuDNN	Loads system cuDNN 9.20	Same
LD_LIBRARY_PATH entry with $ORIGIN or $LIB that has no cuDNN	cuDNN version incompatibility	Loads pip cuDNN 9.26
LD_LIBRARY_PATH entry whose libcudnn.so selects one of two versions	Follows libcudnn.so	Same

Validation

The source builds targeted 100a with CUDA 13.2, PyTorch 2.15.0+cu132, cuDNN 9.26, FlashAttention 2.8.4, and JAX 0.12.0.1. The base was b0219803.

• The new tests in test_sanity.py fail when the search order or the cuRAND load is reverted, or when the LD_LIBRARY_PATH name and token handling is removed, and they pass with this change.
• The token, unversioned symlink, and pip NVRTC cases in the table ran with real libraries on B200. The pip NVRTC case has no unit test, because it depends on which libraries the host has.
• The L0 PyTorch commands and the debug launcher ran on this PR's final code alone and on an untouched tree of the same commit, with each job in its own directory. No test failed on only one tree. The failures on both trees are test_torch_compile.py ([PyTorch] torch.compile fails on FusedAttnBackend IntEnum comparisons #3716), test_dot_product_attention_score_mod ([PyTorch] Keep Flex Attention score_mod graphs on the cuDNN backend #3796), and FP32 numerics cases ([CI] Disable TF32 for the L0 PyTorch numerics tests #3797).
• For JAX, test_layer.py and test_fused_attn_score_mod.py found no failures in the default environment, with the system cuDNN present, before the last commit. test_fused_attn.py found no failures before the last two commits. The last commit only changes LD_LIBRARY_PATH entries that contain tokens, and JAX import was rechecked after it.
• For the no-toolkit case, the toolkit and system cuDNN were moved aside in a disposable container, and the dynamic linker had no CUDA entries left. Transformer Engine took its no-toolkit path, and every loaded CUDA library came from the virtual environment. test_sanity.py, test_numerics.py, tests/pytorch/nvfp4, tests/pytorch/mxfp8, and test_quantized_tensor.py found no failures before the last commit, with NVIDIA_TF32_OVERRIDE=0 set. JAX imported and ran a basic operation. The no-toolkit import and a forward pass were rechecked on the final code.

Pylint reported no issues for the changed module, and the repository hooks passed. A100, H100, aarch64, and NGC images, where the framework uses the system cuDNN, were not tested. Release wheels were not rebuilt with this change.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; the previous token-path finding is fixed.

Summary

The PR puts Python-provided CUDA libraries ahead of default system locations while keeping explicit environment choices first.

• Transformer Engine checks selected CUDA paths and Python packages before host defaults.
• Transformer Engine imports without requiring the cuRAND shared library.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Load CUDA library] --> B[Explicit library and toolkit paths]
  B -->|Not found| C[Selected LD_LIBRARY_PATH entries]
  C -->|Not found| D[Python packages]
  D -->|Not found| E[Default linker search]
  E -->|Not found| F[Raise missing-library error]

Loading

_{Reviews (3) · Last reviewed commit: "[Common] Limit LD_LIBRARY_PATH token han..." · Reviewed by Greptile}