diff --git a/.gitmodules b/.gitmodules index 07647e915e..88d431cba5 100644 --- a/.gitmodules +++ b/.gitmodules @@ -5,5 +5,5 @@ path = 3rdparty/cutlass url = https://github.com/NVIDIA/cutlass.git [submodule "3rdparty/nccl"] - path = 3rdparty/nccl - url = https://github.com/NVIDIA/nccl.git + path = 3rdparty/nccl-extensions + url = https://github.com/NVIDIA/nccl-extensions.git diff --git a/3rdparty/nccl b/3rdparty/nccl deleted file mode 160000 index b87848fbc5..0000000000 --- a/3rdparty/nccl +++ /dev/null @@ -1 +0,0 @@ -Subproject commit b87848fbc52da65b5a898b4ac6633fcf51cec4ed diff --git a/3rdparty/nccl-extensions b/3rdparty/nccl-extensions new file mode 160000 index 0000000000..2c6135a721 --- /dev/null +++ b/3rdparty/nccl-extensions @@ -0,0 +1 @@ +Subproject commit 2c6135a721824ff792af7b72900b0ab758fa1f98 diff --git a/examples/jax/ep/run_test_ep.sh b/examples/jax/ep/run_test_ep.sh index 1305ca6fd1..86aa6ca087 100755 --- a/examples/jax/ep/run_test_ep.sh +++ b/examples/jax/ep/run_test_ep.sh @@ -32,8 +32,8 @@ export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}" COORD="${COORD:-127.0.0.1:12345}" TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-300}" -# Editable installs don't embed rpath; libtransformer_engine.so needs -# libnccl_ep.so.0 from the TE editable location at dlopen time. +# Editable installs don't embed rpath; the TE JAX extension needs +# libtransformer_engine.so from the TE editable location at dlopen time. TE_LIB_PATH=$(pip3 show transformer-engine 2>/dev/null \ | grep -E "Location:|Editable project location:" \ | tail -n 1 | awk '{print $NF}') diff --git a/examples/pytorch/ep/bench/run_ep_bench.sh b/examples/pytorch/ep/bench/run_ep_bench.sh index fefecd7fa9..3b0977e4c3 100755 --- a/examples/pytorch/ep/bench/run_ep_bench.sh +++ b/examples/pytorch/ep/bench/run_ep_bench.sh @@ -26,10 +26,8 @@ if [ "${NSYS}" -eq 1 ] && [ "${KINETO}" -eq 1 ]; then fi SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" RESULTS="${SCRIPT_DIR}/results" mkdir -p "${RESULTS}" -export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) NUM_GPUS="${NUM_GPUS:-${DETECTED_GPUS}}" diff --git a/examples/pytorch/ep/bench/run_nccl_ep_bench.sh b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh index 8f6da04a00..ac5fcacc25 100755 --- a/examples/pytorch/ep/bench/run_nccl_ep_bench.sh +++ b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh @@ -23,8 +23,8 @@ TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" RESULTS="${SCRIPT_DIR}/results" mkdir -p "${RESULTS}" -BIN="${TE_REPO_ROOT}/3rdparty/nccl/build/test/nccl_ep/ep_bench" -LIB="${TE_REPO_ROOT}/3rdparty/nccl/build/lib" +BIN="${TE_REPO_ROOT}/3rdparty/nccl-extensions/build/test/nccl_ep/ep_bench" +LIB="${TE_REPO_ROOT}/3rdparty/nccl-extensions/build/lib" [ -x "${BIN}" ] || { echo "ep_bench not built at ${BIN}" >&2; exit 2; } NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) diff --git a/examples/pytorch/ep/run_test_ep.sh b/examples/pytorch/ep/run_test_ep.sh index 13b41f4cb2..d8e6b50556 100755 --- a/examples/pytorch/ep/run_test_ep.sh +++ b/examples/pytorch/ep/run_test_ep.sh @@ -17,7 +17,6 @@ if [ "${NUM_GPUS}" -gt 8 ]; then NUM_GPUS=8; fi : ${TEST_TIMEOUT_S:=120} SCRIPT="${TE_PATH}/examples/pytorch/ep/ep_moe.py" -export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}" # Stage JIT cubins on tmpfs for fast iteration. : ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"} diff --git a/setup.py b/setup.py index 150d92969c..c0c37364a9 100644 --- a/setup.py +++ b/setup.py @@ -207,11 +207,11 @@ def _discover_nccl_home() -> str: 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(): + """Build libnccl_ep.a from the 3rdparty/nccl-extensions submodule and return NCCL_HOME.""" + nccl_root = current_file_path / "3rdparty" / "nccl-extensions" + if not (nccl_root / "nccl_ep" / "Makefile").exists(): raise RuntimeError( - f"NCCL submodule not found at {nccl_root}. " + f"NCCL EP submodule not found at {nccl_root}. " "Run `git submodule update --init --recursive`." ) @@ -267,13 +267,13 @@ def build_nccl_ep_submodule() -> str: "rebuilding libnccl_ep.a" ) subprocess.check_call( - ["make", "-C", "contrib/nccl_ep", "clean"], + ["make", "-C", "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"], + ["make", "-j", str(nproc), "-C", "nccl_ep", "lib"], cwd=str(nccl_root), env=env, ) diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index 7dbbcdce9d..47732196ed 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -675,7 +675,7 @@ class EPPipelineTest : public EpOpTestBase, public ::testing::WithParamInterface TEST_P(EPPipelineTest, FullForwardBackward) { const DType dtype = GetParam(); // NCCL EP backend currently asserts ncclBfloat16 in ncclEpDispatch - // (contrib/nccl_ep/nccl_ep.cc); skip FP16/FP32 until the backend supports them. + // (nccl_ep/nccl_ep.cc); skip FP16/FP32 until the backend supports them. if (dtype != DType::kBFloat16) { GTEST_SKIP() << test::typeName(dtype) << " not yet supported by NCCL EP backend"; } @@ -750,8 +750,9 @@ class EPZeroCopyTest : public EpOpTestBase { }; TYPED_TEST_SUITE(EPZeroCopyTest, EPBf16Only); -// Identity round-trip with symm-mem on dispatch i/o + combine input. Bit-exact -// vs HBM reference (same routing, same input). +// Identity round-trip with symm-mem on dispatch i/o + combine input. The combined +// result is bit-exact vs the HBM reference; the intermediate recv buffer is not, +// since zero-copy and HBM dispatch use different per-expert layouts. TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { using Tok = TypeParam; EP_PULL_FIXTURE(); @@ -776,10 +777,7 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { ref_t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - std::vector ref_recv(ref_buf.recv_capacity * hidden_dim_); std::vector ref_result(num_tokens_ * hidden_dim_); - NVTE_CHECK_CUDA(cudaMemcpy(ref_recv.data(), ref_buf.recv_tokens.get(), - ref_recv.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(ref_result.data(), ref_buf.result.get(), ref_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); @@ -818,24 +816,17 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { symm_window(sym_recv), sym_t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - std::vector sym_recv_host(sym_buf.recv_capacity * hidden_dim_); std::vector sym_result(num_tokens_ * hidden_dim_); - NVTE_CHECK_CUDA(cudaMemcpy(sym_recv_host.data(), sym_recv.ptr, - sym_recv_host.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(sym_result.data(), sym_buf.result.get(), sym_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); - // Compare per filled recv slot (HBM ref vs symm) and full result. - int total_recv = this->template read_total_recv(sym_buf); - for (int i = 0; i < total_recv * hidden_dim_; ++i) - ASSERT_EQ(tok_to_float(sym_recv_host[i]), tok_to_float(ref_recv[i])) - << "recv mismatch at " << i; + // Combined result is the cross-mode invariant (see note above). for (size_t i = 0; i < sym_result.size(); ++i) ASSERT_EQ(tok_to_float(sym_result[i]), tok_to_float(ref_result[i])) << "result mismatch at " << i; if (g_process_id == 0) - printf(" IdentityAllSymm: passed (recv_slots=%d, bit-exact vs HBM)\n", total_recv); + printf(" IdentityAllSymm: passed (result bit-exact vs HBM)\n"); NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } diff --git a/tests/jax/multi_process_launch_ep.sh b/tests/jax/multi_process_launch_ep.sh index ff89f712eb..8547d77f2b 100755 --- a/tests/jax/multi_process_launch_ep.sh +++ b/tests/jax/multi_process_launch_ep.sh @@ -17,8 +17,8 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" -# Editable installs don't embed rpath; libtransformer_engine.so needs -# libnccl_ep.so.0 from the TE editable location at dlopen time. +# Editable installs don't embed rpath; the TE JAX extension needs +# libtransformer_engine.so from the TE editable location at dlopen time. TE_LIB_PATH=$(pip3 show transformer-engine 2>/dev/null \ | grep -E "Location:|Editable project location:" \ | tail -n 1 | awk '{print $NF}') diff --git a/tests/pytorch/distributed/run_test_ep.sh b/tests/pytorch/distributed/run_test_ep.sh index 68b691f787..62c9a0207f 100755 --- a/tests/pytorch/distributed/run_test_ep.sh +++ b/tests/pytorch/distributed/run_test_ep.sh @@ -9,8 +9,6 @@ set -uo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" -export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) if [ "${DETECTED_GPUS}" -lt 4 ]; then diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index d09876c990..8eba515e5e 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -446,10 +446,10 @@ option(NVTE_WITH_NCCL_EP "Build NCCL EP into libtransformer_engine.so" ON) if(NVTE_WITH_NCCL_EP) # SM>=90 and NCCL>=2.30.4 are gated at runtime in EPBackend::initialize. # -- NCCL EP headers -------------------------------------------------------- -# Headers + libs are produced by the in-tree 3rdparty/nccl submodule build +# Headers + libs are produced by the in-tree 3rdparty/nccl-extensions submodule build # (auto-built by setup.py via build_nccl_ep_submodule). set(NCCL_EP_SUBMODULE_ROOT - "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl") + "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl-extensions") set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/include") if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") message(FATAL_ERROR