From d4cbc5b5ac37803728c4bbdaa864299752eaad88 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 14 Aug 2026 13:41:18 -0700 Subject: [PATCH 1/4] Surface distributed comm-overlap rank errors Comm-overlap launchers captured child stdout and stderr but raised only stderr, which could leave the originating rank context out of pytest and JUnit failures. Preserve the existing result predicates while attaching bounded tails of both streams, and record layer-worker exceptions so torchrun can report the rank-local traceback. Signed-off-by: Sudhakar Singh --- .../distributed/run_layer_with_overlap.py | 2 + .../distributed/test_comm_gemm_overlap.py | 40 ++++++++++++------- 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index 46795415e5..40e826eff6 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -17,6 +17,7 @@ import torch import torch.distributed as dist +from torch.distributed.elastic.multiprocessing.errors import record import transformer_engine.pytorch as te from transformer_engine.common.recipe import ( @@ -348,6 +349,7 @@ def _compare_tensors(name, test, ref, rtol, atol): return numerics_failed, numerics_info +@record def _train(opts): if "OMPI_COMM_WORLD_SIZE" in os.environ: # Execution with `mpirun -np N` diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 6b1ad870e9..7199fd6e31 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -39,6 +39,24 @@ if tex.ubuf_built_with_mpi(): LAUNCH_CMD = ["mpirun", "-np", str(NUM_PROCS), "--oversubscribe", "--quiet", "python3"] +OUTPUT_TAIL_CHARS = 4000 + + +def _assert_subprocess_succeeded(result): + if ( + result.returncode != 0 + or "NUMERICAL CHECK FAILED" in result.stderr + or "NUMERICAL CHECK PASSED" not in result.stdout + ): + raise AssertionError( + f"Distributed test exited with return code {result.returncode}" + f"\n--- stdout (last {OUTPUT_TAIL_CHARS} characters) ---\n" + f"{result.stdout[-OUTPUT_TAIL_CHARS:]}" + f"\n--- stderr (last {OUTPUT_TAIL_CHARS} characters) ---\n" + f"{result.stderr[-OUTPUT_TAIL_CHARS:]}" + ) + + # Fall back on CUDA IPC if the platform does not support CUDA multicast if not tex.device_supports_multicast(): os.environ["UB_SKIPMC"] = "1" @@ -94,13 +112,10 @@ def _run_gemm_with_overlap( ) test_cmd.append("--use-cublasmp") - result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False) - if ( - result.returncode != 0 - or "NUMERICAL CHECK FAILED" in result.stderr.decode() - or "NUMERICAL CHECK PASSED" not in result.stdout.decode() - ): - raise AssertionError(result.stderr.decode()) + result = subprocess.run( + test_cmd, env=os.environ, capture_output=True, text=True, check=False + ) + _assert_subprocess_succeeded(result) def _run_layer_with_overlap( @@ -157,7 +172,9 @@ def _run_layer_with_overlap( # NVTE_ALLOW_NONDETERMINISTIC_ALGO=0, so disable it entirely for this test. os.environ["NVTE_FUSED_ATTN"] = "0" - result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False) + result = subprocess.run( + test_cmd, env=os.environ, capture_output=True, text=True, check=False + ) os.unsetenv("PYTORCH_JIT") os.unsetenv("NVTE_TORCH_COMPILE") @@ -165,12 +182,7 @@ def _run_layer_with_overlap( os.unsetenv("NVTE_FLASH_ATTN") os.unsetenv("NVTE_FUSED_ATTN") - if ( - result.returncode != 0 - or "NUMERICAL CHECK FAILED" in result.stderr.decode() - or "NUMERICAL CHECK PASSED" not in result.stdout.decode() - ): - raise AssertionError(result.stderr.decode()) + _assert_subprocess_succeeded(result) @pytest.mark.parametrize("use_cublasmp", (False, True)) From 602bfdff7b13e7273c56d45debd2edc0a3f0c54c Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 14 Aug 2026 15:27:44 -0700 Subject: [PATCH 2/4] Isolate comm-overlap launcher environment The launcher mutated os.environ and then called os.unsetenv, which leaves Python's environment mapping unchanged. Conditional backend flags could therefore leak into later parameterized children. Build a child-only environment instead so each launch gets its intended overrides while preserving the parent process environment. Signed-off-by: Sudhakar Singh --- .../distributed/test_comm_gemm_overlap.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 7199fd6e31..4827f11a17 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -159,29 +159,24 @@ def _run_layer_with_overlap( pytest.skip("cuBLASMp comm+GEMM overlap does not yet support MXFP8 (block scaling).") test_cmd.append("--use-cublasmp") - os.environ["PYTORCH_JIT"] = "0" - os.environ["NVTE_TORCH_COMPILE"] = "0" - os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" + test_env = os.environ.copy() + test_env["PYTORCH_JIT"] = "0" + test_env["NVTE_TORCH_COMPILE"] = "0" + test_env["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" if te.get_device_compute_capability() <= (8, 0): # We've experienced numerical discrepancies in Flash Attention # backward when running with Userbuffers on A100s. This does # not show up in more recent GPUs. - os.environ["NVTE_FLASH_ATTN"] = "0" + test_env["NVTE_FLASH_ATTN"] = "0" elif fp8: # Fused attention is causing non-deterministic FP8 failures on H100s even with # NVTE_ALLOW_NONDETERMINISTIC_ALGO=0, so disable it entirely for this test. - os.environ["NVTE_FUSED_ATTN"] = "0" + test_env["NVTE_FUSED_ATTN"] = "0" result = subprocess.run( - test_cmd, env=os.environ, capture_output=True, text=True, check=False + test_cmd, env=test_env, capture_output=True, text=True, check=False ) - os.unsetenv("PYTORCH_JIT") - os.unsetenv("NVTE_TORCH_COMPILE") - os.unsetenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO") - os.unsetenv("NVTE_FLASH_ATTN") - os.unsetenv("NVTE_FUSED_ATTN") - _assert_subprocess_succeeded(result) From 3a4023a6cb9449ac19e80b7b9c51803e256b5edb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Aug 2026 22:35:27 +0000 Subject: [PATCH 3/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/test_comm_gemm_overlap.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 4827f11a17..1c50ab405f 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -112,9 +112,7 @@ def _run_gemm_with_overlap( ) test_cmd.append("--use-cublasmp") - result = subprocess.run( - test_cmd, env=os.environ, capture_output=True, text=True, check=False - ) + result = subprocess.run(test_cmd, env=os.environ, capture_output=True, text=True, check=False) _assert_subprocess_succeeded(result) @@ -173,9 +171,7 @@ def _run_layer_with_overlap( # NVTE_ALLOW_NONDETERMINISTIC_ALGO=0, so disable it entirely for this test. test_env["NVTE_FUSED_ATTN"] = "0" - result = subprocess.run( - test_cmd, env=test_env, capture_output=True, text=True, check=False - ) + result = subprocess.run(test_cmd, env=test_env, capture_output=True, text=True, check=False) _assert_subprocess_succeeded(result) From 45713b1600403e70085dd1771086e4553d93c563 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Sat, 15 Aug 2026 01:53:07 -0700 Subject: [PATCH 4/4] Keep FA4 out of L1 distributed tests The moving PyTorch image installs FA4 by default, implicitly expanding an L1 suite that historically covered earlier attention backends. Export the FA4 selector at the suite boundary so every pytest and torchrun child retains the intended backend scope while dedicated attention suites own FA4 coverage. Signed-off-by: Sudhakar Singh --- qa/L1_pytorch_distributed_unittest/test.sh | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index ec19492ee7..e0c92849f8 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -18,28 +18,29 @@ FAILED_CASES="" : ${TE_PATH:=/opt/transformerengine} : ${XML_LOG_DIR:=/logs} +# FA4 coverage belongs to the dedicated attention-backend suites. +export NVTE_FLASH_ATTN_V4=0 mkdir -p "$XML_LOG_DIR" pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" # Run CP tests (deterministic + non-deterministic) first so they can be parallelized. # Each needs 4 GPUs, so >=8 GPUs allows them to run concurrently on disjoint GPU sets. -# Main's CP implementation supports FA2/FA3. Keep FA4 disabled at the suite -# boundary so both the reference and CP halves use the same backend generation. +# Main's CP implementation supports FA2/FA3. NUM_GPUS=$(python3 -c "import torch; print(torch.cuda.device_count())") echo "Detected $NUM_GPUS GPU(s)" if [ "$NUM_GPUS" -ge 8 ]; then echo "Running CP tests in parallel: non-deterministic on GPUs 0-3, deterministic on GPUs 4-7" - CUDA_VISIBLE_DEVICES=0,1,2,3 NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & + CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & PID_CP_NONDET=$! - CUDA_VISIBLE_DEVICES=4,5,6,7 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & + CUDA_VISIBLE_DEVICES=4,5,6,7 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & PID_CP_DET=$! wait $PID_CP_NONDET || test_fail "test_attention_with_cp.py" wait $PID_CP_DET || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" else echo "Running CP tests sequentially: need >=8 GPUs for parallel execution" - NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py" - NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" + python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py" + NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" fi python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/distributed/test_sanity.py || test_fail "test_sanity.py"