[JAX] Support single-process-multi-device EP via XLA-borrowed communicator

examples/jax/ep/ep_moe.py

@@ -4,6 +4,8 @@
 """End-to-end MoE example: dispatch -> batched expert linear -> combine, fwd + bwd.
 
 One process per GPU. Run via run_test_ep.sh.
+Use ``python ep_moe.py --single-process`` for all local GPUs; this requires
+XLA borrowed communicator support.
 """
 
 import argparse
@@ -23,9 +25,19 @@
 
 def _parse_args():
     p = argparse.ArgumentParser(description="TE-JAX EP MoE example (fwd + bwd)")
-    p.add_argument("--coordinator-address", required=True)
-    p.add_argument("--process-id", type=int, required=True)
-    p.add_argument("--num-processes", type=int, required=True)
+    p.add_argument(
+        "--single-process",
+        action="store_true",
+        help=(
+            "Single-controller mode: one process drives every local GPU via the"
+            " XLA-borrowed-comm EP path; no jax.distributed.initialize, no coordinator."
+            " --coordinator-address/--process-id/--num-processes become optional (and are"
+            " ignored) when this is set."
+        ),
+    )
+    p.add_argument("--coordinator-address", default=None)
+    p.add_argument("--process-id", type=int, default=None)
+    p.add_argument("--num-processes", type=int, default=None)
     p.add_argument("--num-tokens", type=int, default=8, help="Per-rank token count.")
     p.add_argument("--top-k", type=int, default=2)
     p.add_argument("--hidden", type=int, default=32)
@@ -55,10 +67,30 @@ def _parse_args():
         default=3,
         help="Number of fwd+bwd iterations to run (same compiled jit, same handle_mem).",
     )
-    return p.parse_args()
+    args = p.parse_args()
+    if not args.single_process:
+        missing = [
+            name
+            for name, val in (
+                ("--coordinator-address", args.coordinator_address),
+                ("--process-id", args.process_id),
+                ("--num-processes", args.num_processes),
+            )
+            if val is None
+        ]
+        if missing:
+            p.error(f"{', '.join(missing)} required unless --single-process is set")
+    return args
 
 
 def _distributed_init(args):
+    if args.single_process:
+        assert (
+            jax.process_count() == 1
+        ), f"--single-process requires jax.process_count() == 1; got {jax.process_count()}"
+        args.process_id = 0
+        args.num_processes = jax.device_count()
+        return
     jax.distributed.initialize(
         coordinator_address=args.coordinator_address,
         num_processes=args.num_processes,
@@ -126,13 +158,14 @@ def _make_inputs(args):
     dp_size = args.dp_size
     ep_size = args.ep_size
     num_procs = args.num_processes
-    dp_color = args.process_id // ep_size
     NLE = args.num_local_experts
 
-    rng_dp = np.random.default_rng(seed=42 + dp_color)
-    tokens_np = (rng_dp.standard_normal((T, H), dtype=np.float32) * 0.5).astype(np.float32)
-    w_np = np.full((T, K), 1.0 / K, dtype=np.float32)
-    idx_np_list = [_make_routing(dp_color, T, K, E, NLE, offset=i) for i in range(L)]
+    if not args.single_process:
+        dp_color = args.process_id // ep_size
+        rng_dp = np.random.default_rng(seed=42 + dp_color)
+        tokens_np = (rng_dp.standard_normal((T, H), dtype=np.float32) * 0.5).astype(np.float32)
+        w_np = np.full((T, K), 1.0 / K, dtype=np.float32)
+        idx_np_list = [_make_routing(dp_color, T, K, E, NLE, offset=i) for i in range(L)]
 
     tokens_global_np = np.concatenate(
         [
@@ -165,14 +198,31 @@ def _make_inputs(args):
     # [num_procs, T, ...] sharded on the first dim across (dp, ep).
     mesh = args.mesh
     dpep_spec = NamedSharding(mesh, PartitionSpec(("dp", "ep"), None, None))
+    if args.single_process:
+        tokens_local = np.broadcast_to(
+            tokens_global_np.reshape(dp_size, T, H)[:, None], (dp_size, ep_size, T, H)
+        ).reshape(num_procs, T, H)
+        idx_local_list = [
+            np.broadcast_to(
+                idx_g.reshape(dp_size, T, K)[:, None], (dp_size, ep_size, T, K)
+            ).reshape(num_procs, T, K)
+            for idx_g in idx_global_np_list
+        ]
+        w_local = np.broadcast_to(
+            w_global_np.reshape(dp_size, T, K)[:, None], (dp_size, ep_size, T, K)
+        ).reshape(num_procs, T, K)
+    else:
+        tokens_local = tokens_np[None, :, :]
+        idx_local_list = [idx_np[None, :, :] for idx_np in idx_np_list]
+        w_local = w_np[None, :, :]
     tokens = jax.make_array_from_process_local_data(
-        dpep_spec, tokens_np[None, :, :].astype(np.float32), (num_procs, T, H)
+        dpep_spec, tokens_local.astype(np.float32), (num_procs, T, H)
     ).astype(jnp.bfloat16)
     topk_idx_list = [
-        jax.make_array_from_process_local_data(dpep_spec, idx_np[None, :, :], (num_procs, T, K))
-        for idx_np in idx_np_list
+        jax.make_array_from_process_local_data(dpep_spec, idx_local, (num_procs, T, K))
+        for idx_local in idx_local_list
     ]
-    topk_w = jax.make_array_from_process_local_data(dpep_spec, w_np[None, :, :], (num_procs, T, K))
+    topk_w = jax.make_array_from_process_local_data(dpep_spec, w_local, (num_procs, T, K))
     kernels_list = [jnp.asarray(k, dtype=jnp.bfloat16) for k in kernels_np_list]
     return (
         tokens_global_np,
@@ -333,6 +383,17 @@ def main():
             print(f"[ep_moe] SKIPPED: NCCL EP requires SM>=90 (got SM{major}{minor})")
             return
 
+    if args.single_process and jax.local_device_count() > 1:
+        from transformer_engine.jax.cpp_extensions.ep import use_nccl_comm_from_xla
+
+        if not use_nccl_comm_from_xla():
+            print(
+                "[ep_moe] SKIPPED: --single-process needs the XLA-borrowed-comm EP path"
+                " (unavailable: build TE with the XLA collectives FFI headers and a"
+                " supporting JAX/jaxlib)"
+            )
+            return
+
     args.mesh, args.mr = _build_mesh_and_resource(args)
 
     with args.mesh, global_shard_guard(args.mr):

examples/jax/ep/run_test_ep.sh

@@ -91,4 +91,35 @@ else
   echo "... ep_moe PASSED"
 fi
 rm -f stdout_rank_*.txt
+
+echo
+echo "*** Executing ep_moe.py --single-process across $NUM_GPUS local GPUs ***"
+# --single-process drives every local GPU from one process; no coordinator/rank
+# loop needed. Pin to exactly NUM_GPUS devices so the default (2,2) mesh applies
+# even on larger boxes, matching the multi-process run above. Take the first
+# NUM_GPUS entries of the caller's existing CUDA_VISIBLE_DEVICES (Slurm/k8s/CI
+# may have bound specific physical GPUs); only fall back to raw ordinals
+# 0..NUM_GPUS-1 if nothing was set.
+if [ -n "${CUDA_VISIBLE_DEVICES:-}" ]; then
+  SINGLE_PROCESS_DEVICES=$(echo "$CUDA_VISIBLE_DEVICES" | cut -d',' -f"1-${NUM_GPUS}")
+else
+  SINGLE_PROCESS_DEVICES=$(seq -s, 0 $((NUM_GPUS - 1)))
+fi
+timeout --foreground --signal=KILL "${TEST_TIMEOUT_S}" \
+  env CUDA_VISIBLE_DEVICES="$SINGLE_PROCESS_DEVICES" python -u "$SCRIPT" --single-process \
+    $EXTRA_ARGS 2>&1 | tee stdout_single_process.txt
+
+if grep -qE "FAILED|Traceback|ERROR" stdout_single_process.txt; then
+  echo "... ep_moe --single-process FAILED"
+  HAS_FAILURE=1
+elif grep -q "SKIPPED" stdout_single_process.txt; then
+  echo "... ep_moe --single-process SKIPPED"
+elif ! grep -qE "\[ep_moe\]" stdout_single_process.txt; then
+  echo "... ep_moe --single-process INVALID (no summary line)"
+  HAS_FAILURE=1
+else
+  echo "... ep_moe --single-process PASSED"
+fi
+rm -f stdout_single_process.txt
+
 exit $HAS_FAILURE

qa/L0_jax_unittest/test.sh

@@ -29,6 +29,10 @@ pip3 install pytest==8.2.1 pytest-timeout==2.4.0 || error_exit "Failed to instal
 mkdir -p "$XML_LOG_DIR"
 
 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax --ignore=$TE_PATH/tests/jax/test_multi_process_ep.py -k 'not distributed' || test_fail "tests/jax/*not_distributed_*"
+# GPU-free EP unit tests (comm-path gating precedence, multi-device guard, comm-path-switch
+# guard); test_multi_process_ep.py is --ignore'd above because the rest of that file needs
+# jax.distributed.initialize (see multi_process_launch_ep.sh), but these three classes don't.
+python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_ep_unit.xml $TE_PATH/tests/jax/test_multi_process_ep.py -k 'TestEpCommSelection or TestEpBootstrapMultiDeviceGuard or TestEpCommPathSwitchGuard' || test_fail "tests/jax/test_multi_process_ep.py (GPU-free EP unit tests)"
 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fused_attn_score_mod.xml $TE_PATH/tests/jax/test_fused_attn_score_mod.py || test_fail "tests/jax/test_fused_attn_score_mod.py"
 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fused_attn_with_determinism.xml $TE_PATH/tests/jax/test_fused_attn.py -k "TestFusedAttnWithDeterminism" || test_fail "tests/jax/test_fused_attn.py"
 

qa/L2_jax_distributed_unittest/test.sh

@@ -15,6 +15,36 @@ mkdir -p "$XML_LOG_DIR"
 XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
 
 # NCCL EP multi-process suite. The launcher skips when fewer than 4 GPUs or no NVLink is detected.
-# Runs the borrowed-comm suite too (L2 only).
 export NVTE_JAX_UNITTEST_LEVEL="L2"
-TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh
+
+# Self-hosted NCCL classes (one process per GPU).
+NVTE_TEST_EP_CLASSES="TestEP,TestEPOverflowDrop,TestEpDomainGrouping" \
+    TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh
+
+# Borrowed-comm classes, in a fresh process group: self-hosted NCCL
+# teardown followed by a borrowed-comm re-init in the *same* process is
+# fragile (observed a spurious combine-backward mismatch), so keep them
+# process-isolated rather than chasing it -- real deployments don't mix
+# the two modes in one process either.
+NVTE_TEST_EP_CLASSES="TestEPBorrowedComm,TestEpDomainGrouping" \
+    TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh
+
+# Same 4 devices split 2 processes x 2 GPUs each (borrowed-comm-only classes;
+# self-hosted NCCL doesn't support >1 local device per process). dp=2,ep=2
+# here keeps each EP domain inside one process (JAX groups devices by
+# process), so this alone doesn't exercise EP traffic crossing a process.
+NVTE_TEST_EP_DEVICES_PER_PROC=2 NVTE_TEST_EP_CLASSES="TestEPBorrowedComm,TestEpDomainGrouping" \
+    TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh
+
+# Same 2x2 split, but ep=4 (dp=1): the single EP domain spans devices
+# [0,1,2,3], forcing NCCL EP traffic across the process boundary.
+NVTE_TEST_EP_MESH=1x4 NVTE_TEST_EP_DEVICES_PER_PROC=2 \
+    NVTE_TEST_EP_CLASSES="TestEPBorrowedComm,TestEpDomainGrouping" \
+    TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh
+
+# All 4 devices in a single process (DEVICES_PER_PROC=4 -> NUM_RUNS=1): the
+# single-controller, multi-domain scenario -- one process bootstraps and
+# drives two independent EP domains (dp=2, ep=2) across its 4 local GPUs.
+NVTE_TEST_EP_DEVICES_PER_PROC=4 \
+    NVTE_TEST_EP_CLASSES="TestEpSingleProcessMultiDomain,TestEpDomainGrouping" \
+    TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh

tests/jax/multi_process_launch_ep.sh

@@ -49,8 +49,18 @@ if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; the
   exit 0
 fi
 
-# Default test mesh is (2, 2); use exactly 4 ranks even on larger boxes.
-NUM_RUNS="${NVTE_TEST_EP_NUM_RANKS:-4}"
+# Devices per process: 1 = one-process-per-GPU (default); >1 tests multiple
+# local devices per process (e.g. one process per MNNVL-connected node).
+DEVICES_PER_PROC="${NVTE_TEST_EP_DEVICES_PER_PROC:-1}"
+if [ "$DEVICES_PER_PROC" -le 0 ] || [ "$((4 % DEVICES_PER_PROC))" -ne 0 ]; then
+  echo "ERROR: NVTE_TEST_EP_DEVICES_PER_PROC=${DEVICES_PER_PROC} must be a positive divisor of 4" \
+       "(the default test mesh's total device count)."
+  exit 1
+fi
+
+# Default test mesh is (2, 2) across 4 devices total; use exactly that many
+# devices even on larger boxes, split into NUM_RUNS processes.
+NUM_RUNS="${NVTE_TEST_EP_NUM_RANKS:-$((4 / DEVICES_PER_PROC))}"
 
 OVERALL_RET=0
 
@@ -66,11 +76,13 @@ for SCRIPT_NAME in $SCRIPT_NAMES; do
   for ((i=1; i<NUM_RUNS; i++))
   do
       timeout --foreground --signal=KILL "${TEST_TIMEOUT_S}" \
-          python "$SCRIPT_PATH" 127.0.0.1:12345 $i $NUM_RUNS > stdout_rank_${i}.txt 2>&1 &
+          python "$SCRIPT_PATH" 127.0.0.1:12345 $i $NUM_RUNS $DEVICES_PER_PROC \
+              > stdout_rank_${i}.txt 2>&1 &
   done
 
   timeout --foreground --signal=KILL "${TEST_TIMEOUT_S}" \
-      python "$SCRIPT_PATH" 127.0.0.1:12345 0 $NUM_RUNS 2>&1 | tee stdout_multi_process.txt
+      python "$SCRIPT_PATH" 127.0.0.1:12345 0 $NUM_RUNS $DEVICES_PER_PROC \
+          2>&1 | tee stdout_multi_process.txt
 
   wait
 
@@ -84,6 +96,15 @@ for SCRIPT_NAME in $SCRIPT_NAMES; do
     echo "       NCCL EP requires NVLS multicast; check NCCL_DEBUG=INFO output."
     RET=1
   fi
+  # "Ran N tests" counts skipped tests too, so a run where every test skips
+  # (e.g. no NCCL EP build) still matches the check above with zero real
+  # coverage. Fail explicitly instead of reporting a green PASS.
+  RAN_N=$(grep -oE "Ran [0-9]+ test" stdout_multi_process.txt | tail -1 | grep -oE '[0-9]+')
+  SKIPPED_N=$(grep -oE "skipped=[0-9]+" stdout_multi_process.txt | tail -1 | grep -oE '[0-9]+')
+  if [ -n "$RAN_N" ] && [ "${SKIPPED_N:-0}" -ge "$RAN_N" ]; then
+    echo "ERROR: all ${RAN_N} test(s) skipped for ${SCRIPT_NAME} — zero real coverage."
+    RET=1
+  fi
   if [ "$RET" -ne 0 ]; then
     for ((i=1; i<NUM_RUNS; i++)); do
       echo "--- rank $i log ---"