diff --git a/build_tools/build_backend.py b/build_tools/build_backend.py new file mode 100644 index 00000000000..cca1848f783 --- /dev/null +++ b/build_tools/build_backend.py @@ -0,0 +1,269 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PEP 517 backend with CUDA-version-, framework-, and feature-specific build requirements.""" + +from contextlib import contextmanager +import os +import re +import subprocess +from typing import Iterator, List, NamedTuple, Optional, Tuple + +from build_tools.build_backend_utils import ( + ConfigSettings, + config_value, + normalize_cuda_version, + requested_cuda_version, + setuptools_hooks, +) +from build_tools.utils import MIN_CUDNN_VERSION, nvcc_path + +# cuda-toolkit pins all components to the same CUDA release. The nvcc extra includes the +# compiler's crt and nvvm components. +_CUDA_TOOLKIT_EXTRAS = ("nvcc", "cudart", "cccl", "cublas", "nvml", "profiler") +_FRAMEWORK_BUILD_PACKAGES = { + "pytorch": ["torch>=2.1"], + "jax": ["jax>=0.5.0", "flax>=0.7.1"], +} +_SUPPORTED_FRAMEWORKS = tuple(_FRAMEWORK_BUILD_PACKAGES) + + +class _Feature(NamedTuple): + """Optional feature that can be enabled with an environment variable or config setting.""" + + environment_variable: str + # Selects an existing installation instead of installing the requirements. + home_variable: Optional[str] + # Formatted with the CUDA major version. + requirements: Tuple[str, ...] + + +# Keyed by config setting, e.g. --config-settings with-cublasmp=1. +_FEATURES = { + "with-cublasmp": _Feature( + "NVTE_WITH_CUBLASMP", + "CUBLASMP_HOME", + # cuBLASMp's overlap algorithms require NCCL 2.30. + ("nvidia-cublasmp-cu{cuda_major}>=0.8.1", "nvidia-nccl-cu{cuda_major}>=2.30"), + ), + "with-cusolvermp": _Feature( + "NVTE_WITH_CUSOLVERMP", "CUSOLVERMP_HOME", ("nvidia-cusolvermp-cu{cuda_major}",) + ), + "enable-nvshmem": _Feature( + "NVTE_ENABLE_NVSHMEM", "NVSHMEM_HOME", ("nvidia-nvshmem-cu{cuda_major}",) + ), + "ub-with-mpi": _Feature("NVTE_UB_WITH_MPI", "MPI_HOME", ("openmpi",)), + "with-cutedsl": _Feature("NVTE_WITH_CUTEDSL", None, ("apache-tvm-ffi>=0.1.12",)), +} +_PRIVATE_SETTINGS = ("cuda-version", "framework", *_FEATURES) + + +def _normalize_frameworks(value: str) -> List[str]: + """Validate and normalize a framework selection.""" + frameworks = list( + dict.fromkeys( + framework.strip().lower() for framework in value.split(",") if framework.strip() + ) + ) + if not frameworks: + raise ValueError("Framework selection cannot be empty") + + special_frameworks = {"all", "none"}.intersection(frameworks) + if special_frameworks: + if len(frameworks) != 1: + raise ValueError("'all' and 'none' cannot be combined with other frameworks") + return list(_SUPPORTED_FRAMEWORKS) if frameworks[0] == "all" else [] + + unsupported = [framework for framework in frameworks if framework not in _SUPPORTED_FRAMEWORKS] + if unsupported: + supported = ", ".join((*_SUPPORTED_FRAMEWORKS, "all", "none")) + raise ValueError( + f"Unsupported framework {unsupported[0]!r}; expected one or more of: {supported}" + ) + + return frameworks + + +def _config_frameworks(config_settings: ConfigSettings) -> Optional[List[str]]: + """Get the target frameworks from PEP 517 config settings.""" + value = config_value(config_settings, "framework") + return _normalize_frameworks(value) if value is not None else None + + +def _environment_frameworks() -> Optional[List[str]]: + """Get the target frameworks from the environment.""" + value = os.getenv("NVTE_FRAMEWORK") + return _normalize_frameworks(value) if value is not None else None + + +def _requested_frameworks(config_settings: ConfigSettings) -> Optional[List[str]]: + """Resolve an explicit framework selection in descending order of precedence.""" + configured = _config_frameworks(config_settings) + return configured if configured is not None else _environment_frameworks() + + +def _config_feature(config_settings: ConfigSettings, setting: str) -> Optional[bool]: + """Get whether PEP 517 config settings enable or disable an optional feature.""" + value = config_value(config_settings, setting) + if value is None: + return None + + if value not in ("0", "1"): + raise ValueError(f"Invalid value {value!r} for {setting}; expected 0 or 1") + return value == "1" + + +def _feature_enabled(config_settings: ConfigSettings, setting: str) -> bool: + """Resolve whether an optional feature is enabled, preferring config settings.""" + configured = _config_feature(config_settings, setting) + if configured is not None: + return configured + return bool(int(os.getenv(_FEATURES[setting].environment_variable, "0"))) + + +@contextmanager +def _build_environment(config_settings: ConfigSettings) -> Iterator[None]: + """Apply config-setting framework and feature selections while running setuptools.""" + overrides = {} + frameworks = _config_frameworks(config_settings) + if frameworks is not None: + overrides["NVTE_FRAMEWORK"] = ",".join(frameworks) if frameworks else "none" + for setting, feature in _FEATURES.items(): + enabled = _config_feature(config_settings, setting) + if enabled is not None: + overrides[feature.environment_variable] = "1" if enabled else "0" + + previous = {name: os.environ.get(name) for name in overrides} + os.environ.update(overrides) + try: + yield + finally: + for name, value in previous.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + +def _nvcc_cuda_version() -> Optional[str]: + """Get the version reported by the selected NVCC.""" + if (cuda_compiler := nvcc_path()) is None: + return None + + try: + result = subprocess.run( + [cuda_compiler, "--version"], + capture_output=True, + check=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + return None + + match = re.search(r"release\s+(\d+\.\d+)", result.stdout) + return normalize_cuda_version(match.group(1)) if match is not None else None + + +def _torch_cuda_version() -> Optional[str]: + """Get the CUDA version required by PyTorch, if available.""" + try: + import torch + except (ImportError, OSError): + return None + value = getattr(getattr(torch, "version", None), "cuda", None) + return normalize_cuda_version(value) if value else None + + +def _framework_cuda_version(config_settings: ConfigSettings) -> Optional[str]: + """Get a CUDA version required by the selected framework.""" + frameworks = _requested_frameworks(config_settings) + + if frameworks is None or "pytorch" in frameworks: + if cuda_version := _torch_cuda_version(): + return cuda_version + + # JAX CUDA plugin names identify the CUDA major version, but not the + # major/minor toolkit release needed to select component packages. + return None + + +def _cuda_version(config_settings: ConfigSettings) -> str: + """Resolve the CUDA version in descending order of precedence.""" + cuda_version = ( + requested_cuda_version(config_settings) + or _nvcc_cuda_version() + or _framework_cuda_version(config_settings) + ) + if cuda_version is None: + raise ValueError( + "Could not determine the CUDA version for an isolated build; install a CUDA " + "toolkit, set CUDA_HOME, or pass --config-settings cuda-version=. " + "(or set NVTE_CUDA_VERSION)" + ) + return cuda_version + + +def _cuda_build_requirements(config_settings: ConfigSettings) -> List[str]: + """Get build requirements for the requested CUDA version. + + The CUDA Toolkit and cuDNN are installed from Python packages unless CUDA_HOME selects an + installed CUDA Toolkit, in which case the installed cuDNN is used as well. + """ + cuda_version = _cuda_version(config_settings) + cuda_major = cuda_version.split(".", maxsplit=1)[0] + requirements = [f"nvidia-nccl-cu{cuda_major}>=2"] + if not os.getenv("CUDA_HOME"): + requirements += [ + f"cuda-toolkit[{','.join(_CUDA_TOOLKIT_EXTRAS)}]=={cuda_version}.*", + f"nvidia-cudnn-cu{cuda_major}>={MIN_CUDNN_VERSION}", + ] + return requirements + + +def _framework_build_requirements(config_settings: ConfigSettings) -> List[str]: + """Get build requirements for the requested frameworks.""" + frameworks = _requested_frameworks(config_settings) + if frameworks is None: + # Without a selection, setup.py detects frameworks in the isolated build + # environment, which has none, so only the core library is built. + return [] + + return [ + requirement + for framework in frameworks + for requirement in _FRAMEWORK_BUILD_PACKAGES[framework] + ] + + +def _feature_build_requirements(config_settings: ConfigSettings) -> List[str]: + """Get build requirements for the enabled optional features.""" + cuda_major = _cuda_version(config_settings).split(".", maxsplit=1)[0] + return [ + requirement.format(cuda_major=cuda_major) + for setting, feature in _FEATURES.items() + if _feature_enabled(config_settings, setting) + and not (feature.home_variable and os.getenv(feature.home_variable)) + for requirement in feature.requirements + ] + + +def _build_requirements(config_settings: ConfigSettings) -> List[str]: + """Get CUDA, framework, and optional feature build requirements.""" + framework_requirements = _framework_build_requirements(config_settings) + return ( + _cuda_build_requirements(config_settings) + + framework_requirements + + _feature_build_requirements(config_settings) + ) + + +# PEP 517 and PEP 660 hooks. setup.py always depends on the build requirements, so they are +# also needed for sdists. +globals().update( + setuptools_hooks( + requirements=_build_requirements, + private_settings=_PRIVATE_SETTINGS, + environment=_build_environment, + ) +) diff --git a/build_tools/build_backend_utils.py b/build_tools/build_backend_utils.py new file mode 100644 index 00000000000..a589cf41bf5 --- /dev/null +++ b/build_tools/build_backend_utils.py @@ -0,0 +1,163 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared helpers for Transformer Engine's PEP 517 build backends. + +Only the standard library and setuptools may be imported here, since the backends import +this module before any other build requirements are installed. +""" + +from contextlib import nullcontext +import os +import re +from typing import Callable, ContextManager, Dict, Iterable, List, Mapping, Optional, Union + +from setuptools import build_meta + +SUPPORTED_CUDA_MAJOR_VERSIONS = ("12", "13") + +ConfigValue = Union[str, List[str]] +ConfigSettings = Optional[Mapping[str, ConfigValue]] + + +def config_value(config_settings: ConfigSettings, setting: str) -> Optional[str]: + """Get a PEP 517 config setting, using the last value if it was given more than once.""" + settings: Mapping[str, ConfigValue] = config_settings or {} + value = settings.get(setting) + if value is None: + return None + + if isinstance(value, list): + if not value: + raise ValueError(f"{setting} cannot be empty") + value = value[-1] + + return str(value) + + +def normalize_cuda_version(value: str) -> str: + """Validate and normalize a CUDA major/minor version.""" + value = value.lower().removeprefix("cu") + if re.fullmatch(r"\d+\.\d+", value) is None: + raise ValueError( + f"Invalid CUDA version {value!r}; expected ., for example 12.8" + ) + + major = value.split(".", maxsplit=1)[0] + if major not in SUPPORTED_CUDA_MAJOR_VERSIONS: + supported = ", ".join(SUPPORTED_CUDA_MAJOR_VERSIONS) + raise ValueError(f"Unsupported CUDA major version {major!r}; expected one of: {supported}") + + return value + + +def requested_cuda_version(config_settings: ConfigSettings) -> Optional[str]: + """Get the CUDA version from the cuda-version config setting or NVTE_CUDA_VERSION.""" + value = config_value(config_settings, "cuda-version") + if value is None: + value = os.getenv("NVTE_CUDA_VERSION") + return normalize_cuda_version(value) if value is not None else None + + +def setuptools_hooks( + requirements: Callable[[ConfigSettings], List[str]], + private_settings: Iterable[str], + environment: Optional[Callable[[ConfigSettings], ContextManager]] = None, +) -> Dict[str, Callable]: + """Create PEP 517 and PEP 660 hooks that build with the legacy setuptools backend. + + Build requirements come only from `requirements`. The setuptools hooks for build + requirements are not used because they run setup.py, which can import build + requirements before they are installed. Private settings are removed before the + config settings are passed to setuptools, and builds run inside `environment`. + """ + backend = build_meta.__legacy__ + private_settings = frozenset(private_settings) + if environment is None: + environment = lambda config_settings: nullcontext() + + def setuptools_config_settings(config_settings: ConfigSettings) -> ConfigSettings: + if config_settings is None: + return None + settings = { + setting: value + for setting, value in config_settings.items() + if setting not in private_settings + } + return settings or None + + def get_requires_for_build_wheel(config_settings: ConfigSettings = None) -> List[str]: + """Get requirements for building a wheel.""" + return requirements(config_settings) + + def get_requires_for_build_sdist(config_settings: ConfigSettings = None) -> List[str]: + """Get requirements needed while evaluating setup.py for an sdist.""" + return requirements(config_settings) + + def build_wheel( + wheel_directory: str, + config_settings: ConfigSettings = None, + metadata_directory: Optional[str] = None, + ) -> str: + """Build a wheel with setuptools.""" + with environment(config_settings): + return backend.build_wheel( + wheel_directory, setuptools_config_settings(config_settings), metadata_directory + ) + + def prepare_metadata_for_build_wheel( + metadata_directory: str, config_settings: ConfigSettings = None + ) -> str: + """Prepare wheel metadata with setuptools.""" + with environment(config_settings): + return backend.prepare_metadata_for_build_wheel( + metadata_directory, setuptools_config_settings(config_settings) + ) + + def build_sdist(sdist_directory: str, config_settings: ConfigSettings = None) -> str: + """Build an sdist with setuptools.""" + with environment(config_settings): + return backend.build_sdist(sdist_directory, setuptools_config_settings(config_settings)) + + hooks = { + "get_requires_for_build_wheel": get_requires_for_build_wheel, + "get_requires_for_build_sdist": get_requires_for_build_sdist, + "build_wheel": build_wheel, + "prepare_metadata_for_build_wheel": prepare_metadata_for_build_wheel, + "build_sdist": build_sdist, + } + + if hasattr(backend, "build_editable"): + + def get_requires_for_build_editable(config_settings: ConfigSettings = None) -> List[str]: + """Get requirements for building an editable wheel.""" + return requirements(config_settings) + + def build_editable( + wheel_directory: str, + config_settings: ConfigSettings = None, + metadata_directory: Optional[str] = None, + ) -> str: + """Build an editable wheel with setuptools.""" + with environment(config_settings): + return backend.build_editable( + wheel_directory, setuptools_config_settings(config_settings), metadata_directory + ) + + def prepare_metadata_for_build_editable( + metadata_directory: str, config_settings: ConfigSettings = None + ) -> str: + """Prepare editable-wheel metadata with setuptools.""" + with environment(config_settings): + return backend.prepare_metadata_for_build_editable( + metadata_directory, setuptools_config_settings(config_settings) + ) + + hooks.update( + get_requires_for_build_editable=get_requires_for_build_editable, + build_editable=build_editable, + prepare_metadata_for_build_editable=prepare_metadata_for_build_editable, + ) + + return hooks diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index 4b7ec153be1..b6559b80d54 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -44,7 +44,12 @@ def __init__( self.cmake_path: Path = cmake_path self.cmake_flags: List[str] = [] if cmake_flags is None else cmake_flags - def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: + def _build_cmake( + self, + build_dir: Path, + install_dir: Path, + site_packages_rpath: Optional[str] = None, + ) -> None: # Make sure paths are str _cmake_bin = str(cmake_bin()) cmake_path = str(self.cmake_path) @@ -70,10 +75,19 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: if discovered_cuda_home is not None: configure_command.append(f"-DCUDAToolkit_ROOT={discovered_cuda_home}") - # CUDA wheels use `lib`, while toolkit installations typically use - # `lib64`. Only override CMake's library discovery for a wheel-style - # layout with the expected libraries present. cuda_lib_dir = discovered_cuda_home / "lib" + if cuda_lib_dir.is_dir(): + # NVCC searches `lib64` for its static runtime libraries, while + # CUDA wheels install them in `lib`. + cuda_static_libraries = ( + cuda_lib_dir / "libcudadevrt.a", + cuda_lib_dir / "libcudart_static.a", + ) + if all(library.is_file() for library in cuda_static_libraries): + configure_command.append(f"-DCMAKE_CUDA_FLAGS=-L{cuda_lib_dir}") + + # Override CMake's shared-library discovery for each library that is + # present in the wheel-style layout. if cuda_lib_dir.is_dir() and (cuda_full_version := cuda_version()): cuda_major_version = cuda_full_version[0] cuda_libraries = { @@ -82,11 +96,11 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: "CUDA_cublas_LIBRARY": cuda_lib_dir / f"libcublas.so.{cuda_major_version}", "CUDA_cublasLt_LIBRARY": cuda_lib_dir / f"libcublasLt.so.{cuda_major_version}", } - if all(library.is_file() for library in cuda_libraries.values()): - configure_command.append(f"-DCMAKE_CUDA_FLAGS=-L{cuda_lib_dir}") - configure_command.extend( - f"-D{variable}={library}" for variable, library in cuda_libraries.items() - ) + configure_command.extend( + f"-D{variable}={library}" + for variable, library in cuda_libraries.items() + if library.is_file() + ) discovered_nvcc_path = nvcc_path() if discovered_nvcc_path is not None: @@ -99,6 +113,8 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: f"-DCMAKE_CUDA_COMPILER_LAUNCHER={ccache_bin}", ] configure_command += self.cmake_flags + if site_packages_rpath is not None: + configure_command.append(f"-DNVTE_SITE_PACKAGES_RPATH={site_packages_rpath}") import pybind11 @@ -138,6 +154,19 @@ class _CMakeBuildExtension(extension_cls): """Setuptools command with support for CMake extension modules""" def run(self) -> None: + # Ensure that shared objects files for source and PyPI installations live + # in separate directories to avoid conflicts during install and runtime. + lib_dir = ( + "wheel_lib" + if bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))) or framework_extension_only + else "" + ) + + # Non-inplace builds move shared objects into site-packages/transformer_engine/. + site_packages_rpath = None + if not self.inplace: + site_packages_rpath = "$ORIGIN/../.." if lib_dir else "$ORIGIN/.." + # Build CMake extensions for ext in self.extensions: package_path = Path(self.get_ext_fullpath(ext.name)) @@ -158,6 +187,7 @@ def run(self) -> None: ext._build_cmake( build_dir=build_dir, install_dir=install_dir, + site_packages_rpath=site_packages_rpath, ) continue @@ -172,6 +202,7 @@ def run(self) -> None: ext._build_cmake( build_dir=Path(build_dir), install_dir=install_dir, + site_packages_rpath=site_packages_rpath, ) # Build non-CMake extensions as usual @@ -182,14 +213,6 @@ def run(self) -> None: super().run() self.extensions = all_extensions - # Ensure that shared objects files for source and PyPI installations live - # in separate directories to avoid conflicts during install and runtime. - lib_dir = ( - "wheel_lib" - if bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))) or framework_extension_only - else "" - ) - # Ensure that binaries are not in global package space. # For editable/inplace builds this is not a concern as # the SOs will be in a local directory anyway. diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 173f9cee708..5819c6937cf 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -17,6 +17,7 @@ debug_build_enabled, get_bolt_build_flags, nccl_ep_enabled, + nvshmem_home_path, setup_mpi_flags, ) from typing import List @@ -109,13 +110,18 @@ def setup_pytorch_extension( library_dirs = [] libraries = [] if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", 0))): - assert ( - os.getenv("NVSHMEM_HOME") is not None - ), "NVSHMEM_HOME must be set when compiling with NVTE_ENABLE_NVSHMEM=1" - nvshmem_home = Path(os.getenv("NVSHMEM_HOME")) + nvshmem_home = nvshmem_home_path() + assert nvshmem_home is not None, ( + "NVSHMEM_HOME must be set, or the nvidia-nvshmem-cu* Python package installed, when" + " compiling with NVTE_ENABLE_NVSHMEM=1" + ) include_dirs.append(nvshmem_home / "include") library_dirs.append(nvshmem_home / "lib") - libraries.append("nvshmem_host") + if (nvshmem_home / "lib" / "libnvshmem_host.so").is_file(): + libraries.append("nvshmem_host") + else: + # NVSHMEM Python packages only provide the versioned shared library. + libraries.append(":libnvshmem_host.so.3") cxx_flags.append("-DNVTE_ENABLE_NVSHMEM") if bool(int(os.getenv("NVTE_WITH_CUBLASMP", 0))): diff --git a/build_tools/utils.py b/build_tools/utils.py index 6841a595802..d4af10a55a5 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -366,6 +366,81 @@ def nccl_lib_path() -> Optional[Path]: return None +# Keep in sync with MIN_CUDNN_VERSION in transformer_engine/common/__init__.py. +MIN_CUDNN_VERSION = "9.12" + + +def _python_package_path(package_name: str, relative_path: str) -> Optional[Path]: + """Return a path inside an installed Python package, or `None` if it does not exist.""" + try: + package_distribution = distribution(package_name) + except PackageNotFoundError: + return None + + path = Path(package_distribution.locate_file(relative_path)) + return path if path.exists() else None + + +def cublasmp_home_path() -> Optional[Path]: + """Return the cuBLASMp installation root. + + Prefers CUBLASMP_HOME, then the nvidia-cublasmp-cu* Python package. + Returns `None` if cuBLASMp is not found.""" + if cublasmp_home := os.getenv("CUBLASMP_HOME"): + return Path(cublasmp_home) + + cuda_major = cuda_version()[0] + return _python_package_path( + f"nvidia-cublasmp-cu{cuda_major}", f"nvidia/cublasmp/cu{cuda_major}" + ) + + +def cusolvermp_home_path() -> Path: + """Return the cuSOLVERMp installation root. + + Prefers CUSOLVERMP_HOME, then the nvidia-cusolvermp-cu* Python package, + then `/usr`.""" + if cusolvermp_home := os.getenv("CUSOLVERMP_HOME"): + return Path(cusolvermp_home) + + cuda_major = cuda_version()[0] + package_path = _python_package_path( + f"nvidia-cusolvermp-cu{cuda_major}", f"nvidia/cu{cuda_major}" + ) + return package_path if package_path is not None else Path("/usr") + + +def nvshmem_home_path() -> Optional[Path]: + """Return the NVSHMEM installation root. + + Prefers NVSHMEM_HOME, then the nvidia-nvshmem-cu* Python package. + Returns `None` if NVSHMEM is not found.""" + if nvshmem_home := os.getenv("NVSHMEM_HOME"): + return Path(nvshmem_home) + + return _python_package_path(f"nvidia-nvshmem-cu{cuda_version()[0]}", "nvidia/nvshmem") + + +def mpi_home_path() -> Optional[Path]: + """Return the MPI installation root. + + Prefers MPI_HOME, then the openmpi Python package, which installs into the + environment prefix rather than site-packages. Returns `None` if MPI is not found.""" + if mpi_home := os.getenv("MPI_HOME"): + return Path(mpi_home) + + try: + openmpi_distribution = distribution("openmpi") + except PackageNotFoundError: + return None + + for package_path in openmpi_distribution.files or []: + if package_path.name == "mpi.h" and package_path.parent.name == "include": + return Path(openmpi_distribution.locate_file(package_path)).resolve().parent.parent + + return None + + @functools.lru_cache(maxsize=None) def cuda_toolkit_include_path() -> Optional[Path]: """Return the CUDA Toolkit include directory.""" @@ -376,31 +451,45 @@ def cuda_toolkit_include_path() -> Optional[Path]: return None +def _nvidia_package_roots() -> List[Path]: + """Return the directories of the `nvidia` namespace package for NVIDIA Python packages.""" + try: + import nvidia + except ModuleNotFoundError: + return [] + return [Path(path).resolve() for path in nvidia.__path__] + + +def cuda_toolkit_from_python_packages() -> bool: + """Whether the CUDA Toolkit used for the build is installed from Python packages.""" + cuda_home = cuda_home_path() + return cuda_home is not None and any( + cuda_home.is_relative_to(root) for root in _nvidia_package_roots() + ) + + @functools.lru_cache(maxsize=None) -def get_cuda_include_dirs() -> Tuple[str, str]: - """Returns the CUDA header directory.""" +def get_cuda_include_dirs() -> List[Path]: + """Returns the CUDA header directories.""" force_wheels = bool(int(os.getenv("NVTE_BUILD_USE_NVIDIA_WHEELS", "0"))) - # If cuda is installed via toolkit, all necessary headers - # are bundled inside the top level cuda directory. - if not force_wheels and cuda_toolkit_include_path() is not None: - return [cuda_toolkit_include_path()] + nvidia_roots = _nvidia_package_roots() - # Use pip wheels to include all headers. - try: - import nvidia - except ModuleNotFoundError as e: - raise RuntimeError("CUDA not found.") + # A system toolkit contains all necessary headers under one include + # directory. NVIDIA wheels split them across namespace packages. + if not force_wheels and not cuda_toolkit_from_python_packages(): + if (include_path := cuda_toolkit_include_path()) is not None: + return [include_path] - if nvidia.__file__ is not None: - cuda_root = Path(nvidia.__file__).parent - else: - cuda_root = Path(nvidia.__path__[0]) # namespace - return [ + include_dirs = [ subdir / "include" - for subdir in cuda_root.iterdir() + for root in nvidia_roots + for subdir in root.iterdir() if subdir.is_dir() and (subdir / "include").is_dir() ] + if not include_dirs: + raise RuntimeError("CUDA not found.") + return include_dirs @functools.lru_cache(maxsize=None) @@ -551,10 +640,11 @@ def get_frameworks() -> List[str]: def setup_mpi_flags(include_dirs: List, cxx_flags: List) -> None: """Add MPI include path and compile definition if NVTE_UB_WITH_MPI is enabled.""" if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))): - assert ( - os.getenv("MPI_HOME") is not None - ), "MPI_HOME=/path/to/mpi must be set when compiling with NVTE_UB_WITH_MPI=1!" - mpi_path = Path(os.getenv("MPI_HOME")) + mpi_path = mpi_home_path() + assert mpi_path is not None, ( + "MPI_HOME=/path/to/mpi must be set, or the openmpi Python package installed, " + "when compiling with NVTE_UB_WITH_MPI=1!" + ) include_dirs.append(mpi_path / "include") cxx_flags.append("-DNVTE_UB_WITH_MPI") diff --git a/build_tools/wheel_utils/build_wheels.sh b/build_tools/wheel_utils/build_wheels.sh index 761f12390e1..72467caf200 100644 --- a/build_tools/wheel_utils/build_wheels.sh +++ b/build_tools/wheel_utils/build_wheels.sh @@ -23,8 +23,9 @@ cd /TransformerEngine git checkout $TARGET_BRANCH git submodule update --init --recursive -# Install deps -/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel 'nvidia-cudnn-frontend>=1.25.0' 'apache-tvm-ffi>=0.1.12' +# Install deps. The core library wheel is built in an isolated environment, but the +# metapackage, wheel repacking, and framework sdists still use this interpreter. +/opt/python/cp310-cp310/bin/pip install pybind11[global] setuptools wheel if $BUILD_METAPACKAGE ; then cd /TransformerEngine @@ -44,8 +45,12 @@ if $BUILD_COMMON ; then WHEEL_ABI_TAG=none fi - # Create the wheel. - /opt/python/cp310-cp310/bin/python setup.py bdist_wheel --verbose --python-tag=$WHEEL_PYTHON_TAG --plat-name=$PLATFORM 2>&1 | tee /wheelhouse/logs/common.txt + # Create the wheel in an isolated build environment. build_tools/build_backend.py + # installs the CUDA and optional feature (NVTE_WITH_*) build requirements. + /opt/python/cp310-cp310/bin/python -m pip wheel . --no-deps --wheel-dir dist --verbose \ + --config-settings framework=none \ + --config-settings=--build-option=--python-tag=$WHEEL_PYTHON_TAG \ + --config-settings=--build-option=--plat-name=$PLATFORM 2>&1 | tee /wheelhouse/logs/common.txt # Repack the wheel for specific cuda version. /opt/python/cp310-cp310/bin/wheel unpack dist/* diff --git a/docs/envvars.rst b/docs/envvars.rst index 5ff430eabea..e9a3483472d 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -39,8 +39,8 @@ Build Configuration .. envvar:: NVTE_FRAMEWORK :Type: ``str`` - :Default: Auto-detected - :Description: Comma-separated list of frameworks to build support for (``pytorch``, ``jax``, ``all``, or ``none``). If not specified, automatically detects installed frameworks. + :Default: Auto-detected without build isolation; ``none`` with build isolation + :Description: Comma-separated list of frameworks to build support for (``pytorch``, ``jax``, ``all``, or ``none``). Isolated PEP 517 builds cannot inspect frameworks in the caller's environment, so they only build the core library unless frameworks are selected through this variable or pip's ``--config-settings framework=`` option. The config setting takes precedence over this variable. .. envvar:: NVTE_USE_CCACHE @@ -78,27 +78,53 @@ Build Configuration :Default: ``10`` :Description: Number of Philox4x32 rounds used by stochastic rounding kernels. Must be a positive integer. +.. envvar:: NVTE_CUDA_VERSION + + :Type: ``str`` + :Default: Auto-detected based on system CUDA version + :Description: CUDA version to build for. Valid values: ``"12.*"``, ``"13.*"``. If not set, automatically determined based on the installed CUDA Toolkit version. Isolated builds fail if no CUDA Toolkit is found and neither this variable nor pip's ``--config-settings cuda-version=.`` option is set. When the JAX extension package is built with build isolation, this variable or the config setting selects the CUDA support of the JAX installed in the build environment, which defaults to CUDA 12. + Optional Dependencies ^^^^^^^^^^^^^^^^^^^^^ +Each of the following options can also be set for an isolated PEP 517 build with pip's +``--config-settings =<0 or 1>`` option, which takes precedence over the variable. +In an isolated build, the Python package listed for an option is installed as a build +requirement unless the option's installation directory variable is set. Without build +isolation, the installation directory variable is used if set, and the package is used if +installed. ``libtransformer_engine.so`` finds these packages at runtime when they are installed +in the same environment. + .. envvar:: NVTE_UB_WITH_MPI :Type: ``int`` (0 or 1) :Default: ``0`` - :Description: Enable MPI support for userbuffers. When set to ``1``, requires ``MPI_HOME`` to be set to the MPI installation directory. + :Description: Enable MPI support for userbuffers. Requires the MPI installation directory in ``MPI_HOME`` or the ``openmpi`` Python package. Config setting: ``ub-with-mpi``. .. envvar:: NVTE_ENABLE_NVSHMEM :Type: ``int`` (0 or 1) :Default: ``0`` - :Description: Enable NVSHMEM support. When set to ``1``, requires ``NVSHMEM_HOME`` to be set to the NVSHMEM installation directory. + :Description: Enable NVSHMEM support. Requires the NVSHMEM installation directory in ``NVSHMEM_HOME`` or the ``nvidia-nvshmem-cu`` Python package. Config setting: ``enable-nvshmem``. + +.. envvar:: NVTE_WITH_CUBLASMP + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Use cuBLASMp for tensor-parallel GEMMs. Requires cuBLASMp 0.8.1 or newer, from the installation directory in ``CUBLASMP_HOME`` or the ``nvidia-cublasmp-cu`` Python package, and NCCL 2.30 or newer. Config setting: ``with-cublasmp``. + +.. envvar:: NVTE_WITH_CUSOLVERMP + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Use cuSOLVERMp for distributed Newton-Schulz. Requires cuSOLVERMp from the installation directory in ``CUSOLVERMP_HOME``, the ``nvidia-cusolvermp-cu`` Python package, or ``/usr``, in that order. Config setting: ``with-cusolvermp``. .. envvar:: NVTE_WITH_CUTEDSL :Type: ``int`` (0 or 1) :Default: ``0`` :Description: Enable CuTeDSL support when building ``libtransformer_engine.so``. When set to ``1``, requires - CPython Stable ABIs from Python 3.10 and above, and ``apache-tvm-ffi`` installed. + CPython Stable ABIs from Python 3.10 and above, and ``apache-tvm-ffi`` installed. Config setting: ``with-cutedsl``. Note: when linking ``libtransformer_engine.so`` built with ``NVTE_WITH_CUTEDSL=1`` to a standalone C++ executable, you must provide the definition of these CPython symbols. See ``tests/cpp/CMakeLists.txt`` for an example. diff --git a/docs/installation.rst b/docs/installation.rst index 9d3eada827c..791f47929b7 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -36,7 +36,7 @@ Transformer Engine can be directly installed from `our PyPI build-constraints.txt + + # JAX + python3 -c "import jax, jaxlib; print(f'jax=={jax.__version__}'); print(f'jaxlib=={jaxlib.__version__}')" > build-constraints.txt + +If the framework was installed from a separate package index, e.g. the PyTorch index for a specific CUDA version, pass the same `--index-url` or `--extra-index-url` options to pip so that the build environment can find the same version. Framework builds that are not published to a package index, such as the ones in the NGC containers, cannot be used with build isolation. Use `--no-build-isolation` for these. + +Installation (from PyPI) +^^^^^^^^^^^^^^^^^^^^^^^^ + +The core library is shipped as wheels and does not need to be built. Execute the following command to install the PyTorch extension with build isolation: + +.. code-block:: bash + + pip3 install --build-constraint build-constraints.txt transformer_engine[pytorch] + +The JAX extension is installed in the same way: + +.. code-block:: bash + + pip3 install --build-constraint build-constraints.txt transformer_engine[jax] + +With build isolation, the JAX extension is built against JAX with CUDA 12 support and depends on the core library compiled for CUDA 12 by default. For JAX with CUDA 13 support, set the `cuda-version` config setting: + +.. code-block:: bash + + pip3 install --build-constraint build-constraints.txt --config-settings cuda-version=13.0 transformer_engine[jax] + +Installation (from GitHub or source) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +With build isolation, frameworks installed in the current environment are not detected. To build Transformer Engine support for them, set the `framework` config setting to a comma-separated list (e.g. `framework=jax,pytorch`). If no frameworks are specified, only the framework-agnostic C++ API is built. + +.. code-block:: bash + + cd TransformerEngine + pip3 install --build-constraint build-constraints.txt --config-settings framework=pytorch . + +The same options can be used to install from GitHub: + +.. code-block:: bash + + pip3 install --build-constraint build-constraints.txt --config-settings framework=pytorch git+https://github.com/NVIDIA/TransformerEngine.git@stable + +Unless `CUDA_HOME` is set, the CUDA Toolkit and cuDNN for the build are installed into the build environment from PyPI, matching the version of the installed CUDA Toolkit. If no CUDA Toolkit is installed, specify the CUDA version with the `cuda-version` config setting (e.g. `cuda-version=13.0`). These packages are only used for the build, so the CUDA libraries and cuDNN must also be installed at runtime, e.g. as dependencies of PyTorch or JAX. When building framework support, choose the CUDA version that PyTorch or JAX was built with. To build with an installed CUDA Toolkit and cuDNN instead, set `CUDA_HOME` to the location of the CUDA Toolkit. + +Optional dependencies, such as cuBLASMp and NVSHMEM, are enabled with config settings as well (e.g. `with-cublasmp=1`). In an isolated build, their Python packages are installed into the build environment. To use the optional dependencies at runtime, install the same Python packages in the environment where Transformer Engine is installed. Each config setting has an equivalent environment variable, e.g. `NVTE_FRAMEWORK` for `framework` and `NVTE_CUDA_VERSION` for `cuda-version`. See :doc:`envvars` for the complete list. + +Other installers support build constraints as well. For example, with `uv`: + +.. code-block:: bash + + uv pip install --build-constraints build-constraints.txt --config-setting framework=pytorch . + .. include:: troubleshooting.rst diff --git a/pyproject.toml b/pyproject.toml index 2f5f8979703..7be88f7988b 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,19 @@ # See LICENSE for license information. [build-system] -requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.29.0"] +requires = [ + "setuptools>=61.0", + # CMake 3.25.0 breaks FindCUDAToolkit when nvptxcompiler_static is absent from CUDA wheels. + "cmake>=3.21,!=3.25.0", + "wheel", + "pybind11[global]", + "ninja", + "nvidia-cudnn-frontend>=1.29.0", +] -# Use legacy backend to import local packages in setup.py -build-backend = "setuptools.build_meta:__legacy__" +# CUDA, framework, and optional feature requirements are defined in +# `build_tools/build_backend.py`. Isolated builds only build framework support +# selected with --config-settings framework= or NVTE_FRAMEWORK. Select CUDA with +# --config-settings cuda-version=. +build-backend = "build_tools.build_backend" +backend-path = ["."] diff --git a/setup.py b/setup.py index 2dabf37b322..020902aa042 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ import subprocess import time from pathlib import Path -from typing import List, Tuple +from typing import List, Optional, Tuple import setuptools from wheel.bdist_wheel import bdist_wheel @@ -18,17 +18,21 @@ from build_tools.build_ext import CMakeExtension, get_build_ext from build_tools.te_version import te_version from build_tools.utils import ( + cublasmp_home_path, cuda_archs, cuda_home_path, - cuda_version, + cuda_toolkit_from_python_packages, cudnn_frontend_include_path, + cusolvermp_home_path, get_bolt_build_flags, get_frameworks, remove_dups, min_python_version_str, + mpi_home_path, nccl_ep_enabled, get_max_jobs_for_parallel_build, nvcc_path, + nvshmem_home_path, ) frameworks = get_frameworks() @@ -75,31 +79,38 @@ def setup_common_extension() -> CMakeExtension: cmake_flags.append("-DNVTE_WITH_CUTEDSL=OFF") if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))): - assert ( - os.getenv("MPI_HOME") is not None - ), "MPI_HOME must be set when compiling with NVTE_UB_WITH_MPI=1" + mpi_home = mpi_home_path() + assert mpi_home is not None, ( + "MPI_HOME must be set, or the openmpi Python package installed, when compiling with" + " NVTE_UB_WITH_MPI=1" + ) cmake_flags.append("-DNVTE_UB_WITH_MPI=ON") + cmake_flags.append(f"-DMPI_HOME={mpi_home}") if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", "0"))): - assert ( - os.getenv("NVSHMEM_HOME") is not None - ), "NVSHMEM_HOME must be set when compiling with NVTE_ENABLE_NVSHMEM=1" + nvshmem_home = nvshmem_home_path() + assert nvshmem_home is not None, ( + "NVSHMEM_HOME must be set, or the nvidia-nvshmem-cu* Python package installed, when" + " compiling with NVTE_ENABLE_NVSHMEM=1" + ) cmake_flags.append("-DNVTE_ENABLE_NVSHMEM=ON") + cmake_flags.append(f"-DNVSHMEM_HOME={nvshmem_home}") if bool(int(os.getenv("NVTE_BUILD_ACTIVATION_WITH_FAST_MATH", "0"))): cmake_flags.append("-DNVTE_BUILD_ACTIVATION_WITH_FAST_MATH=ON") if bool(int(os.getenv("NVTE_WITH_CUBLASMP", "0"))): + cublasmp_dir = cublasmp_home_path() + assert cublasmp_dir is not None, ( + "CUBLASMP_HOME must be set, or the nvidia-cublasmp-cu* Python package installed, when" + " compiling with NVTE_WITH_CUBLASMP=1" + ) cmake_flags.append("-DNVTE_WITH_CUBLASMP=ON") - cublasmp_dir = os.getenv("CUBLASMP_HOME") or metadata.distribution( - f"nvidia-cublasmp-cu{cuda_version()[0]}" - ).locate_file(f"nvidia/cublasmp/cu{cuda_version()[0]}") cmake_flags.append(f"-DCUBLASMP_DIR={cublasmp_dir}") if bool(int(os.getenv("NVTE_WITH_CUSOLVERMP", "0"))): cmake_flags.append("-DNVTE_WITH_CUSOLVERMP=ON") - cusolvermp_dir = os.getenv("CUSOLVERMP_HOME", "/usr") - cmake_flags.append(f"-DCUSOLVERMP_DIR={cusolvermp_dir}") + cmake_flags.append(f"-DCUSOLVERMP_DIR={cusolvermp_home_path()}") # NCCL EP (Hopper+): on by default; auto-skipped when no arch >= 90 is # targeted. Set NVTE_WITH_NCCL_EP=0 to force off. @@ -170,15 +181,21 @@ def _discover_nccl_home() -> str: f"'{env_home}/include/nccl.h' was not found; falling back to system probes." ) + # With a CUDA Toolkit from Python packages, the NCCL loaded at runtime is most likely the + # nvidia-nccl wheel, which libtransformer_engine's RUNPATH prefers when it is installed. + if cuda_toolkit_from_python_packages() and (pip_root := _python_package_nccl_home()): + return pip_root + lib_names = ("libnccl.so", "libnccl.so.2") # Include Debian/Ubuntu multiarch subdirs (e.g. lib/aarch64-linux-gnu). lib_subdirs = ("lib", "lib64", "lib/aarch64-linux-gnu", "lib/x86_64-linux-gnu") # Prefer the NCCL the dynamic loader will actually resolve at runtime so the - # EP build links against the same libnccl that gets loaded. libtransformer_engine - # carries no NCCL RUNPATH, so the loader uses ldconfig/system paths; building - # against a different NCCL (e.g. a pip wheel) causes ABI mismatches. ldconfig is - # the ground truth for runtime resolution, so consult it before well-known prefixes. + # EP build links against the same libnccl that gets loaded. Unless an nvidia-nccl + # wheel is installed next to it, which libtransformer_engine's RUNPATH prefers, + # the loader uses ldconfig/system paths; building against a different NCCL causes + # ABI mismatches. ldconfig is the ground truth for runtime resolution, so consult + # it before well-known prefixes. try: out = subprocess.check_output(["ldconfig", "-p"], stderr=subprocess.DEVNULL).decode() for line in out.splitlines(): @@ -199,24 +216,30 @@ def _discover_nccl_home() -> str: ): return str(p) - # Fall back to a pip-installed NCCL (nvidia-nccl-cu* wheel) under nvidia/nccl - # in site-packages, used only when no system NCCL is present. + # Fall back to a pip-installed NCCL when no system NCCL is present. + if pip_root := _python_package_nccl_home(): + return pip_root + + raise RuntimeError( + "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix." + ) + + +def _python_package_nccl_home() -> Optional[str]: + """Return the NCCL installed from an nvidia-nccl-cu* wheel under nvidia/nccl, if any.""" try: import importlib.util spec = importlib.util.find_spec("nvidia.nccl") - if spec is not None and spec.submodule_search_locations: - pip_root = Path(next(iter(spec.submodule_search_locations))) - if (pip_root / "include" / "nccl.h").exists() and any( - (pip_root / sub / name).exists() for sub in lib_subdirs for name in lib_names - ): - return str(pip_root) except (ImportError, ValueError): - pass - - raise RuntimeError( - "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix." - ) + return None + if spec is None or not spec.submodule_search_locations: + return None + + pip_root = Path(next(iter(spec.submodule_search_locations))) + if (pip_root / "include" / "nccl.h").exists() and (pip_root / "lib" / "libnccl.so.2").exists(): + return str(pip_root) + return None def build_nccl_ep_submodule() -> str: diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index f90f74449c5..b1e3608e84f 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -534,9 +534,11 @@ if (NVTE_WITH_CUBLASMP) PATHS ${CUBLASMP_DIR} PATH_SUFFIXES lib lib64 lib/aarch64-linux-gnu lib/sbsa-linux-gnu lib/x86_64-linux-gnu REQUIRED) + # NCCL Python packages only provide the versioned shared library. find_library(NCCL_LIB - NAMES nccl libnccl - PATH_SUFFIXES lib + NAMES nccl libnccl.so.2 + HINTS "${NCCL_INCLUDE_DIR}/.." + PATH_SUFFIXES lib lib64 REQUIRED) # cuBLASMp 0.8.1 adds MXFP8 support to the CUDA-graph-safe overlap algos, # and NCCL 2.30 is the first release with graph-safe one-sided RMA @@ -561,14 +563,16 @@ option(NVTE_WITH_CUSOLVERMP "Use cuSolverMp for distributed Newton-Schulz" OFF) if (NVTE_WITH_CUSOLVERMP) target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_CUSOLVERMP) target_include_directories(transformer_engine PRIVATE ${CUSOLVERMP_DIR}/include) + # cuSOLVERMp and NCCL Python packages only provide the versioned shared libraries. find_library(CUSOLVERMP_LIB - NAMES cusolverMp libcusolverMp + NAMES cusolverMp libcusolverMp.so.0 PATHS ${CUSOLVERMP_DIR} PATH_SUFFIXES lib REQUIRED) find_library(NCCL_LIB - NAMES nccl libnccl - PATH_SUFFIXES lib + NAMES nccl libnccl.so.2 + HINTS "${NCCL_INCLUDE_DIR}/.." + PATH_SUFFIXES lib lib64 REQUIRED) target_link_libraries(transformer_engine PRIVATE ${NCCL_LIB} ${CUSOLVERMP_LIB}) message(STATUS "Using cuSolverMp at: ${CUSOLVERMP_DIR}") @@ -608,8 +612,10 @@ file(REAL_PATH "${NCCL_EP_SHARED_LINK}" NCCL_EP_SHARED_LIB) # -- NCCL core library ------------------------------------------------------- if(NOT NCCL_LIB) + # NCCL Python packages only provide the versioned shared library. find_library(NCCL_LIB - NAMES nccl libnccl + NAMES nccl libnccl.so.2 + HINTS "${NCCL_INCLUDE_DIR}/.." PATH_SUFFIXES lib lib64 REQUIRED) endif() @@ -838,6 +844,34 @@ endif() set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --threads ${BUILD_THREADS_PER_JOB}") message(STATUS "Threads per parallel build job: ${BUILD_THREADS_PER_JOB}") +# Find optional dependencies that are installed from Python packages into the +# same environment as the installed library. +set(NVTE_SITE_PACKAGES_RPATH "" CACHE STRING + "Python site-packages directory relative to the installed library, e.g. \$ORIGIN/..") +if (NVTE_SITE_PACKAGES_RPATH) + set(nvte_python_package_rpath) + if (NVTE_WITH_CUBLASMP) + list(APPEND nvte_python_package_rpath + "${NVTE_SITE_PACKAGES_RPATH}/nvidia/cublasmp/cu${CUDAToolkit_VERSION_MAJOR}/lib") + endif() + if (NVTE_WITH_CUSOLVERMP) + list(APPEND nvte_python_package_rpath + "${NVTE_SITE_PACKAGES_RPATH}/nvidia/cu${CUDAToolkit_VERSION_MAJOR}/lib") + endif() + if (NVTE_WITH_NCCL_EP OR NVTE_WITH_CUBLASMP OR NVTE_WITH_CUSOLVERMP) + list(APPEND nvte_python_package_rpath "${NVTE_SITE_PACKAGES_RPATH}/nvidia/nccl/lib") + endif() + if (NVTE_ENABLE_NVSHMEM) + list(APPEND nvte_python_package_rpath "${NVTE_SITE_PACKAGES_RPATH}/nvidia/nvshmem/lib") + endif() + if (NVTE_UB_WITH_MPI) + # The openmpi package installs into /lib, two levels above site-packages. + list(APPEND nvte_python_package_rpath "${NVTE_SITE_PACKAGES_RPATH}/../..") + endif() + set_target_properties(transformer_engine PROPERTIES + INSTALL_RPATH "${nvte_python_package_rpath}") +endif() + # Install library install(TARGETS transformer_engine DESTINATION .) if(NVTE_WITH_NCCL_EP) diff --git a/transformer_engine/common/nvshmem_api/CMakeLists.txt b/transformer_engine/common/nvshmem_api/CMakeLists.txt index 3d9b6b5ec44..54d22732485 100644 --- a/transformer_engine/common/nvshmem_api/CMakeLists.txt +++ b/transformer_engine/common/nvshmem_api/CMakeLists.txt @@ -13,8 +13,15 @@ set(NVSHMEM_HOME "$ENV{NVSHMEM_HOME}" CACHE STRING "Location of NVSHMEM installa add_library(nvshmemapi STATIC nvshmem_waitkernel.cu) set(NVSHMEMAPI_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}" PARENT_SCOPE) +# NVSHMEM Python packages only provide the versioned host library. +find_library(NVSHMEM_HOST_LIB + NAMES nvshmem_host libnvshmem_host.so.3 + PATHS ${NVSHMEM_HOME} + PATH_SUFFIXES lib lib64 + NO_DEFAULT_PATH + REQUIRED) target_link_directories(nvshmemapi PUBLIC ${NVSHMEM_HOME}/lib) -target_link_libraries(nvshmemapi PUBLIC -static-libstdc++ nvshmem_device nvshmem_host CUDA::nvml CUDA::cublas CUDA::cuda_driver) +target_link_libraries(nvshmemapi PUBLIC -static-libstdc++ nvshmem_device ${NVSHMEM_HOST_LIB} CUDA::nvml CUDA::cublas CUDA::cuda_driver) target_include_directories(nvshmemapi PRIVATE ${NVSHMEM_HOME}/include/ ${CMAKE_CURRENT_SOURCE_DIR}/../include) diff --git a/transformer_engine/jax/MANIFEST.in b/transformer_engine/jax/MANIFEST.in index 0c814f95da5..4924c920f51 100644 --- a/transformer_engine/jax/MANIFEST.in +++ b/transformer_engine/jax/MANIFEST.in @@ -1,3 +1,4 @@ +include _build_backend.py recursive-include build_tools *.* recursive-include common_headers *.* recursive-include csrc *.* diff --git a/transformer_engine/jax/_build_backend.py b/transformer_engine/jax/_build_backend.py new file mode 100644 index 00000000000..7a981eefb07 --- /dev/null +++ b/transformer_engine/jax/_build_backend.py @@ -0,0 +1,41 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PEP 517 backend that builds against JAX with support for the requested CUDA version.""" + +from pathlib import Path +import sys +from typing import List + +# Like setup.py, prefer the repository's build_tools in a source checkout. The sdist contains +# a copy next to this file. +_REPOSITORY_ROOT = Path(__file__).resolve().parent.parent.parent +if (_REPOSITORY_ROOT / "build_tools").is_dir(): + sys.path.insert(0, str(_REPOSITORY_ROOT)) + +# pylint: disable=wrong-import-position +from build_tools.build_backend_utils import ( + ConfigSettings, + requested_cuda_version, + setuptools_hooks, +) + +# Isolated builds have always used JAX with CUDA 12 support. +_DEFAULT_CUDA_MAJOR_VERSION = "12" + + +def _build_requirements(config_settings: ConfigSettings) -> List[str]: + """Get JAX with support for the requested CUDA version.""" + # setup.py detects the CUDA version of the installed JAX to select the core library. + cuda_version = requested_cuda_version(config_settings) + cuda_major = ( + cuda_version.split(".", maxsplit=1)[0] if cuda_version else _DEFAULT_CUDA_MAJOR_VERSION + ) + return [f"jax[cuda{cuda_major}]"] + + +# PEP 517 and PEP 660 hooks +globals().update( + setuptools_hooks(requirements=_build_requirements, private_settings=("cuda-version",)) +) diff --git a/transformer_engine/jax/pyproject.toml b/transformer_engine/jax/pyproject.toml index d5c15b4a8fb..4c9648430cf 100755 --- a/transformer_engine/jax/pyproject.toml +++ b/transformer_engine/jax/pyproject.toml @@ -3,7 +3,9 @@ # See LICENSE for license information. [build-system] -requires = ["setuptools>=61.0", "pybind11[global]", "pip", "jax[cuda12]", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"] +requires = ["setuptools>=61.0", "pybind11[global]", "pip", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"] -# Use legacy backend to import local packages in setup.py -build-backend = "setuptools.build_meta:__legacy__" +# JAX with CUDA support is added in `_build_backend.py`. Select CUDA with +# --config-settings cuda-version= or NVTE_CUDA_VERSION (default: CUDA 12). +build-backend = "_build_backend" +backend-path = ["."]