TE EP integration to MoEBlock

qa/L0_jax_distributed_unittest/test.sh

@@ -41,12 +41,12 @@ wait
 TE_PATH=$TE_PATH bash $TE_PATH/examples/jax/ep/run_test_ep.sh || test_fail "run_test_ep.sh"
 wait
 
-# MoE custom_vjp distributed suite. Runs one Python process per GPU
-# via tests/jax/run_multiprocess_moe_vjp.sh (mirrors the pattern in
+# TE-EP MoE custom_vjp distributed suite. Runs one Python process per
+# GPU via tests/jax/run_te_ep_moe.sh (mirrors the pattern in
 # examples/jax/encoder/run_test_multiprocessing_encoder.sh). Requires
 # >=4 visible GPUs.
-TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_multiprocess_moe_vjp.sh \
-    || test_fail "test_multiprocess_moe_vjp.py"
+TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_te_ep_moe.sh \
+    || test_fail "test_te_ep_moe.py"
 # Exercise the multi-GPU tutorial in docs/examples/jax (needs >= 4 GPUs;
 # auto-skips otherwise).
 CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_docs_examples_jax_distributed.xml -k multi_gpu $TE_PATH/docs/examples/jax/ || test_fail "docs/examples/jax (multi-GPU)"

tests/jax/conftest.py

@@ -90,8 +90,8 @@ def pytest_addoption(parser):
     """CLI options used by multiprocess JAX tests.
 
     ``--num-process`` and ``--process-id`` let a multiprocess launcher
-    (see ``tests/jax/run_multiprocess_moe_vjp.sh``) fork one pytest
-    process per GPU and tell each child its rank, so the test module
+    (see ``tests/jax/run_te_ep_moe.sh``) fork one pytest process per
+    GPU and tell each child its rank, so the test module
     can call ``jax.distributed.initialize(...)`` with the right
     ``local_device_ids``. Both default to 0; non-multiprocess tests
     ignore them.

tests/jax/run_multiprocess_moe_vjp.sh → tests/jax/run_te_ep_moe.sh

@@ -3,46 +3,43 @@
 #
 # See LICENSE for license information.
 #
-# Multiprocess (one-GPU-per-process) launcher for the unified MoE VJP
+# Multiprocess (one-GPU-per-process) launcher for the TE-EP MoE custom_vjp
 # test suite. Forks one pytest invocation per visible GPU, passing each
-# its own --num-process=N --process-id=i, and waits for all of them.
-# Each child calls jax.distributed.initialize(..., local_device_ids=
-# process_id) so each Python process only sees its one GPU as a local
-# device and the participating processes form a global mesh.
+# its own --num-process=N --process-id=i, and waits for all of them. Each
+# child calls jax.distributed.initialize(..., local_device_ids=process_id)
+# so each Python process only sees its one GPU as a local device and the
+# participating processes form a global (ep, fsdp) mesh.
 
 set -euo pipefail
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 TE_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
-TEST_FILE="$TE_ROOT/tests/jax/test_multiprocess_moe_vjp.py"
+TEST_FILE="$TE_ROOT/tests/jax/test_te_ep_moe.py"
 PYTEST_INI="$TE_ROOT/tests/jax/pytest.ini"
 
 NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}"
 if [ "$NUM_GPUS" -lt 4 ]; then
-    echo "[run_multiprocess_moe_vjp.sh] need >=4 GPUs (got $NUM_GPUS); aborting" >&2
+    echo "[run_te_ep_moe.sh] need >=4 GPUs (got $NUM_GPUS); aborting" >&2
     exit 1
 fi
 
 export XLA_PYTHON_CLIENT_PREALLOCATE="${XLA_PYTHON_CLIENT_PREALLOCATE:-false}"
 export XLA_PYTHON_CLIENT_MEM_FRACTION="${XLA_PYTHON_CLIENT_MEM_FRACTION:-0.5}"
-export MOE_VJP_COORDINATOR_ADDRESS="${MOE_VJP_COORDINATOR_ADDRESS:-127.0.0.1:13456}"
+export TE_EP_MOE_COORDINATOR_ADDRESS="${TE_EP_MOE_COORDINATOR_ADDRESS:-127.0.0.1:13457}"
 
 echo "============================================================"
-echo "MoE VJP MULTIPROCESS test (one process per GPU, ${NUM_GPUS} GPUs)"
+echo "TE-EP MoE MULTIPROCESS test (one process per GPU, ${NUM_GPUS} GPUs)"
 echo "  test file          : $TEST_FILE"
-echo "  coordinator        : $MOE_VJP_COORDINATOR_ADDRESS"
+echo "  coordinator        : $TE_EP_MOE_COORDINATOR_ADDRESS"
 echo "  XLA_PYTHON_CLIENT_PREALLOCATE: $XLA_PYTHON_CLIENT_PREALLOCATE"
 echo "  XLA_PYTHON_CLIENT_MEM_FRACTION: $XLA_PYTHON_CLIENT_MEM_FRACTION"
 echo "============================================================"
 
-# Per-process logs. MOE_VJP_MP_LOG_DIR can be set to a host-mounted dir
-# (e.g. when running inside a container that throws away /tmp on exit)
-# so logs survive for postmortem inspection. Defaults to a fresh /tmp.
-if [ -n "${MOE_VJP_MP_LOG_DIR:-}" ]; then
-    LOG_DIR="$MOE_VJP_MP_LOG_DIR"
+if [ -n "${TE_EP_MOE_MP_LOG_DIR:-}" ]; then
+    LOG_DIR="$TE_EP_MOE_MP_LOG_DIR"
     mkdir -p "$LOG_DIR"
 else
-    LOG_DIR=$(mktemp -d -t moe_vjp_mp_XXXXXX)
+    LOG_DIR=$(mktemp -d -t te_ep_moe_mp_XXXXXX)
 fi
 echo "Per-process logs: $LOG_DIR"
 
@@ -63,8 +60,6 @@ cleanup() {
 }
 trap cleanup EXIT INT TERM
 
-# Launch one pytest per GPU. Process 0 streams to stdout; others log
-# only to file so the live output isn't a mosaic.
 for i in $(seq 0 $((NUM_GPUS - 1))); do
     LOG_FILE="$LOG_DIR/proc_${i}.log"
     PYTEST_CMD=(
@@ -84,7 +79,6 @@ for i in $(seq 0 $((NUM_GPUS - 1))); do
     PIDS+=("$!")
 done
 
-# Wait for all and collect exit codes.
 EXITS=()
 for pid in "${PIDS[@]}"; do
     if wait "$pid"; then
@@ -94,20 +88,16 @@ for pid in "${PIDS[@]}"; do
     fi
 done
 
-# Summary.
 echo
 echo "============================================================"
 echo "Per-process exit codes:"
 for i in "${!EXITS[@]}"; do
     echo "  proc $i -> ${EXITS[$i]}"
 done
 
-# Final pass/fail. Any non-zero in any process fails the suite, but
-# we tolerate non-zero on the non-zero processes only if proc 0
-# reports PASS (this matches the encoder launcher's logic). Simplest
-# Treat exit 0 (pass) and exit 5 (pytest "no tests collected", which
-# the file emits via ``pytest.skip(allow_module_level=True)`` on
-# pre-Blackwell GPUs) as success. Anything else is a failure.
+# Treat exit 0 (pass) and exit 5 (pytest "no tests collected", which the
+# file emits via pytest.skip(allow_module_level=True) on pre-Blackwell
+# GPUs) as success.
 FAILED=0
 for e in "${EXITS[@]}"; do
     if [ "$e" != "0" ] && [ "$e" != "5" ]; then
@@ -118,14 +108,14 @@ done
 
 echo
 if [ "$FAILED" -eq 0 ]; then
-    echo "[run_multiprocess_moe_vjp.sh] all processes PASSED"
-    if [ -z "${MOE_VJP_MP_LOG_DIR:-}" ]; then
+    echo "[run_te_ep_moe.sh] all processes PASSED"
+    if [ -z "${TE_EP_MOE_MP_LOG_DIR:-}" ]; then
         rm -rf "$LOG_DIR"
     fi
     exit 0
 fi
 
-echo "[run_multiprocess_moe_vjp.sh] at least one process FAILED"
+echo "[run_te_ep_moe.sh] at least one process FAILED"
 echo "  retaining logs at $LOG_DIR for diagnosis"
 echo "  process 0 tail:"
 tail -20 "$LOG_DIR/proc_0.log" 2>/dev/null || true