Expert Parallelism: common C API + NCCL EP backend #3034
+2,450
−7
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -84,6 +84,31 @@ def setup_common_extension() -> CMakeExtension: | |||||
| cusolvermp_dir = os.getenv("CUSOLVERMP_HOME", "/usr") | ||||||
| cmake_flags.append(f"-DCUSOLVERMP_DIR={cusolvermp_dir}") | ||||||
|
|
||||||
| # NCCL EP (Hopper+): on by default; auto-skipped when no arch >= 90 is | ||||||
| # targeted. Set NVTE_WITH_NCCL_EP=0 to force off. | ||||||
| nccl_ep_env = os.getenv("NVTE_WITH_NCCL_EP") | ||||||
| nccl_ep_explicit = nccl_ep_env is not None | ||||||
| build_with_nccl_ep = bool(int(nccl_ep_env if nccl_ep_explicit else "1")) | ||||||
| if build_with_nccl_ep: | ||||||
| arch_tokens = [a.strip() for a in str(archs or "").split(";") if a.strip()] | ||||||
| has_hopper_or_newer = any( | ||||||
| t.lower() == "native" or (t.rstrip("af").isdigit() and int(t.rstrip("af")) >= 90) | ||||||
| for t in arch_tokens | ||||||
| ) | ||||||
| if not has_hopper_or_newer: | ||||||
| if nccl_ep_explicit: | ||||||
| raise RuntimeError( | ||||||
| f"NVTE_WITH_NCCL_EP=1 was set but NVTE_CUDA_ARCHS ('{archs}') " | ||||||
| "contains no arch >= 90. NCCL EP requires Hopper or newer." | ||||||
| ) | ||||||
| print(f"[NCCL EP] No arch >= 90 in NVTE_CUDA_ARCHS ('{archs}'); skipping build.") | ||||||
| build_with_nccl_ep = False | ||||||
|
Comment on lines
+92
to
+105
Member
There was a problem hiding this comment. Is this needed? Considering that we are fine if they compile for anything AND Hopper+, then we should also be fine if they compile for the same things without Hopper+ (and then it should just say unsupported at runtime once somebody tries to run it on that hardware).
Collaborator
Author
There was a problem hiding this comment. Cleaned up this fallback. Skipped NCCL EP build when arch < 90.
phu0ngng marked this conversation as resolved.
|
||||||
| if build_with_nccl_ep: | ||||||
| nccl_home = build_nccl_ep_submodule() | ||||||
|
Member
There was a problem hiding this comment. Flagging that we compile
It's not obviously wrong, but it does seem messy. I suppose there are also some edge case problems (if NCCL is not installed but is listed as a dependency, compiling NCCL EP in the configuration stage will fail but building during the build stage can succeed). It seems the "proper" approach would be to treat |
||||||
| cmake_flags.append(f"-DNCCL_INCLUDE_DIR={nccl_home}/include") | ||||||
| else: | ||||||
| cmake_flags.append("-DNVTE_WITH_NCCL_EP=OFF") | ||||||
|
|
||||||
| # Add custom CMake arguments from environment variable | ||||||
| nvte_cmake_extra_args = os.getenv("NVTE_CMAKE_EXTRA_ARGS") | ||||||
| if nvte_cmake_extra_args: | ||||||
|
|
@@ -130,6 +155,138 @@ def setup_requirements() -> Tuple[List[str], List[str]]: | |||||
| return [remove_dups(reqs) for reqs in [install_reqs, test_reqs]] | ||||||
|
|
||||||
|
|
||||||
| def _discover_nccl_home() -> str: | ||||||
|
phu0ngng marked this conversation as resolved.
Member
There was a problem hiding this comment. This could be a standalone utility funciton in |
||||||
| """Resolve NCCL_HOME: honor env var, else probe well-known prefixes, else ldconfig.""" | ||||||
| env_home = os.environ.get("NCCL_HOME") | ||||||
| if env_home: | ||||||
| if (Path(env_home) / "include" / "nccl.h").exists(): | ||||||
| return env_home | ||||||
| print( | ||||||
| f"[NCCL EP] WARNING: NCCL_HOME='{env_home}' is set but " | ||||||
| f"'{env_home}/include/nccl.h' was not found; falling back to system probes." | ||||||
| ) | ||||||
|
|
||||||
| 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") | ||||||
|
|
||||||
| # pip-installed NCCL (nvidia-nccl-cu* wheel) lives under nvidia/nccl in | ||||||
| # site-packages and has no top-level include/lib layout. | ||||||
| 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 | ||||||
|
|
||||||
| for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"): | ||||||
| p = Path(cand) | ||||||
| if (p / "include" / "nccl.h").exists() and any( | ||||||
| (p / sub / name).exists() for sub in lib_subdirs for name in lib_names | ||||||
| ): | ||||||
| return str(p) | ||||||
|
|
||||||
| try: | ||||||
| out = subprocess.check_output(["ldconfig", "-p"], stderr=subprocess.DEVNULL).decode() | ||||||
| for line in out.splitlines(): | ||||||
| if "libnccl.so" in line and "=>" in line: | ||||||
| lib_path = Path(line.split("=>")[-1].strip()) | ||||||
| # Walk upward so multiarch layouts (.../lib/<triplet>/libnccl.so) | ||||||
| # resolve to the prefix that contains include/nccl.h. | ||||||
| for root in (lib_path.parent.parent, lib_path.parent.parent.parent): | ||||||
| if (root / "include" / "nccl.h").exists(): | ||||||
| return str(root) | ||||||
| except (subprocess.CalledProcessError, FileNotFoundError): | ||||||
| pass | ||||||
|
|
||||||
| raise RuntimeError( | ||||||
| "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix." | ||||||
| ) | ||||||
|
|
||||||
|
|
||||||
| def build_nccl_ep_submodule() -> str: | ||||||
| """Build libnccl_ep.a from the 3rdparty/nccl submodule and return NCCL_HOME.""" | ||||||
| nccl_root = current_file_path / "3rdparty" / "nccl" | ||||||
| if not (nccl_root / "Makefile").exists(): | ||||||
| raise RuntimeError( | ||||||
| f"NCCL submodule not found at {nccl_root}. " | ||||||
| "Run `git submodule update --init --recursive`." | ||||||
| ) | ||||||
|
|
||||||
| build_dir = nccl_root / "build" | ||||||
| nccl_ep_lib = build_dir / "lib" / "libnccl_ep.a" | ||||||
| gencode_stamp = build_dir / "lib" / "libnccl_ep.gencode" | ||||||
|
|
||||||
| # Caller gates on arch >= 90 or "native"; expand "native" to the host's | ||||||
| # actual sm_XX so the build stamp distinguishes machines. | ||||||
| arch_tokens = [a.strip() for a in str(cuda_archs() or "").split(";") if a.strip()] | ||||||
| arch_list: list[str] = [] | ||||||
| for t in arch_tokens: | ||||||
| if t.lower() == "native": | ||||||
| try: | ||||||
| out = subprocess.check_output( | ||||||
| ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], | ||||||
| stderr=subprocess.DEVNULL, | ||||||
| ).decode() | ||||||
| except (subprocess.CalledProcessError, FileNotFoundError) as e: | ||||||
| raise RuntimeError( | ||||||
| "NVTE_CUDA_ARCHS=native requires nvidia-smi to resolve the host arch." | ||||||
| ) from e | ||||||
| for line in out.splitlines(): | ||||||
| cap = line.strip().replace(".", "") | ||||||
| if cap.isdigit() and int(cap) >= 90 and cap not in arch_list: | ||||||
| arch_list.append(cap) | ||||||
| else: | ||||||
| bare = t.rstrip("af") | ||||||
| if bare.isdigit() and int(bare) >= 90 and bare not in arch_list: | ||||||
| arch_list.append(bare) | ||||||
| if not arch_list: | ||||||
| raise RuntimeError( | ||||||
| "NCCL EP requires Hopper or newer (SM >= 90); none found in" | ||||||
| f" NVTE_CUDA_ARCHS={cuda_archs()!r}. Re-run with NVTE_WITH_NCCL_EP=0 to skip the NCCL" | ||||||
| " EP build (the rest of TE still builds)." | ||||||
| ) | ||||||
| gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list) | ||||||
|
phu0ngng marked this conversation as resolved.
|
||||||
|
|
||||||
| nproc = os.cpu_count() or 8 | ||||||
|
Member
There was a problem hiding this comment. This seems likely to break users with wimpy nodes. We can reuse an existing API to control build parallelism:
Suggested change
|
||||||
| env = os.environ.copy() | ||||||
| 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. | ||||||
| nccl_home = _discover_nccl_home() | ||||||
| env["NCCL_HOME"] = nccl_home | ||||||
| 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: | ||||||
| print( | ||||||
| f"[NCCL EP] gencode changed ('{prev_gencode}' -> '{gencode}'); " | ||||||
| "rebuilding libnccl_ep.a" | ||||||
| ) | ||||||
| subprocess.check_call( | ||||||
| ["make", "-C", "contrib/nccl_ep", "clean"], | ||||||
| cwd=str(nccl_root), | ||||||
| env=env, | ||||||
| ) | ||||||
| print(f"[NCCL EP] Building libnccl_ep.a (gencode='{gencode}')") | ||||||
| subprocess.check_call( | ||||||
| ["make", "-j", str(nproc), "-C", "contrib/nccl_ep", "lib"], | ||||||
| cwd=str(nccl_root), | ||||||
| env=env, | ||||||
| ) | ||||||
| gencode_stamp.parent.mkdir(parents=True, exist_ok=True) | ||||||
| gencode_stamp.write_text(gencode) | ||||||
|
|
||||||
| return nccl_home | ||||||
|
|
||||||
|
|
||||||
| def git_check_submodules() -> None: | ||||||
| """ | ||||||
| Attempt to checkout git submodules automatically during setup. | ||||||
|
|
||||||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,63 @@ | ||
| #!/usr/bin/env bash | ||
| # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| # | ||
| # See LICENSE for license information. | ||
| # | ||
| # Run TE EP distributed unit tests via mpirun. Each MPI rank pins to one GPU | ||
| # (rank % device_count) and exchanges ncclUniqueId through MPI_Bcast. | ||
| # | ||
| # Usage: | ||
| # bash run_test_ep.sh [num_gpus] [build_dir] | ||
| # | ||
| # Defaults: | ||
| # num_gpus = number of GPUs visible to nvidia-smi | ||
| # build_dir = <script_dir>/build | ||
| # | ||
| # Environment variables: | ||
| # GTEST_FILTER - forwarded to all processes (e.g., "EPPipelineTest.*") | ||
| # GTEST_XML_PREFIX - if set, each rank writes JUnit XML to | ||
| # ${GTEST_XML_PREFIX}.rank<N>.xml | ||
| # MPIRUN - override the mpirun binary (default: mpirun) | ||
| # MPIRUN_EXTRA - extra flags forwarded to mpirun | ||
|
|
||
| set -euo pipefail | ||
|
|
||
| SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" | ||
| BUILD_DIR="${2:-${SCRIPT_DIR}/build}" | ||
| NUM_GPUS="${1:-$(nvidia-smi -L 2>/dev/null | wc -l)}" | ||
| MPIRUN="${MPIRUN:-mpirun}" | ||
|
|
||
| # Skip cleanly on pre-Hopper: NCCL EP requires SM>=90. | ||
| MIN_SM=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null \ | ||
| | awk -F. 'NR==1 || ($1*10+$2)<min { min=$1*10+$2 } END { print min+0 }') | ||
| if (( MIN_SM > 0 && MIN_SM < 90 )); then | ||
| echo "NCCL EP requires SM>=90 (lowest visible GPU is SM${MIN_SM}); SKIPPING." | ||
| exit 0 | ||
| fi | ||
|
|
||
| TEST_BIN="${BUILD_DIR}/test_ep" | ||
| if [[ ! -x "${TEST_BIN}" ]]; then | ||
| echo "ERROR: binary not found: ${TEST_BIN}" | ||
| echo "Build: cd ${SCRIPT_DIR} && mkdir -p build && cd build && cmake .. && make" | ||
| exit 1 | ||
| fi | ||
|
|
||
| if (( NUM_GPUS < 2 )); then | ||
| echo "EP Tests: requires at least 2 GPUs, found ${NUM_GPUS}. Skipping." | ||
| exit 0 | ||
| fi | ||
|
|
||
| GTEST_ARGS="${GTEST_FILTER:+--gtest_filter=${GTEST_FILTER}}" | ||
|
|
||
| echo "=== EP Tests ===" | ||
| echo " GPUs: ${NUM_GPUS} Binary: ${TEST_BIN}" | ||
| echo | ||
|
|
||
| if [[ -n "${GTEST_XML_PREFIX:-}" ]]; then | ||
| # bash -c so OMPI_COMM_WORLD_RANK expands per-rank, avoiding a write race | ||
| # on a single shared output path. | ||
| "${MPIRUN}" --allow-run-as-root --oversubscribe -n "${NUM_GPUS}" ${MPIRUN_EXTRA:-} bash -c \ | ||
| "exec '${TEST_BIN}' ${GTEST_ARGS} --gtest_output=xml:${GTEST_XML_PREFIX}.rank\${OMPI_COMM_WORLD_RANK}.xml" | ||
| else | ||
| "${MPIRUN}" --allow-run-as-root --oversubscribe -n "${NUM_GPUS}" ${MPIRUN_EXTRA:-} "${TEST_BIN}" ${GTEST_ARGS} | ||
| fi |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In general this is fine, but still interested to know why you needed to do those changes?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When I run the suite on a local node, I don't have cublasMp in the image. So the comm_gemm tests failed, and the EP tests are skipped. The change is to make sure EP test runs even comm_gemm tests failed.