surface comm overlap rank errors

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"

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`

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,8 @@ 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(
@@ -144,33 +157,23 @@ 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"
-
-    result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False)
+        test_env["NVTE_FUSED_ATTN"] = "0"
 
-    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")
+    result = subprocess.run(test_cmd, env=test_env, capture_output=True, text=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())
+    _assert_subprocess_succeeded(result)
 
 
 @pytest.mark.parametrize("use_cublasmp", (False, True))