diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index 30ea83a55b..5a063723fa 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -6,6 +6,7 @@ import copy import os +import shutil import subprocess import sys import sysconfig @@ -171,6 +172,17 @@ def run(self) -> None: self.copy_file(ext, target_dir) os.remove(ext) + nccl_ep_dir = Path(self.build_lib) / "nccl_ep" + if nccl_ep_dir.is_dir(): + target_nccl_ep_dir = target_dir / "nccl_ep" + if target_nccl_ep_dir.exists(): + shutil.rmtree(target_nccl_ep_dir) + shutil.copytree( + nccl_ep_dir, + target_nccl_ep_dir, + ) + shutil.rmtree(nccl_ep_dir) + def build_extensions(self): # For core lib + JAX install, fix build_ext from pybind11.setup_helpers # to handle CUDA files correctly. diff --git a/qa/L0_jax_wheel/test.sh b/qa/L0_jax_wheel/test.sh index fa50a6de68..1d83946918 100644 --- a/qa/L0_jax_wheel/test.sh +++ b/qa/L0_jax_wheel/test.sh @@ -30,6 +30,12 @@ WHL_BASE="transformer_engine-${VERSION}" # Core wheel. NVTE_RELEASE_BUILD=1 pip3 wheel --no-build-isolation -vvv --wheel-dir ./dist . || error_exit "Failed to setup bdist_wheel" wheel unpack dist/${WHL_BASE}-* || error_exit "Failed to unpack dist/${WHL_BASE}-*.whl" +if python3 -c "from build_tools.utils import nccl_ep_enabled; raise SystemExit(not nccl_ep_enabled())"; then + test -f "${WHL_BASE}/transformer_engine/wheel_lib/libnccl_ep.so" || error_exit "Core wheel is missing libnccl_ep.so" + test -f "${WHL_BASE}/transformer_engine/wheel_lib/nccl_ep/include/nccl_ep.h" || error_exit "Core wheel is missing nccl_ep.h" + test -d "${WHL_BASE}/transformer_engine/wheel_lib/nccl_ep/include/nccl_ep" || error_exit "Core wheel is missing NCCL EP JIT headers" + python3 "$TE_PATH/qa/check_nccl_ep_headers.py" "${WHL_BASE}/transformer_engine/wheel_lib/nccl_ep/include" || error_exit "Core wheel has incomplete NCCL EP JIT headers" +fi sed -i "s/Name: transformer-engine/Name: transformer-engine-cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" sed -i "s/Name: transformer_engine/Name: transformer_engine_cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" mv "${WHL_BASE}/${WHL_BASE}.dist-info" "${WHL_BASE}/transformer_engine_cu12-${VERSION}.dist-info" || error_exit "Failed to move ${WHL_BASE}.dist-info to transformer_engine_cu12-${VERSION}.dist-info" diff --git a/qa/L0_pytorch_wheel/test.sh b/qa/L0_pytorch_wheel/test.sh index fe4aab456e..07d651d6df 100644 --- a/qa/L0_pytorch_wheel/test.sh +++ b/qa/L0_pytorch_wheel/test.sh @@ -29,6 +29,12 @@ WHL_BASE="transformer_engine-${VERSION}" # Core wheel. NVTE_RELEASE_BUILD=1 pip3 wheel --no-build-isolation -vvv --wheel-dir ./dist . || error_exit "Failed to setup bdist_wheel" python3 -m wheel unpack dist/${WHL_BASE}-* || error_exit "Failed to unpack dist/${WHL_BASE}-*.whl" +if python3 -c "from build_tools.utils import nccl_ep_enabled; raise SystemExit(not nccl_ep_enabled())"; then + test -f "${WHL_BASE}/transformer_engine/wheel_lib/libnccl_ep.so" || error_exit "Core wheel is missing libnccl_ep.so" + test -f "${WHL_BASE}/transformer_engine/wheel_lib/nccl_ep/include/nccl_ep.h" || error_exit "Core wheel is missing nccl_ep.h" + test -d "${WHL_BASE}/transformer_engine/wheel_lib/nccl_ep/include/nccl_ep" || error_exit "Core wheel is missing NCCL EP JIT headers" + python3 "$TE_PATH/qa/check_nccl_ep_headers.py" "${WHL_BASE}/transformer_engine/wheel_lib/nccl_ep/include" || error_exit "Core wheel has incomplete NCCL EP JIT headers" +fi sed -i "s/Name: transformer-engine/Name: transformer-engine-cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" sed -i "s/Name: transformer_engine/Name: transformer_engine_cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" mv "${WHL_BASE}/${WHL_BASE}.dist-info" "${WHL_BASE}/transformer_engine_cu12-${VERSION}.dist-info" || error_exit "Failed to move ${WHL_BASE}.dist-info to transformer_engine_cu12-${VERSION}.dist-info" diff --git a/qa/check_nccl_ep_headers.py b/qa/check_nccl_ep_headers.py new file mode 100644 index 0000000000..e5fad5afdf --- /dev/null +++ b/qa/check_nccl_ep_headers.py @@ -0,0 +1,53 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Verify that packaged NCCL EP JIT headers contain all local quoted includes.""" + +import re +import sys +from pathlib import Path + + +_INCLUDE_PATTERN = re.compile(r'^\s*#\s*include\s*"([^"]+)"') +_EXTERNAL_HEADERS = {"nccl.h", "nccl_device.h"} +_HEADER_SUFFIXES = {".cuh", ".h", ".hh", ".hpp", ".inc", ".inl"} + + +def main() -> None: + include_root = Path(sys.argv[1]) + jit_root = include_root / "nccl_ep" + public_header = include_root / "nccl_ep.h" + if not public_header.is_file(): + raise RuntimeError(f"Missing NCCL EP public header: {public_header}") + if not jit_root.is_dir(): + raise RuntimeError(f"Missing NCCL EP JIT header directory: {jit_root}") + + failures = [] + + headers = [public_header, *jit_root.rglob("*")] + for header in headers: + if not header.is_file() or header.suffix not in _HEADER_SUFFIXES: + continue + for line_number, line in enumerate(header.read_text().splitlines(), 1): + match = _INCLUDE_PATTERN.match(line) + if match is None: + continue + include = match.group(1) + if include in _EXTERNAL_HEADERS: + continue + candidates = ( + header.parent / include, + include_root / include, + jit_root / include, + jit_root / "device" / include, + ) + if not any(candidate.is_file() for candidate in candidates): + failures.append(f"{header.relative_to(include_root)}:{line_number}: {include}") + + if failures: + raise RuntimeError("Missing local NCCL EP JIT headers:\n" + "\n".join(failures)) + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index 2a8a9d7688..bd76776959 100644 --- a/setup.py +++ b/setup.py @@ -208,7 +208,7 @@ def _discover_nccl_home() -> str: def build_nccl_ep_submodule() -> str: - """Build libnccl_ep.a from the 3rdparty/nccl-extensions submodule and return NCCL_HOME.""" + """Build NCCL EP libraries from 3rdparty/nccl-extensions and return NCCL_HOME.""" nccl_root = current_file_path / "3rdparty" / "nccl-extensions" if not (nccl_root / "nccl_ep" / "Makefile").exists(): raise RuntimeError( @@ -217,7 +217,7 @@ def build_nccl_ep_submodule() -> str: ) build_dir = nccl_root / "build" - nccl_ep_lib = build_dir / "lib" / "libnccl_ep.a" + nccl_ep_shared_lib = build_dir / "lib" / "libnccl_ep.so" gencode_stamp = build_dir / "lib" / "libnccl_ep.gencode" # Caller gates on arch >= 90 or "native"; expand "native" to the host's @@ -261,18 +261,18 @@ def build_nccl_ep_submodule() -> str: env["NCCL_EP_BUILDDIR"] = str(build_dir) prev_gencode = gencode_stamp.read_text().strip() if gencode_stamp.exists() else None - if not nccl_ep_lib.exists() or prev_gencode != gencode: - if nccl_ep_lib.exists() and prev_gencode != gencode: + if not nccl_ep_shared_lib.exists() or prev_gencode != gencode: + if nccl_ep_shared_lib.exists() and prev_gencode != gencode: print( f"[NCCL EP] gencode changed ('{prev_gencode}' -> '{gencode}'); " - "rebuilding libnccl_ep.a" + "rebuilding NCCL EP libraries" ) subprocess.check_call( ["make", "-C", "nccl_ep", "clean"], cwd=str(nccl_root), env=env, ) - print(f"[NCCL EP] Building libnccl_ep.a (gencode='{gencode}')") + print(f"[NCCL EP] Building static and shared libraries (gencode='{gencode}')") make_jobs = f"-j{nproc}" if nproc else "-j" subprocess.check_call( ["make", make_jobs, "-C", "nccl_ep", "lib"], diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index 13b6242816..9ca071edfa 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -101,7 +101,7 @@ gtest_discover_tests(test_comm_gemm DISCOVERY_TIMEOUT 600) # Launched via mpirun; ncclUniqueId exchange uses MPI_Bcast (see test_ep_common.h). # The test binary only uses NCCL core symbols (ncclMemAlloc, ncclCommWindow*); # all ncclEp* calls live behind TE's public , which is -# statically linked into libtransformer_engine.so. +# backed by the runtime-loaded libnccl_ep.so library. message(STATUS "EP test: NCCL headers: ${NCCL_INCLUDE_DIR}") set(EP_TEST_COMMON_INCLUDES ${NCCL_INCLUDE_DIR} @@ -129,4 +129,4 @@ target_link_libraries(test_ep PUBLIC ${EP_TEST_COMMON_LIBS}) # Do NOT use gtest_discover_tests - these binaries require multi-process # launch via run_test_ep.sh, not direct single-process execution. -message(STATUS "EP distributed tests enabled (NCCL EP statically linked into libtransformer_engine.so)") +message(STATUS "EP distributed tests enabled (NCCL EP loaded from libnccl_ep.so at runtime)") diff --git a/tests/cpp_distributed/run_test_ep.sh b/tests/cpp_distributed/run_test_ep.sh index da293dadfd..02a66ed47b 100755 --- a/tests/cpp_distributed/run_test_ep.sh +++ b/tests/cpp_distributed/run_test_ep.sh @@ -53,6 +53,12 @@ if (( NUM_GPUS < 2 )); then exit 0 fi +# Force this test run to compile at least one NCCL EP JIT kernel instead of +# succeeding from a cache populated by an earlier process. +NCCL_EP_JIT_CACHE_DIR=$(mktemp -d "${TMPDIR:-/tmp}/te-nccl-ep-jit.XXXXXX") +export NCCL_EP_JIT_CACHE_DIR +trap 'rm -rf "${NCCL_EP_JIT_CACHE_DIR}"' EXIT + GTEST_ARGS="${GTEST_FILTER:+--gtest_filter=${GTEST_FILTER}}" echo "=== EP Tests ===" @@ -67,3 +73,8 @@ if [[ -n "${GTEST_XML_PREFIX:-}" ]]; then else "${MPIRUN}" --allow-run-as-root --oversubscribe -n "${NUM_GPUS}" ${MPIRUN_EXTRA:-} "${TEST_BIN}" ${GTEST_ARGS} fi + +if ! compgen -G "${NCCL_EP_JIT_CACHE_DIR}/*/*.cubin" > /dev/null; then + echo "ERROR: NCCL EP tests did not produce a JIT-compiled cubin." + exit 1 +fi diff --git a/tests/pytorch/test_nccl_ep_discovery.py b/tests/pytorch/test_nccl_ep_discovery.py new file mode 100644 index 0000000000..ac7edbc3d0 --- /dev/null +++ b/tests/pytorch/test_nccl_ep_discovery.py @@ -0,0 +1,60 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import transformer_engine + + +def _hide_packaged_library(monkeypatch): + def _missing_library(_): + raise FileNotFoundError + + monkeypatch.setattr( + transformer_engine.common, + "_get_shared_object_file", + _missing_library, + ) + + +def test_nccl_ep_library_found_from_home(monkeypatch, tmp_path): + home = tmp_path / "nccl_ep" + library_dir = home / "lib" + library_dir.mkdir(parents=True) + (library_dir / "libnccl_ep.so.0.1").touch() + + monkeypatch.setenv("NCCL_EP_HOME", str(home)) + _hide_packaged_library(monkeypatch) + monkeypatch.setattr(transformer_engine, "find_library", lambda _: None) + + assert transformer_engine._nccl_ep_library_installed() + + +def test_nccl_ep_library_found_in_package(monkeypatch, tmp_path): + library = tmp_path / "libnccl_ep.so" + library.touch() + + monkeypatch.delenv("NCCL_EP_HOME", raising=False) + monkeypatch.setattr( + transformer_engine.common, + "_get_shared_object_file", + lambda _: library, + ) + monkeypatch.setattr(transformer_engine, "find_library", lambda _: None) + + assert transformer_engine._nccl_ep_library_installed() + + +def test_nccl_ep_library_found_by_dynamic_loader(monkeypatch): + monkeypatch.delenv("NCCL_EP_HOME", raising=False) + _hide_packaged_library(monkeypatch) + monkeypatch.setattr(transformer_engine, "find_library", lambda _: "libnccl_ep.so.0") + + assert transformer_engine._nccl_ep_library_installed() + + +def test_nccl_ep_library_not_found(monkeypatch): + monkeypatch.delenv("NCCL_EP_HOME", raising=False) + _hide_packaged_library(monkeypatch) + monkeypatch.setattr(transformer_engine, "find_library", lambda _: None) + + assert not transformer_engine._nccl_ep_library_installed() diff --git a/transformer_engine/__init__.py b/transformer_engine/__init__.py index 480a2e9a06..1c35b1262b 100644 --- a/transformer_engine/__init__.py +++ b/transformer_engine/__init__.py @@ -9,11 +9,13 @@ import ctypes import functools import os +from ctypes.util import find_library from importlib import metadata +from pathlib import Path from typing import Optional, Tuple import transformer_engine.common -# Minimum NCCL version for the statically-linked NCCL EP backend. +# Minimum NCCL version for the runtime-loaded NCCL EP backend. _NCCL_EP_MIN_VERSION = (2, 30, 4) @@ -32,14 +34,39 @@ def _nccl_runtime_version() -> Optional[Tuple[int, int, int]]: return (v // 10000, (v // 100) % 100, v % 100) +def _nccl_ep_library_installed() -> bool: + if nccl_ep_home := os.getenv("NCCL_EP_HOME"): + home = Path(nccl_ep_home) + if any( + library.is_file() + for lib_dir in ("lib", "lib64") + for library in (home / lib_dir).glob("libnccl_ep.so*") + ): + return True + + try: + transformer_engine.common._get_shared_object_file("nccl_ep") + except FileNotFoundError: + return find_library("nccl_ep") is not None + else: + return True + + def is_nccl_ep_available() -> bool: - """Return True if the runtime libnccl.so meets the NCCL EP minimum.""" + """Return True if the NCCL EP library and a compatible NCCL runtime are available.""" + if not _nccl_ep_library_installed(): + return False cur = _nccl_runtime_version() return cur is not None and cur >= _NCCL_EP_MIN_VERSION def require_nccl_ep() -> None: """Raise RuntimeError if NCCL EP cannot run on the current libnccl.""" + if not _nccl_ep_library_installed(): + raise RuntimeError( + "NCCL EP library libnccl_ep.so is not installed. Build Transformer Engine " + "with NVTE_WITH_NCCL_EP=1." + ) mn = ".".join(str(x) for x in _NCCL_EP_MIN_VERSION) cur = _nccl_runtime_version() if cur is None: diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index e7aaf78f6c..3503941cf4 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -455,7 +455,7 @@ endif() # 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) +option(NVTE_WITH_NCCL_EP "Build and package the runtime-loaded NCCL EP library" ON) if(NVTE_WITH_NCCL_EP) # SM>=90 and NCCL>=2.30.4 are gated at runtime in EPBackend::initialize. # -- NCCL EP headers -------------------------------------------------------- @@ -471,17 +471,17 @@ if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") 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. +# -- libnccl_ep.so ---------------------------------------------------------- +# Loaded on demand by EPBackend only after it verifies NCCL >= 2.30.4. Keeping +# NCCL EP out of libtransformer_engine.so allows non-EP users to run with older +# NCCL versions even when the core library uses eager symbol binding. set(NCCL_EP_LIB_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/lib") -find_file(NCCL_EP_LIB - NAMES libnccl_ep.a +find_file(NCCL_EP_SHARED_LINK + NAMES libnccl_ep.so HINTS ${NCCL_EP_LIB_DIR} NO_DEFAULT_PATH REQUIRED) +file(REAL_PATH "${NCCL_EP_SHARED_LINK}" NCCL_EP_SHARED_LIB) # -- NCCL core library ------------------------------------------------------- if(NOT NCCL_LIB) @@ -494,20 +494,17 @@ endif() target_include_directories(transformer_engine PRIVATE ${NCCL_EP_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. +# libtransformer_engine.so uses only the baseline NCCL API directly. NCCL EP's +# newer symbols remain isolated in libnccl_ep.so until EP is initialized. 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) + ep/ep_api.cpp + ep/nccl_ep_provider.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 enabled (runtime-loaded library): ${NCCL_EP_SHARED_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. @@ -677,3 +674,16 @@ message(STATUS "Threads per parallel build job: ${BUILD_THREADS_PER_JOB}") # Install library install(TARGETS transformer_engine DESTINATION .) +if(NVTE_WITH_NCCL_EP) + # Install a real file rather than the build-tree symlink so wheel packaging + # does not need to preserve the versioned symlink chain. + install(FILES "${NCCL_EP_SHARED_LIB}" + DESTINATION . + RENAME libnccl_ep.so) + # NCCL EP JIT compilation requires both the public header root and staged + # device headers under NCCL_EP_HOME/include/nccl_ep. + install(FILES "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h" + DESTINATION nccl_ep/include) + install(DIRECTORY "${NCCL_EP_INCLUDE_DIR}/nccl_ep" + DESTINATION nccl_ep/include) +endif() diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 6b276d22a8..e47a6f11b2 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -107,9 +107,11 @@ def _get_shared_object_file(library: str) -> Path: """ # Check provided input and determine the correct prefix for .so. - assert library in ("core", "torch", "jax"), f"Unsupported TE library {library}." + assert library in ("core", "torch", "jax", "nccl_ep"), f"Unsupported TE library {library}." if library == "core": so_prefix = "libtransformer_engine" + elif library == "nccl_ep": + so_prefix = "libnccl_ep" else: so_prefix = f"transformer_engine_{library}" diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index fb608639a8..3dd3d6e972 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -23,6 +23,7 @@ #include "../common.h" #include "../util/cuda_runtime.h" #include "../util/logging.h" +#include "nccl_ep_provider.h" namespace transformer_engine { namespace ep { @@ -163,6 +164,9 @@ void EPBackend::initialize(ncclComm_t ep_comm, NVTEEpGroupConfig config) { nccl_version / 10000, ".", (nccl_version / 100) % 100, ".", nccl_version % 100, " at runtime."); + // Load libnccl_ep only after the runtime NCCL version has been validated. + nccl_ep::initialize(); + validate_config(config); int comm_size = 0; @@ -178,14 +182,14 @@ void EPBackend::shutdown() { std::lock_guard lock(inst.mutex_); if (!inst.initialized_) return; for (auto& e : inst.lru_) { - if (e.handle != nullptr) ncclEpHandleDestroy(e.handle); + if (e.handle != nullptr) nccl_ep::handle_destroy(e.handle); } inst.lru_.clear(); inst.index_.clear(); inst.fallback_layer_cfg_.reset(); // ncclEpGroupDestroy reads from ep_comm_; destroy group while comm is still alive. if (inst.ep_group_ != nullptr) { - ncclEpGroupDestroy(inst.ep_group_); + nccl_ep::group_destroy(inst.ep_group_); inst.ep_group_ = nullptr; } inst.ep_comm_ = nullptr; // borrowed; caller destroys @@ -203,8 +207,8 @@ ncclEpHandle_t EPBackend::open_handle(void* handle_mem, size_t handle_mem_size, ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; hcfg.dispatch_output_per_expert_alignment = dispatch_output_per_expert_alignment; ncclEpHandle_t handle; - NVTE_CHECK_NCCL(ncclEpInitHandle(&handle, ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, num_topk, - &routing_desc)); + NVTE_CHECK_NCCL(nccl_ep::init_handle(&handle, ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, + num_topk, &routing_desc)); return handle; } @@ -261,7 +265,7 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) { setenv("NCCL_EP_SHUFFLE_SMS", sm_buf, /*overwrite=*/0); setenv("NCCL_EP_PREPROCESS_NUM_SMS", sm_buf, /*overwrite=*/0); - NVTE_CHECK_NCCL(ncclEpCreateGroup(&ep_group_, ep_comm, &cfg)); + NVTE_CHECK_NCCL(nccl_ep::create_group(&ep_group_, ep_comm, &cfg)); ep_comm_ = ep_comm; @@ -319,15 +323,15 @@ ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerCon ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; hcfg.dispatch_output_per_expert_alignment = layer_cfg.dispatch_output_per_expert_alignment; size_t hm_size = 0; - NVTE_CHECK_NCCL(ncclEpHandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, &hm_size, - layer_cfg.top_k)); + NVTE_CHECK_NCCL(nccl_ep::handle_mem_size(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, &hm_size, + layer_cfg.top_k)); ncclEpHandle_t h = open_handle(handle_mem, hm_size, layer_cfg.top_k, layer_cfg.dispatch_output_per_expert_alignment); lru_.push_front(HandleEntry{handle_mem, h, layer_cfg, hm_size}); index_.emplace(handle_mem, lru_.begin()); while (lru_.size() > cache_cap_locked()) { HandleEntry& victim = lru_.back(); - if (victim.handle != nullptr) ncclEpHandleDestroy(victim.handle); + if (victim.handle != nullptr) nccl_ep::handle_destroy(victim.handle); index_.erase(victim.handle_mem); lru_.pop_back(); } @@ -361,8 +365,8 @@ size_t EPBackend::handle_mem_size(NVTEEpLayerConfig layer_cfg) { ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; hcfg.dispatch_output_per_expert_alignment = layer_cfg.dispatch_output_per_expert_alignment; size_t hm_size = 0; - NVTE_CHECK_NCCL(ncclEpHandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, &hm_size, - layer_cfg.top_k)); + NVTE_CHECK_NCCL(nccl_ep::handle_mem_size(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, &hm_size, + layer_cfg.top_k)); return hm_size; } @@ -398,7 +402,7 @@ void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); ncclEpHandle_t h = prepare_handle_locked(handle_mem, layer_cfg); - NVTE_CHECK_NCCL(ncclEpUpdateHandle(h, &nccl_topk_idx, &layout_info, stream)); + NVTE_CHECK_NCCL(nccl_ep::update_handle(h, &nccl_topk_idx, &layout_info, stream)); } void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTETensor tokens, @@ -489,8 +493,8 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); ncclEpHandle_t h = lookup_handle_locked(handle_mem); - NVTE_CHECK_NCCL(ncclEpDispatch(h, &in_struct, &out_struct, - /*layout_info=*/nullptr, &dispatch_cfg, stream)); + NVTE_CHECK_NCCL(nccl_ep::dispatch(h, &in_struct, &out_struct, + /*layout_info=*/nullptr, &dispatch_cfg, stream)); } void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, @@ -513,7 +517,7 @@ void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); ncclEpHandle_t h = lookup_handle_locked(handle_mem); - NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, /*config=*/nullptr, stream)); + NVTE_CHECK_NCCL(nccl_ep::combine(h, &in_struct, &out_struct, /*config=*/nullptr, stream)); } void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, @@ -551,7 +555,7 @@ void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); ncclEpHandle_t h = lookup_handle_locked(handle_mem); - NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, &cfg, stream)); + NVTE_CHECK_NCCL(nccl_ep::combine(h, &in_struct, &out_struct, &cfg, stream)); } void EPBackend::combine_bwd(void* handle_mem, const NVTETensor grad, const NVTECommWindow& grad_win, diff --git a/transformer_engine/common/ep/nccl_ep_provider.cpp b/transformer_engine/common/ep/nccl_ep_provider.cpp new file mode 100644 index 0000000000..f91d06ea58 --- /dev/null +++ b/transformer_engine/common/ep/nccl_ep_provider.cpp @@ -0,0 +1,304 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "nccl_ep_provider.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../util/cuda_runtime.h" +#include "../util/logging.h" +#include "../util/shared_lib_wrapper.h" + +namespace transformer_engine { +namespace ep { +namespace nccl_ep { +namespace { + +constexpr char kNCCLEPLibraryName[] = "libnccl_ep.so"; +const char kLibraryAnchor = 0; + +bool env_is_set(const char* name) { + const char* value = std::getenv(name); + return value != nullptr && value[0] != '\0'; +} + +void set_default_env(const char* name, const std::filesystem::path& value) { + NVTE_CHECK(setenv(name, value.c_str(), 0) == 0, "Could not set ", name, ": ", + std::strerror(errno)); +} + +void append_unique(std::vector* paths, const std::filesystem::path& path) { + if (path.empty()) return; + for (const auto& existing : *paths) { + if (existing == path) return; + } + paths->push_back(path); +} + +// Add the conventional unversioned and major-version paths first, then discover +// fully versioned files (for example, libnccl_ep.so.0.1) for installations that +// omit symlinks. Preserve lookup order while avoiding duplicate dlopen attempts. +void append_library_dir(std::vector* candidates, + const std::filesystem::path& directory, const std::string& versioned_name) { + append_unique(candidates, directory / kNCCLEPLibraryName); + append_unique(candidates, directory / versioned_name); + + std::error_code error; + if (!std::filesystem::is_directory(directory, error) || error) return; + + std::vector fully_versioned; + const std::string prefix = versioned_name + "."; + for (std::filesystem::directory_iterator iterator(directory, error), end; + !error && iterator != end; iterator.increment(error)) { + const std::filesystem::directory_entry& entry = *iterator; + const std::string filename = entry.path().filename().string(); + std::error_code file_error; + if (filename.rfind(prefix, 0) == 0 && entry.is_regular_file(file_error) && !file_error) { + fully_versioned.push_back(entry.path()); + } + } + std::sort(fully_versioned.rbegin(), fully_versioned.rend()); + for (const auto& path : fully_versioned) append_unique(candidates, path); +} + +std::vector nccl_ep_library_candidates() { + using Path = std::filesystem::path; + std::vector candidates; + const std::string versioned_name = + std::string(kNCCLEPLibraryName) + "." + std::to_string(NCCL_EP_MAJOR); + + if (env_is_set("NCCL_EP_HOME")) { + const Path home = std::getenv("NCCL_EP_HOME"); + append_library_dir(&candidates, home / "lib", versioned_name); + append_library_dir(&candidates, home / "lib64", versioned_name); + } + + const Path library_dir = shared_library_directory(static_cast(&kLibraryAnchor)); + append_library_dir(&candidates, library_dir, versioned_name); + append_unique(&candidates, versioned_name); + append_unique(&candidates, kNCCLEPLibraryName); + return candidates; +} + +std::optional nccl_header_version(const std::filesystem::path& include_dir) { + std::ifstream header(include_dir / "nccl.h"); + if (!header) return std::nullopt; + + int major = -1; + int minor = -1; + int patch = -1; + for (std::string line; std::getline(header, line);) { + if (major < 0) std::sscanf(line.c_str(), "#define NCCL_MAJOR %d", &major); + if (minor < 0) std::sscanf(line.c_str(), "#define NCCL_MINOR %d", &minor); + if (patch < 0) std::sscanf(line.c_str(), "#define NCCL_PATCH %d", &patch); + } + if (major < 0 || minor < 0 || patch < 0) return std::nullopt; + return major * 10000 + minor * 100 + patch; +} + +void append_ancestor_include_dirs(std::vector* candidates, + std::filesystem::path directory) { + while (!directory.empty()) { + const std::filesystem::path parent = directory.parent_path(); + if (parent == directory) break; + append_unique(candidates, directory / "include"); + append_unique(candidates, directory / "nvidia" / "nccl" / "include"); + directory = parent; + } +} + +void configure_nccl_ep_source_dir(const std::filesystem::path& library_path) { + if (env_is_set("NCCL_EP_HOME") || env_is_set("NCCL_EP_JIT_SOURCE_DIR")) { + return; + } + + const std::filesystem::path library_dir = library_path.parent_path(); + const std::filesystem::path homes[] = { + library_dir / "nccl_ep", + library_dir, + library_dir.parent_path(), + }; + for (const auto& home : homes) { + std::error_code error; + if (std::filesystem::is_directory(home / "include" / "nccl_ep", error) && !error) { + set_default_env("NCCL_EP_HOME", home); + return; + } + } +} + +void configure_nccl_include_dir() { + if (env_is_set("NCCL_EP_JIT_BUILD_INCLUDE_DIR") || env_is_set("NCCL_HOME")) { + return; + } + + int runtime_version = 0; + NVTE_CHECK_NCCL(ncclGetVersion(&runtime_version)); + + std::vector candidates; + append_ancestor_include_dirs( + &candidates, shared_library_directory(reinterpret_cast(&ncclGetVersion))); + append_ancestor_include_dirs(&candidates, + shared_library_directory(static_cast(&kLibraryAnchor))); + append_unique(&candidates, "/opt/nvidia/nccl/include"); + append_unique(&candidates, "/usr/local/nccl/include"); + append_unique(&candidates, "/usr/include"); + + for (const auto& candidate : candidates) { + if (nccl_header_version(candidate) == runtime_version) { + set_default_env("NCCL_EP_JIT_BUILD_INCLUDE_DIR", candidate); + return; + } + } +} + +void configure_cuda_include_dir() { + if (env_is_set("NCCL_EP_JIT_CUDA_INCLUDE_DIR") || env_is_set("CUDA_HOME") || + env_is_set("CUDA_PATH")) { + return; + } + + const std::string& include_dir = cuda::include_directory(false); + if (!include_dir.empty()) { + set_default_env("NCCL_EP_JIT_CUDA_INCLUDE_DIR", include_dir); + } +} + +void configure_jit_environment(const std::filesystem::path& library_path) { + configure_nccl_ep_source_dir(library_path); + configure_nccl_include_dir(); + configure_cuda_include_dir(); +} + +void* load_symbol(void* handle, const char* name) { + dlerror(); + void* symbol = dlsym(handle, name); + const char* error = dlerror(); + NVTE_CHECK(error == nullptr && symbol != nullptr, "Could not load ", name, " from ", + kNCCLEPLibraryName, ": ", error == nullptr ? "symbol not found" : error); + return symbol; +} + +void* open_nccl_ep_library() { + std::string failures; + for (const auto& candidate : nccl_ep_library_candidates()) { + dlerror(); + if (void* handle = dlopen(candidate.c_str(), RTLD_NOW | RTLD_LOCAL)) { + return handle; + } + const char* error = dlerror(); + if (!failures.empty()) failures += "; "; + failures += + candidate.string() + ": " + (error == nullptr ? "unknown dynamic loader error" : error); + } + NVTE_ERROR("Could not load ", kNCCLEPLibraryName, ". Tried ", failures); + return nullptr; // Unreachable. +} + +void* library_handle() { + // Deliberately keep libnccl_ep loaded until process exit. EPBackend owns + // objects whose implementation lives in this library. + static void* handle = [] { + void* result = open_nccl_ep_library(); + + try { + using GetVersion = decltype(&ncclEpGetVersion); + auto get_version = reinterpret_cast(load_symbol(result, "ncclEpGetVersion")); + int runtime_version = 0; + NVTE_CHECK_NCCL(get_version(&runtime_version)); + NVTE_CHECK(runtime_version / 10000 == NCCL_EP_MAJOR, + "Incompatible NCCL EP library major version ", runtime_version / 10000, + "; expected ", NCCL_EP_MAJOR); + const std::filesystem::path library_path = + shared_library_path(reinterpret_cast(get_version)); + NVTE_CHECK(!library_path.empty(), "Could not determine the NCCL EP library path"); + configure_jit_environment(library_path); + } catch (...) { + dlclose(result); + throw; + } + return result; + }(); + return handle; +} + +} // namespace + +void initialize() { (void)library_handle(); } + +void* get_symbol(const char* symbol) { return load_symbol(library_handle(), symbol); } + +namespace { + +template +auto call_symbol(const char* name, Args&&... args) { + using FuncT = decltype(FuncPtr); + static FuncT func = reinterpret_cast(get_symbol(name)); + return func(std::forward(args)...); +} + +} // namespace + +ncclResult_t create_group(ncclEpGroup_t* ep_group, ncclComm_t comm, + const ncclEpGroupConfig_t* config) { + return call_symbol<&ncclEpCreateGroup>("ncclEpCreateGroup", ep_group, comm, config); +} + +ncclResult_t group_destroy(ncclEpGroup_t ep_group) { + return call_symbol<&ncclEpGroupDestroy>("ncclEpGroupDestroy", ep_group); +} + +ncclResult_t handle_destroy(ncclEpHandle_t handle) { + return call_symbol<&ncclEpHandleDestroy>("ncclEpHandleDestroy", handle); +} + +ncclResult_t init_handle(ncclEpHandle_t* handle, ncclEpGroup_t ep_group, ncclEpLayout_t layout, + const ncclEpHandleConfig_t* config, int num_topk, + const ncclEpTensor_t* handle_mem) { + return call_symbol<&ncclEpInitHandle>("ncclEpInitHandle", handle, ep_group, layout, config, + num_topk, handle_mem); +} + +ncclResult_t handle_mem_size(ncclEpGroup_t ep_group, ncclEpLayout_t layout, + const ncclEpHandleConfig_t* config, size_t* size_out, int num_topk) { + return call_symbol<&ncclEpHandleMemSize>("ncclEpHandleMemSize", ep_group, layout, config, + size_out, num_topk); +} + +ncclResult_t update_handle(ncclEpHandle_t handle, const ncclEpTensor_t* topk_idx, + const ncclEpLayoutInfo_t* layout_info, cudaStream_t stream) { + return call_symbol<&ncclEpUpdateHandle>("ncclEpUpdateHandle", handle, topk_idx, layout_info, + stream); +} + +ncclResult_t dispatch(ncclEpHandle_t handle, const ncclEpDispatchInputs_t* inputs, + const ncclEpDispatchOutputs_t* outputs, const ncclEpLayoutInfo_t* layout_info, + const ncclEpDispatchConfig_t* config, cudaStream_t stream) { + return call_symbol<&ncclEpDispatch>("ncclEpDispatch", handle, inputs, outputs, layout_info, + config, stream); +} + +ncclResult_t combine(ncclEpHandle_t handle, const ncclEpCombineInputs_t* inputs, + const ncclEpCombineOutputs_t* outputs, const ncclEpCombineConfig_t* config, + cudaStream_t stream) { + return call_symbol<&ncclEpCombine>("ncclEpCombine", handle, inputs, outputs, config, stream); +} + +} // namespace nccl_ep +} // namespace ep +} // namespace transformer_engine diff --git a/transformer_engine/common/ep/nccl_ep_provider.h b/transformer_engine/common/ep/nccl_ep_provider.h new file mode 100644 index 0000000000..d806c656f9 --- /dev/null +++ b/transformer_engine/common/ep/nccl_ep_provider.h @@ -0,0 +1,46 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_EP_NCCL_EP_PROVIDER_H_ +#define TRANSFORMER_ENGINE_COMMON_EP_NCCL_EP_PROVIDER_H_ + +#include + +namespace transformer_engine { +namespace ep { + +namespace nccl_ep { + +/*! \brief Load and validate libnccl_ep.so. Idempotent and thread-safe. */ +void initialize(); + +/*! \brief Get a function pointer from the runtime NCCL EP library. */ +void* get_symbol(const char* symbol); + +ncclResult_t create_group(ncclEpGroup_t* ep_group, ncclComm_t comm, + const ncclEpGroupConfig_t* config); +ncclResult_t group_destroy(ncclEpGroup_t ep_group); +ncclResult_t handle_destroy(ncclEpHandle_t handle); +ncclResult_t init_handle(ncclEpHandle_t* handle, ncclEpGroup_t ep_group, ncclEpLayout_t layout, + const ncclEpHandleConfig_t* config, int num_topk, + const ncclEpTensor_t* handle_mem); +ncclResult_t handle_mem_size(ncclEpGroup_t ep_group, ncclEpLayout_t layout, + const ncclEpHandleConfig_t* config, size_t* size_out, int num_topk); +ncclResult_t update_handle(ncclEpHandle_t handle, const ncclEpTensor_t* topk_idx, + const ncclEpLayoutInfo_t* layout_info, cudaStream_t stream); +ncclResult_t dispatch(ncclEpHandle_t handle, const ncclEpDispatchInputs_t* inputs, + const ncclEpDispatchOutputs_t* outputs, const ncclEpLayoutInfo_t* layout_info, + const ncclEpDispatchConfig_t* config, cudaStream_t stream); +ncclResult_t combine(ncclEpHandle_t handle, const ncclEpCombineInputs_t* inputs, + const ncclEpCombineOutputs_t* outputs, const ncclEpCombineConfig_t* config, + cudaStream_t stream); + +} // namespace nccl_ep + +} // namespace ep +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_EP_NCCL_EP_PROVIDER_H_ diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index 9c4b23912a..93a855284c 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -94,6 +94,9 @@ typedef struct { /*! \brief Bootstrap the EP backend from an existing NCCL EP sub-communicator. * Requires SM>=90. * + * This call validates that the runtime NCCL is >=2.30.4 and then loads the + * optional libnccl_ep.so library. Non-EP users do not load the library. + * * ep_comm is borrowed and must span exactly group_config.ep_size ranks. The * caller retains ownership and must keep it alive until nvte_ep_shutdown() * returns. Re-init after shutdown is allowed; double-init throws. One EP diff --git a/transformer_engine/common/util/cuda_runtime.cpp b/transformer_engine/common/util/cuda_runtime.cpp index 5690ef5e3c..d23f11ccc6 100644 --- a/transformer_engine/common/util/cuda_runtime.cpp +++ b/transformer_engine/common/util/cuda_runtime.cpp @@ -7,7 +7,6 @@ #include "../util/cuda_runtime.h" #include -#include #include #include @@ -15,6 +14,7 @@ #include "../common.h" #include "../util/cuda_driver.h" +#include "../util/shared_lib_wrapper.h" #include "../util/system.h" #include "common/util/cuda_runtime.h" @@ -27,27 +27,6 @@ namespace { // String with build-time CUDA include path #include "string_path_cuda_include.h" -// Get the runtime directory of the shared library that contains this code -std::filesystem::path shared_library_directory() { - static const char library_anchor = 0; - Dl_info library_info{}; - if (dladdr(static_cast(&library_anchor), &library_info) == 0 || - library_info.dli_fname == nullptr) { - return {}; - } - - std::filesystem::path library_path = library_info.dli_fname; - if (library_path.is_relative()) { - std::error_code error; - library_path = std::filesystem::absolute(library_path, error); - if (error) { - return {}; - } - } - - return library_path.parent_path(); -} - std::string runtime_cuda_major_version() { int runtime_version = 0; // Header discovery is best-effort, so do not throw if the runtime cannot @@ -65,7 +44,9 @@ std::filesystem::path python_cuda_directory() { // Find the Python package root from the installed Transformer Engine package. Do not // assume that the root is named site-packages or dist-packages since valid installs // may use an arbitrary target directory. - Path te_package_directory = shared_library_directory(); + static const char library_anchor = 0; + Path te_package_directory = + transformer_engine::shared_library_directory(static_cast(&library_anchor)); while (true) { if (te_package_directory.filename() == "transformer_engine") { break; diff --git a/transformer_engine/common/util/shared_lib_wrapper.h b/transformer_engine/common/util/shared_lib_wrapper.h index e8abe68a2a..79e5a598eb 100644 --- a/transformer_engine/common/util/shared_lib_wrapper.h +++ b/transformer_engine/common/util/shared_lib_wrapper.h @@ -9,8 +9,34 @@ #include +#include + namespace transformer_engine { +/*! \brief Return the shared object containing an address. */ +inline std::filesystem::path shared_library_path(const void *anchor) { + Dl_info library_info{}; + if (anchor == nullptr || dladdr(anchor, &library_info) == 0 || + library_info.dli_fname == nullptr) { + return {}; + } + + std::filesystem::path library_path = library_info.dli_fname; + if (library_path.is_relative()) { + std::error_code error; + library_path = std::filesystem::absolute(library_path, error); + if (error) { + return {}; + } + } + return library_path; +} + +/*! \brief Return the directory containing the shared object for an address. */ +inline std::filesystem::path shared_library_directory(const void *anchor) { + return shared_library_path(anchor).parent_path(); +} + /*! \brief Wrapper class for a shared library * * \todo Windows support