diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index cbb8838b00..d6c8ddb315 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,31 @@ 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 += [ + f"-DCUDAToolkit_ROOT={discovered_cuda_home}", + # CUDA wheels use `lib`, while nvcc's host linker expects `lib64`. + f"-DCMAKE_CUDA_FLAGS=-L{discovered_cuda_home}/lib", + ] + + if cuda_full_version := cuda_version(): + cuda_major_version = cuda_full_version[0] + cuda_lib_dir = discovered_cuda_home / "lib" + configure_command += [ + f"-DCUDA_CUDART={cuda_lib_dir / f'libcudart.so.{cuda_major_version}'}", + f"-DCUDA_cudart_LIBRARY={cuda_lib_dir / f'libcudart.so.{cuda_major_version}'}", + f"-DCUDA_cublas_LIBRARY={cuda_lib_dir / f'libcublas.so.{cuda_major_version}'}", + ( + f"-DCUDA_cublasLt_LIBRARY={cuda_lib_dir / f'libcublasLt.so.{cuda_major_version}'}" + ), + ] + + 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 +212,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..da4ecf957b 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. @@ -175,6 +175,90 @@ def found_pybind11() -> bool: return False +@functools.lru_cache(maxsize=None) +def cuda_home_path() -> Optional[Path]: + """Returns the CUDA home path. This path should contain binaries (e.g. nvcc), headers, and libraries. + + Returns `None` if CUDA is not found.""" + if cuda_home := os.getenv("CUDA_HOME"): + return Path(cuda_home) + + # Check if site-packages contains an `nvidia` directory + for site_package in sys.path: + if not Path(site_package).is_dir(): + continue + + nvidia_dir = Path(site_package) / "nvidia" + if not nvidia_dir.is_dir(): + continue + + if Path(nvidia_dir / "bin").is_dir(): + return nvidia_dir + + # In this case there must be a "CUDA version directory" like `cu12` or `cu13` + # that contains the binaries, headers, and libraries + for cuda_version_dir in nvidia_dir.iterdir(): + if not cuda_version_dir.is_dir(): + continue + + # Verify that the directory name matches the expected `cu##` format + if not re.match(r"cu\d+", cuda_version_dir.name): + continue + + if not (cuda_version_dir / "bin").is_dir(): + continue + + return cuda_version_dir + + 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() -> Tuple[str, str]: """Returns root path for cuda toolkit includes. @@ -198,30 +282,49 @@ def cuda_toolkit_include_path() -> Tuple[str, str]: @functools.lru_cache(maxsize=None) -def nvcc_path() -> Tuple[str, str]: - """Returns the NVCC binary path. +def nvcc_path() -> Optional[Path]: + """Get 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}") + Returns `None` if NVCC is not found. + """ + + def lookup_via_cuda_home() -> Optional[str]: + if (cuda_home := cuda_home_path()) is not None: + return cuda_home / "bin" / "nvcc" + return None + + def lookup_via_path() -> Optional[str]: + if (nvcc_bin := shutil.which("nvcc")) is not None: + return nvcc_bin + return None + + def lookup_via_distribution() -> Optional[str]: + try: + return distribution("nvidia-cuda-nvcc").locate_file("bin/nvcc") + except PackageNotFoundError: + return None + + def lookup_via_local_cuda() -> Optional[str]: + return "/usr/local/cuda/bin/nvcc" + + nvcc_lookup_funcs: List[Callable[[], Optional[str]]] = [ + lookup_via_cuda_home, + lookup_via_path, + lookup_via_distribution, + lookup_via_local_cuda, + ] + + for nvcc_lookup_func in nvcc_lookup_funcs: + nvcc_bin = nvcc_lookup_func() - return nvcc_bin + if nvcc_bin is None: + continue + + nvcc_bin_path = Path(nvcc_bin) + if nvcc_bin_path.is_file(): + return nvcc_bin_path + + return None @functools.lru_cache(maxsize=None) @@ -324,13 +427,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 +438,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]: