NCCL EP hard link makes `import transformer_engine` require `libcuda.so.1`

[cuichenx]

Describe the bug

When Transformer Engine is built for Hopper or newer with NCCL EP enabled,
libtransformer_engine.so has a direct ELF DT_NEEDED dependency on
libcuda.so.1. As a result, plain import transformer_engine fails on a host
that has the CUDA toolkit/runtime libraries but intentionally has no NVIDIA
driver or GPU device.

This happens before any Transformer Engine operation, CUDA execution, or NCCL
EP API is requested. It prevents CPU-only workflows that need to import TE and
construct TE modules with parameters on device="cpu", while never running a
TE forward pass. One downstream example is CPU-only checkpoint conversion:
using the normal TE-backed model spec preserves the exact parameter and
checkpoint schema expected by the GPU model, whereas substituting local
PyTorch modules can change fused parameter names and TE _extra_state entries.

The direct driver dependency was added by
#3127, commit
4955320121e500db98f1c08d8d90075c17d9469e. The NCCL EP CMake path
whole-archives libnccl_ep.a into libtransformer_engine.so and explicitly
links CUDA::cuda_driver. The accompanying comment says this ordering is
intended to make --as-needed record libcuda.so.1.

This regresses the loading property established by
#1240, which removed
the core library's direct CUDA driver link and used TE's existing indirect
driver-entry-point infrastructure instead. The hard link is still present on
TE main at 2d80391f77e542c06d0281688ddd17b6e34adec8.

Relevant source:

• NCCL EP link at the affected revision:
  

  TransformerEngine/transformer_engine/common/CMakeLists.txt

  Lines 440 to 510 in 4329ff8

  	# -- NCCL EP (on by default, HT mode only) ---------------------------------
  	# Set -DNVTE_WITH_NCCL_EP=OFF (or NVTE_WITH_NCCL_EP=0 in setup.py) to
  	# skip NCCL EP entirely - useful on older images whose system NCCL is below
  	# the 2.30.4 EP minimum.
  	option(NVTE_WITH_NCCL_EP "Build NCCL EP into libtransformer_engine.so" ON)
  	if(NVTE_WITH_NCCL_EP)
  	# SM>=90 and NCCL>=2.30.4 are gated at runtime in EPBackend::initialize.
  	# -- NCCL EP headers --------------------------------------------------------
  	# Headers + libs are produced by the in-tree 3rdparty/nccl submodule build
  	# (auto-built by setup.py via build_nccl_ep_submodule).
  	set(NCCL_EP_SUBMODULE_ROOT
  	"${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl")
  	set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/include")
  	if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h")
  	message(FATAL_ERROR
  	"NCCL EP header not found at ${NCCL_EP_INCLUDE_DIR}/nccl_ep.h. "
  	"Run `git submodule update --init --recursive` and rebuild TE.")
  	endif()
  	message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}")
  	# -- libnccl_ep.a -----------------------------------------------------------
  	# Statically linked into libtransformer_engine.so. EPBackend::initialize checks
  	# NCCL >= 2.30.4 before any nccl_ep call, so the newer NCCL symbols nccl_ep
  	# imports stay unresolved (and harmless) under default ELF lazy binding when
  	# the gate trips. LD_BIND_NOW environments lose this property.
  	set(NCCL_EP_LIB_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/lib")
  	find_file(NCCL_EP_LIB
  	NAMES libnccl_ep.a
  	HINTS ${NCCL_EP_LIB_DIR}
  	NO_DEFAULT_PATH
  	REQUIRED)
  	# -- NCCL core: nccl.h + libnccl.so -----------------------------------------
  	# setup.py passes -DNCCL_INCLUDE_DIR; standalone CMake falls back to probing
  	# well-known NCCL install prefixes.
  	find_path(NCCL_INCLUDE_DIR nccl.h
  	HINTS /opt/nvidia/nccl/include /usr/local/nccl/include)
  	if(NOT NCCL_INCLUDE_DIR)
  	message(FATAL_ERROR
  	"nccl.h not found. Pass -DNCCL_INCLUDE_DIR=<prefix>/include.")
  	endif()
  	if(NOT NCCL_LIB)
  	find_library(NCCL_LIB
  	NAMES nccl libnccl
  	PATH_SUFFIXES lib lib64
  	REQUIRED)
  	endif()
  	target_include_directories(transformer_engine PRIVATE
  	${NCCL_EP_INCLUDE_DIR}
  	${NCCL_INCLUDE_DIR})
  	# libnccl.so direct symbols (ncclGetVersion etc.) come from libnccl_ep.a's
  	# DT_NEEDED chain plus this TU's own references. CUDA::cuda_driver must follow
  	# the static archive on the link line so --as-needed records libcuda.so.1.
  	target_link_libraries(transformer_engine PUBLIC ${NCCL_LIB})
  	target_link_libraries(transformer_engine PRIVATE
  	-Wl,--whole-archive ${NCCL_EP_LIB} -Wl,--no-whole-archive
  	CUDA::cuda_driver)
  	target_sources(transformer_engine PRIVATE
  	ep/ep_backend.cpp
  	ep/ep_api.cpp)
  	target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_NCCL_EP)
  	message(STATUS "NCCL EP enabled (static link): ${NCCL_EP_LIB}")
  	message(STATUS "NCCL EP include: ${NCCL_EP_INCLUDE_DIR}")
  	else()
  	# NCCL EP off: ep_api.cpp's #else branch exports throwing nvte_ep_* stubs.
  	target_sources(transformer_engine PRIVATE ep/ep_api.cpp)
  	message(STATUS "NCCL EP disabled (NVTE_WITH_NCCL_EP=OFF) - using nvte_ep_* stubs")

• Import-time core-library load:
  

  TransformerEngine/transformer_engine/common/__init__.py

  Lines 357 to 382 in 4329ff8

  	@functools.lru_cache(maxsize=None)
  	def _load_core_library():
  	"""Load shared library with Transformer Engine C extensions"""
  	return ctypes.CDLL(_get_shared_object_file("core"), mode=ctypes.RTLD_GLOBAL | os.RTLD_LAZY)
  	if "NVTE_PROJECT_BUILDING" not in os.environ or bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))):
  	sanity_checks_for_pypi_installation()
  	# `_load_cuda_library` is used for packages that must be loaded
  	# during runtime. Both system and pypi packages are searched
  	# 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
  	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)
  	_TE_LIB_CTYPES = _load_core_library()

• Current-main NCCL EP link:
  

  TransformerEngine/transformer_engine/common/CMakeLists.txt

  Lines 496 to 514 in 2d80391

  	# libnccl.so direct symbols (ncclGetVersion etc.) come from libnccl_ep.a's
  	# DT_NEEDED chain plus this TU's own references. CUDA::cuda_driver must follow
  	# the static archive on the link line so --as-needed records libcuda.so.1.
  	target_link_libraries(transformer_engine PUBLIC ${NCCL_LIB})
  	target_link_libraries(transformer_engine PRIVATE
  	-Wl,--whole-archive ${NCCL_EP_LIB} -Wl,--no-whole-archive
  	CUDA::cuda_driver)
  	target_sources(transformer_engine PRIVATE
  	ep/ep_backend.cpp
  	ep/ep_api.cpp)
  	target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_NCCL_EP)
  	message(STATUS "NCCL EP enabled (static link): ${NCCL_EP_LIB}")
  	message(STATUS "NCCL EP include: ${NCCL_EP_INCLUDE_DIR}")
  	else()
  	# NCCL EP off: ep_api.cpp's #else branch exports throwing nvte_ep_* stubs.
  	target_sources(transformer_engine PRIVATE ep/ep_api.cpp)
  	message(STATUS "NCCL EP disabled (NVTE_WITH_NCCL_EP=OFF) - using nvte_ep_* stubs")

Steps/Code to reproduce bug

Use a TE build that targets SM90 or newer and has NCCL EP enabled. Run it on a
Linux host/container with the required CUDA toolkit libraries installed, but
with no libcuda.so.1 and no /dev/nvidia* devices.

$ ls /dev/nvidia*
ls: cannot access '/dev/nvidia*': No such file or directory

$ ldconfig -p | grep -E 'libcuda\.so|libnvidia-ml'
# no output

$ python - <<'PY'
import torch

print("CUDA initialized before TE import:", torch.cuda.is_initialized())
import transformer_engine
PY
CUDA initialized before TE import: False
Traceback (most recent call last):
  ...
  File "transformer_engine/common/__init__.py", line 382, in <module>
    _TE_LIB_CTYPES = _load_core_library()
  File "transformer_engine/common/__init__.py", line 360, in _load_core_library
    return ctypes.CDLL(..., mode=ctypes.RTLD_GLOBAL | os.RTLD_LAZY)
OSError: libcuda.so.1: cannot open shared object file: No such file or directory

The binary dependency is visible without importing TE:

$ readelf -d /path/to/libtransformer_engine.so | grep NEEDED | grep libcuda
 0x0000000000000001 (NEEDED) Shared library: [libcuda.so.1]

RTLD_LAZY does not help because the ELF loader must resolve direct
DT_NEEDED libraries when libtransformer_engine.so is loaded.

The observed PyTorch extension did not itself have a direct libcuda.so.1
entry; the failing dependency was on the TE core library.

Expected behavior

On a system where TE's non-driver shared-library dependencies are present,
importing transformer_engine and transformer_engine.pytorch should not
require the NVIDIA driver merely because NCCL EP was included at build time.

Constructing TE modules with parameters on device="cpu" should remain
possible without initializing or using CUDA. CUDA execution and NCCL EP may
still require a driver and GPU, and should fail with a clear error only when
those capabilities are actually requested.

Suggested fix

Prefer making the NCCL EP backend an optional, lazily loaded component instead
of whole-archiving it into the always-loaded TE core library:

1. Keep the public nvte_ep_* C API in libtransformer_engine.so as thin
   forwarding entry points.
2. Put ep_backend.cpp, libnccl_ep.a, and their NCCL/CUDA-driver link
   dependencies in a separate shared object, for example
   libtransformer_engine_nccl_ep.so.
3. Load that backend with dlopen and resolve a versioned function table on
   the first nvte_ep_initialize() call, not during Python package import.
4. If the backend or driver is unavailable, raise an actionable NCCL EP error
   at that point. Other TE imports and CPU parameter construction should remain
   usable.
5. Preserve the current throwing stubs when TE is built with
   NVTE_WITH_NCCL_EP=0.

An alternative is to remove direct CUDA driver references from the NCCL EP
objects and route them through TE's existing cudaGetDriverEntryPoint-based
loader. The key requirement is that the always-loaded
libtransformer_engine.so no longer records libcuda.so.1 solely because the
optional NCCL EP backend was compiled.

Current workaround

Building TE with NVTE_WITH_NCCL_EP=0 selects the existing throwing
nvte_ep_* stubs and avoids the NCCL EP CMake link branch. This is suitable
for a CPU-only conversion image, but it disables NCCL EP for GPU/MoE workloads
and therefore is not a general solution for a shared training image.

Adding a CUDA stub library to LD_LIBRARY_PATH is not a safe workaround. It
masks the import-time dependency and defers failure until an accidental driver
call.

Proposed acceptance tests

1. Build for SM90+ with NVTE_WITH_NCCL_EP=1 and verify that
   libtransformer_engine.so has no DT_NEEDED entry for libcuda.so.1.
   A separately loaded NCCL EP backend may retain that dependency.

2. In a Linux container with CUDA toolkit/runtime libraries but no NVIDIA
   driver or /dev/nvidia*, verify:

   import torch
   import transformer_engine
   import transformer_engine.pytorch as te
   
   assert not torch.cuda.is_initialized()
   module = te.Linear(16, 16, device="cpu")
   assert module.weight.device.type == "cpu"
   assert not torch.cuda.is_initialized()

3. Verify that requesting NCCL EP without a usable driver/backend produces a
   clear runtime exception at the EP API boundary rather than breaking package
   import.

4. Run the existing NCCL EP tests on H100 or newer and confirm there is no
   functional or performance regression after the backend is loaded.

5. Retain coverage for NVTE_WITH_NCCL_EP=0 and its existing stubs.

Environment overview

• Environment location: Linux container on a genuinely driverless CPU-only
  host
• Installation: prebuilt container package from the exact TE source revision
  below
• Transformer Engine: 2.17.1+4329ff84
• Transformer Engine source: 4329ff84bfbdaa778a33cba02a15fb0807c64689
• Python: 3.12.3
• PyTorch: 2.13.0a0+8145d630e8.nv26.06
• CUDA toolkit/runtime: 13.3
• NVIDIA driver library: absent
• GPU devices: none

Device details

No GPU is intentionally present for the failing import. The TE binary was
built for Hopper-or-newer targets, which enables the NCCL EP build path by
default.

Additional context

This report does not request CPU execution of TE kernels. The required contract
is narrower: importing TE and constructing its parameter schema on CPU should
not load the NVIDIA driver. Normal TE forward execution and NCCL EP remain GPU
operations.

[timmoon10]

I wonder if we should consider this as a NCCL EP bug. It's common practice to avoid directly linking to libcuda.so because it can differ between build-time (a stub included with the CUDA Toolkit to allow compilation on non-GPU nodes) and run-time (environment-dependent driver, possibly very old). TE has a special wrapper for calling CUDA driver APIs (see here), as does PyTorch (see here). As it is, NCCL EP cannot really exist as a standalone library, but must rely on Python to manually load the right libcuda.so.