Revert "Expert Parallelism: common C API + NCCL EP backend" (#3126)

Revert "Expert Parallelism: common C API + NCCL EP backend (#3034)"

This reverts commit c3396ee.

Signed-off-by: Tim Moon <tmoon@nvidia.com>

Author: phu0ngng

.gitmodules

@@ -7,6 +7,3 @@
 [submodule "3rdparty/cutlass"]
 	path = 3rdparty/cutlass
 	url = https://github.com/NVIDIA/cutlass.git
-[submodule "3rdparty/nccl"]
-	path = 3rdparty/nccl
-	url = https://github.com/NVIDIA/nccl.git

3rdparty/nccl

@@ -2,52 +2,16 @@
 #
 # See LICENSE for license information.
 
-function test_fail() {
-    RET=1
-    FAILED_CASES="$FAILED_CASES $1"
-    echo "Error: sub-test failed: $1"
-}
-
-RET=0
-FAILED_CASES=""
+set -e
 
 # Find TE
 : ${TE_PATH:=/opt/transformerengine}
-: ${XML_LOG_DIR:=/logs}
-mkdir -p "$XML_LOG_DIR"
-
 TE_LIB_PATH=$(pip3 show transformer-engine | grep -E "Location:|Editable project location:" | tail -n 1 | awk '{print $NF}')
 export LD_LIBRARY_PATH=$TE_LIB_PATH:$LD_LIBRARY_PATH
 
 if [[ $(nvidia-smi --list-gpus | wc -l) -ge 4 ]]; then
-    cd $TE_PATH/tests/cpp_distributed
-    configure_ok=1
-    cmake -GNinja -S. -Bbuild || { test_fail "configure"; configure_ok=0; }
-
-    # Build each suite independently so one suite's build failure does not
-    # mask the other's results. Skip mpirun when the binary is missing.
-    if [[ $configure_ok -eq 1 ]]; then
-        comm_gemm_ok=1
-        ep_ok=1
-        cmake --build build --target test_comm_gemm || { test_fail "test_comm_gemm_build"; comm_gemm_ok=0; }
-        cmake --build build --target test_ep        || { test_fail "test_ep_build";        ep_ok=0; }
-
-        if [[ $comm_gemm_ok -eq 1 ]]; then
-            # Per-rank XML to avoid a write race on a shared path.
-            mpirun --allow-run-as-root --np 4 --oversubscribe bash -c \
-                "exec ./build/test_comm_gemm --gtest_output=xml:$XML_LOG_DIR/cpp_distributed_test_comm_gemm.rank\${OMPI_COMM_WORLD_RANK}.xml" \
-                || test_fail "test_comm_gemm"
-        fi
-
-        if [[ $ep_ok -eq 1 ]]; then
-            # EP suites; runner self-skips on pre-Hopper GPUs.
-            GTEST_XML_PREFIX="$XML_LOG_DIR/cpp_distributed_test_ep" \
-                bash ./run_test_ep.sh 4 ./build || test_fail "test_ep"
-        fi
-    fi
-fi
-
-if [ "$RET" -ne 0 ]; then
-    echo "FAILED sub-tests:$FAILED_CASES"
+	cd $TE_PATH/tests/cpp_distributed
+	cmake -GNinja -S. -Bbuild
+	cmake --build build
+	mpirun --allow-run-as-root --np 4 --oversubscribe ./build/test_comm_gemm
 fi
-exit $RET

qa/L1_cpp_distributed/test.sh

@@ -84,31 +84,6 @@ 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
-    if build_with_nccl_ep:
-        nccl_home = build_nccl_ep_submodule()
-        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:
@@ -155,138 +130,6 @@ def setup_requirements() -> Tuple[List[str], List[str]]:
     return [remove_dups(reqs) for reqs in [install_reqs, test_reqs]]
 
 
-def _discover_nccl_home() -> str:
-    """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)
-
-    nproc = os.cpu_count() or 8
-    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.

setup.py

@@ -55,32 +55,10 @@ target_include_directories(test_comm_gemm PRIVATE ${test_comm_gemm_INCLUDES})
 find_package(CUDAToolkit REQUIRED)
 find_package(OpenMP REQUIRED)
 find_package(MPI REQUIRED)
-
-# -- NCCL core ----------------------------------------------------------------
-# Anchor on libnccl and derive nccl.h from the same install prefix so the
-# header and library can't drift across installs.
 find_library(NCCL_LIB
              NAMES nccl libnccl
-             HINTS /opt/nvidia/nccl/lib /opt/nvidia/nccl/lib64
-                   /usr/local/nccl/lib /usr/local/nccl/lib64
-             PATH_SUFFIXES lib lib64
+             PATH_SUFFIXES lib
              REQUIRED)
-get_filename_component(_nccl_lib_dir "${NCCL_LIB}" DIRECTORY)
-set(NCCL_PREFIX "${_nccl_lib_dir}")
-while(NCCL_PREFIX AND NOT EXISTS "${NCCL_PREFIX}/include/nccl.h")
-  get_filename_component(_nccl_parent "${NCCL_PREFIX}" DIRECTORY)
-  if(_nccl_parent STREQUAL NCCL_PREFIX)
-    break()
-  endif()
-  set(NCCL_PREFIX "${_nccl_parent}")
-endwhile()
-find_path(NCCL_INCLUDE_DIR nccl.h
-          HINTS "${NCCL_PREFIX}/include"
-          NO_DEFAULT_PATH)
-if(NOT NCCL_INCLUDE_DIR)
-  message(FATAL_ERROR
-    "nccl.h not found under the prefix of ${NCCL_LIB}.")
-endif()
 list(APPEND test_comm_gemm_LINKER_LIBS
      CUDA::cuda_driver
      CUDA::cudart
@@ -96,37 +74,3 @@ target_compile_options(test_comm_gemm PRIVATE -O2 -fopenmp)
 
 include(GoogleTest)
 gtest_discover_tests(test_comm_gemm DISCOVERY_TIMEOUT 600)
-
-# -- EP distributed tests ------------------------------------------------------
-# 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 <transformer_engine/ep.h>, which is
-# statically linked into libtransformer_engine.so.
-message(STATUS "EP test: NCCL headers: ${NCCL_INCLUDE_DIR}")
-set(EP_TEST_COMMON_INCLUDES
-    ${NCCL_INCLUDE_DIR}
-    ${MPI_CXX_INCLUDE_PATH}
-    ../../transformer_engine/common/include
-    ../../transformer_engine/common
-    ${CMAKE_CURRENT_SOURCE_DIR})
-
-# nvrtc must follow TE_LIB so symbols referenced from libtransformer_engine.so
-# (loaded via dlopen in Python; not in its DT_NEEDED) resolve through nvrtc.
-set(EP_TEST_COMMON_LIBS
-    CUDA::cuda_driver
-    CUDA::cudart
-    GTest::gtest
-    ${TE_LIB}
-    CUDA::nvrtc
-    ${NCCL_LIB}
-    MPI::MPI_CXX
-    OpenMP::OpenMP_CXX)
-
-# -- EP distributed tests (per-op + full pipeline + zero-copy symm) -----------
-add_executable(test_ep test_ep.cu ../cpp/test_common.cu)
-target_include_directories(test_ep PRIVATE ${EP_TEST_COMMON_INCLUDES})
-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)")