diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index cbb8838b00..e8a60f91e7 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -19,6 +19,8 @@ from .utils import ( cmake_bin, + cuda_home_path, + cuda_version, debug_build_enabled, found_ninja, get_frameworks, @@ -61,6 +63,33 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: f"-DCMAKE_BUILD_TYPE={build_type}", f"-DCMAKE_INSTALL_PREFIX={install_dir}", ] + + discovered_cuda_home = cuda_home_path() + 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() and (cuda_full_version := cuda_version()): + cuda_major_version = cuda_full_version[0] + cuda_libraries = { + "CUDA_CUDART": cuda_lib_dir / f"libcudart.so.{cuda_major_version}", + "CUDA_cudart_LIBRARY": cuda_lib_dir / f"libcudart.so.{cuda_major_version}", + "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() + ) + + discovered_nvcc_path = nvcc_path() + if discovered_nvcc_path is not None: + configure_command.append(f"-DCMAKE_CUDA_COMPILER={discovered_nvcc_path}") + if bool(int(os.getenv("NVTE_USE_CCACHE", "0"))): ccache_bin = os.getenv("NVTE_CCACHE_BIN", "ccache") configure_command += [ @@ -185,6 +214,11 @@ def _compile_fn(obj, src, ext, cc_args, extra_postargs, pp_opts) -> None: and not framework_extension_only ): nvcc_bin = nvcc_path() + if nvcc_bin is None: + raise RuntimeError( + f"NVCC not found and is required for building CUDA source {src}" + ) + self.compiler.set_executable("compiler_so", str(nvcc_bin)) if isinstance(cflags, dict): cflags = cflags["nvcc"] diff --git a/build_tools/jax.py b/build_tools/jax.py index d61c26c128..031432e6f9 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -16,6 +16,8 @@ cudnn_frontend_include_path, debug_build_enabled, setup_mpi_flags, + nccl_include_path, + nccl_lib_path, nccl_ep_enabled, ) from typing import List @@ -90,6 +92,8 @@ def setup_jax_extension( # Header files include_dirs = get_cuda_include_dirs() + if (discovered_nccl_include_path := nccl_include_path()) is not None: + include_dirs.append(discovered_nccl_include_path) include_dirs.append(cudnn_frontend_include_path()) include_dirs.extend( [ @@ -117,6 +121,12 @@ def setup_jax_extension( if nccl_ep_enabled(): cxx_flags.append("-DNVTE_WITH_NCCL_EP") + kwargs = {} + if (discovered_nccl_lib_path := nccl_lib_path()) is not None: + kwargs["extra_objects"] = [str(discovered_nccl_lib_path)] + else: + kwargs["libraries"] = ["nccl"] + # Define TE/JAX as a Pybind11Extension from pybind11.setup_helpers import Pybind11Extension @@ -125,5 +135,5 @@ def setup_jax_extension( sources=[str(path) for path in sources], include_dirs=[str(path) for path in include_dirs], extra_compile_args=cxx_flags, - libraries=["nccl"], + **kwargs, ) diff --git a/build_tools/utils.py b/build_tools/utils.py index 6a8168e7d0..3b02fcea2b 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -16,7 +16,7 @@ from pathlib import Path from importlib.metadata import PackageNotFoundError, distribution, version as get_version from subprocess import CalledProcessError -from typing import List, Optional, Tuple, Union +from typing import Callable, List, Optional, Tuple, Union # Needs to stay consistent with .pre-commit-config.yaml config. @@ -176,52 +176,137 @@ def found_pybind11() -> bool: @functools.lru_cache(maxsize=None) -def cuda_toolkit_include_path() -> Tuple[str, str]: - """Returns root path for cuda toolkit includes. - - return `None` if CUDA is not found.""" - # Try finding CUDA - cuda_home: Optional[Path] = None - if cuda_home is None and os.getenv("CUDA_HOME"): - # Check in CUDA_HOME - cuda_home = Path(os.getenv("CUDA_HOME")) / "include" - if cuda_home is None: - # Check in NVCC - nvcc_bin = shutil.which("nvcc") - if nvcc_bin is not None: - cuda_home = Path(nvcc_bin.rstrip("/bin/nvcc")) / "include" - if cuda_home is None: - # Last-ditch guess in /usr/local/cuda - if Path("/usr/local/cuda").is_dir(): - cuda_home = Path("/usr/local/cuda") / "include" - return cuda_home +def nvcc_path() -> Optional[Path]: + """Get the NVCC binary path. + + Returns `None` if NVCC is not found. + """ + + def lookup_via_cuda_home() -> Optional[Path]: + if cuda_home := os.getenv("CUDA_HOME"): + return Path(cuda_home) / "bin" / "nvcc" + return None + + def lookup_via_python_path() -> Optional[Path]: + for python_path in sys.path: + nvidia_dir = Path(python_path) / "nvidia" + if not nvidia_dir.is_dir(): + continue + + cuda_roots = [nvidia_dir] + cuda_version_dirs = [ + path + for path in nvidia_dir.iterdir() + if path.is_dir() and re.fullmatch(r"cu\d+", path.name) + ] + cuda_version_dirs.sort(key=lambda path: int(path.name[2:]), reverse=True) + cuda_roots.extend(cuda_version_dirs) + + for cuda_root in cuda_roots: + nvcc_bin = cuda_root / "bin" / "nvcc" + if nvcc_bin.is_file(): + return nvcc_bin + + return None + + def lookup_via_distribution() -> Optional[Path]: + try: + cuda_nvcc_distribution = distribution("nvidia-cuda-nvcc") + except PackageNotFoundError: + return None + + for package_path in cuda_nvcc_distribution.files or []: + package_path = Path(package_path) + if package_path.name == "nvcc" and package_path.parent.name == "bin": + return Path(cuda_nvcc_distribution.locate_file(package_path)) + + return None + + def lookup_via_path() -> Optional[Path]: + if (nvcc_bin := shutil.which("nvcc")) is not None: + return Path(nvcc_bin) + return None + + def lookup_via_local_cuda() -> Path: + return Path("/usr/local/cuda/bin/nvcc") + + nvcc_lookup_funcs: List[Callable[[], Optional[Path]]] = [ + lookup_via_cuda_home, + lookup_via_python_path, + lookup_via_distribution, + lookup_via_path, + lookup_via_local_cuda, + ] + + for nvcc_lookup_func in nvcc_lookup_funcs: + if (nvcc_bin := nvcc_lookup_func()) is not None and nvcc_bin.is_file(): + return nvcc_bin.resolve() + + return None @functools.lru_cache(maxsize=None) -def nvcc_path() -> Tuple[str, str]: - """Returns the NVCC binary path. - - Throws FileNotFoundError if NVCC is not found.""" - # Try finding NVCC - nvcc_bin: Optional[Path] = None - if nvcc_bin is None and os.getenv("CUDA_HOME"): - # Check in CUDA_HOME - cuda_home = Path(os.getenv("CUDA_HOME")) - nvcc_bin = cuda_home / "bin" / "nvcc" - if nvcc_bin is None: - # Check if nvcc is in path - nvcc_bin = shutil.which("nvcc") - if nvcc_bin is not None: - cuda_home = Path(nvcc_bin.rstrip("/bin/nvcc")) - nvcc_bin = Path(nvcc_bin) - if nvcc_bin is None: - # Last-ditch guess in /usr/local/cuda - cuda_home = Path("/usr/local/cuda") - nvcc_bin = cuda_home / "bin" / "nvcc" - if not nvcc_bin.is_file(): - raise FileNotFoundError(f"Could not find NVCC at {nvcc_bin}") - - return nvcc_bin +def cuda_home_path() -> Optional[Path]: + """Return the CUDA Toolkit root containing NVCC.""" + if (nvcc_bin := nvcc_path()) is not None: + return nvcc_bin.parent.parent + return None + + +@functools.lru_cache(maxsize=None) +def nccl_root_path() -> Optional[Path]: + """Return the NCCL installation root. + + Returns `None` if NCCL is not found.""" + if (cuda_home := cuda_home_path()) is not None: + nccl_root = cuda_home.parent / "nccl" + if nccl_root.is_dir(): + return nccl_root + + # Check NCCL Python packages + for package_name in ["nvidia-nccl-cu13", "nvidia-nccl-cu12"]: + try: + nccl_distribution = distribution(package_name) + except PackageNotFoundError: + continue + + nccl_root = Path(nccl_distribution.locate_file("nvidia/nccl")) + if nccl_root.is_dir(): + return nccl_root + + return None + + +@functools.lru_cache(maxsize=None) +def nccl_include_path() -> Optional[Path]: + """Return the NCCL include directory.""" + if (nccl_root := nccl_root_path()) is not None: + include_path = nccl_root / "include" + if include_path.is_dir(): + return include_path + + return None + + +@functools.lru_cache(maxsize=None) +def nccl_lib_path() -> Optional[Path]: + """Return the NCCL shared library path.""" + if (nccl_root := nccl_root_path()) is not None: + lib_path = nccl_root / "lib" / "libnccl.so.2" + if lib_path.is_file(): + return lib_path + + return None + + +@functools.lru_cache(maxsize=None) +def cuda_toolkit_include_path() -> Optional[Path]: + """Return the CUDA Toolkit include directory.""" + if (cuda_home := cuda_home_path()) is not None: + include_path = cuda_home / "include" + if include_path.is_dir(): + return include_path + return None @functools.lru_cache(maxsize=None) @@ -324,13 +409,9 @@ def cuda_version() -> Tuple[int, ...]: and check pip version. """ - try: - nvcc_bin = nvcc_path() - except FileNotFoundError as e: - pass - else: + if (nvcc_bin := nvcc_path()) is not None: output = subprocess.run( - [nvcc_bin, "-V"], + [str(nvcc_bin), "-V"], capture_output=True, check=True, universal_newlines=True, @@ -339,12 +420,20 @@ def cuda_version() -> Tuple[int, ...]: version = match.group(1).split(".") return tuple(int(v) for v in version) - try: - version_str = get_version("nvidia-cuda-runtime-cu12") - version_tuple = tuple(int(part) for part in version_str.split(".") if part.isdigit()) - return version_tuple - except importlib.metadata.PackageNotFoundError: - raise RuntimeError("Could neither find NVCC executable nor CUDA runtime Python package.") + version_str: Optional[str] = None + package_names = ["nvidia-cuda-runtime", "nvidia-cuda-runtime-cu13", "nvidia-cuda-runtime-cu12"] + + for package_name in package_names: + try: + version_str = get_version(package_name) + except PackageNotFoundError: + pass + else: + return tuple(int(part) for part in version_str.split(".") if part.isdigit()) + + raise RuntimeError( + f"Could neither find NVCC executable nor CUDA runtime Python package for {package_names}." + ) def get_frameworks() -> List[str]: diff --git a/setup.py b/setup.py index 2a8a9d7688..5187179498 100644 --- a/setup.py +++ b/setup.py @@ -19,6 +19,7 @@ from build_tools.te_version import te_version from build_tools.utils import ( cuda_archs, + cuda_home_path, cuda_version, cudnn_frontend_include_path, get_frameworks, @@ -26,6 +27,7 @@ min_python_version_str, nccl_ep_enabled, get_max_jobs_for_parallel_build, + nvcc_path, ) frameworks = get_frameworks() @@ -253,6 +255,10 @@ def build_nccl_ep_submodule() -> str: nproc = get_max_jobs_for_parallel_build() env = os.environ.copy() + if (cuda_home := cuda_home_path()) is not None: + env.setdefault("CUDA_HOME", str(cuda_home)) + if (nvcc_bin := nvcc_path()) is not None: + env.setdefault("NVCC", str(nvcc_bin)) env["NVCC_GENCODE"] = gencode # NCCL EP needs the core NCCL headers + libnccl.so; write NCCL EP build # outputs to the submodule's local build/ tree.