Skip to content

[JAX] Keep the routing map format alive and EP multiprocess tests in L2 Jax dist #3159

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
KshitijLakhani merged 2 commits into from
Jul 1, 2026
Merged

[JAX] Keep the routing map format alive and EP multiprocess tests in L2 Jax dist #3159

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4448ab7
Keep the routing map format alive
KshitijLakhani Jul 1, 2026
76a1a7d
Fix incorrectly launched multi process EP tests in L2 Jax instead of …
KshitijLakhani Jul 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions qa/L2_jax_distributed_unittest/test.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,3 +13,6 @@ mkdir -p "$XML_LOG_DIR"

# Use --xla_gpu_enable_triton_gemm=false to ensure the reference JAX implementation we are using is accurate.
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.
TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh
2 changes: 1 addition & 1 deletion qa/L2_jax_unittest/test.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
: ${XML_LOG_DIR:=/logs}
mkdir -p "$XML_LOG_DIR"

NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax -k 'not distributed' || test_fail "tests/jax/*not_distributed_*"
NVTE_JAX_UNITTEST_LEVEL="L2" 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_*"

pip3 install -r $TE_PATH/examples/jax/mnist/requirements.txt || error_exit "Failed to install mnist requirements"
# Note: mnist intentionally does NOT set --xla_gpu_deterministic_ops because it
Expand Down
2 changes: 1 addition & 1 deletion transformer_engine/jax/cpp_extensions/router.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -412,7 +412,7 @@ def partition(
arg_infos,
result_infos,
):
del result_infos, routing_map_format
del result_infos
grad_spec = get_padded_spec(arg_infos[2])
out_sharding = NamedSharding(mesh, PartitionSpec(*grad_spec))
arg_shardings = (arg_infos[0].sharding, arg_infos[1].sharding, arg_infos[2].sharding)
Expand Down
Loading