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" 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..1c50ab405f 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,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))