From 390529aa02591c037f1a9c3a88f0b3d31492d0ea Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Tue, 27 Oct 2026 10:37:21 -0700 Subject: [PATCH 1/9] Skip the DSA KV-cache quantization test on Blackwell It intermittently hits a CUDA illegal memory access in the DSA forward during calibration on the Blackwell CI runners. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- tests/gpu_megatron/torch/quantization/plugins/test_megatron.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 24bbe809b44..47ae586b0ec 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -2306,6 +2306,8 @@ def _test_dsa_kv_cache_quant_helper(tmp_path, rank, size): assert all(m.k_bmm_quantizer.is_enabled and m.v_bmm_quantizer.is_enabled for m in restored) +# Intermittent since 2026-10-27 (5 of ~45 CI runs), always in the DSA forward during calibration. +@skip_flaky_on_blackwell def test_dsa_kv_cache_quant(dist_workers_size_1, tmp_path): """DSAttention gets calibrated K/V KV-cache quantizers that survive a checkpoint round-trip.""" pytest.importorskip("megatron.core.transformer.experimental_attention_variant.dsa") From b93f6437ab1e9767f1a6d186cc476e5cd01901f4 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Tue, 27 Oct 2026 10:59:57 -0700 Subject: [PATCH 2/9] Run the flaky Blackwell megatron tests, tolerating only known CUDA crashes Replace the Blackwell skip with a flaky_cuda_crash_on_blackwell marker: on sm_120, a marked test whose failure is a known CUDA crash (illegal memory access, CUBLAS_STATUS_INTERNAL_ERROR, dead worker) is reported as xfail, and any other failure still fails. The worker pool is replaced after a crash, so it no longer cascades into later tests. test_regular_state_dict[320] stays skipped since it runs in the pytest process itself. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- tests/gpu_megatron/conftest.py | 35 +++++++++++++++++++ .../torch/peft/plugins/test_megatron_peft.py | 2 ++ .../quantization/plugins/test_megatron.py | 34 +++++++++--------- 3 files changed, 53 insertions(+), 18 deletions(-) diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index 1af63e694e5..95f6b27497c 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -61,6 +61,41 @@ def _prebuild_quant_cuda_extensions(): ext.precompile() +_IS_BLACKWELL = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 12 +# Errors the nemo:26.08 TE/CUDA 13 stack intermittently raises on Blackwell (sm_120); not modelopt bugs +_BLACKWELL_CUDA_CRASHES = ( + "illegal memory access", + "CUBLAS_STATUS_INTERNAL_ERROR", + "Worker process(es) died", +) + + +def pytest_configure(config): + config.addinivalue_line( + "markers", + "flaky_cuda_crash_on_blackwell: on Blackwell, report a known CUDA crash as xfail, other failures still fail", + ) + + +@pytest.hookimpl(wrapper=True) +def pytest_runtest_makereport(item, call): + """Report a known Blackwell CUDA crash in a marked test as xfail instead of a failure. + + Only for tests run in a worker pool: a crashed pool is replaced, so later tests are unaffected. + """ + report = yield + if ( + _IS_BLACKWELL + and report.when == "call" + and report.failed + and item.get_closest_marker("flaky_cuda_crash_on_blackwell") + and any(err in report.longreprtext for err in _BLACKWELL_CUDA_CRASHES) + ): + report.outcome = "skipped" + report.wasxfail = "known flaky CUDA crash on Blackwell (nemo:26.08 TE/CUDA 13 stack)" + return report + + @pytest.fixture(scope="session", autouse=True) def _prestart_preprocess_forkserver(): """Start the forkserver that ``megatron_preprocess_data`` workers fork from, at session start. diff --git a/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py b/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py index 4568fa10168..9fd6f8d255a 100644 --- a/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py +++ b/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py @@ -483,6 +483,8 @@ def _test_adapter_gradient_flow_freeze_lora_model(lora_config, tmp_path, rank, s LARGE_LORA_CFG_RANDOM_INIT_TEST, # Use random init so gradients flow to both lora_a and lora_b ], ) +# Intermittent CUBLAS_STATUS_INTERNAL_ERROR in backward on Blackwell; see conftest.py +@pytest.mark.flaky_cuda_crash_on_blackwell def test_adapter_gradient_flow_freeze_lora_model(dist_workers, lora_config, tmp_path): dist_workers.run( partial(_test_adapter_gradient_flow_freeze_lora_model, lora_config, str(tmp_path)) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 47ae586b0ec..50f2b4b7e59 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -119,14 +119,8 @@ SEED = 1234 -# TODO: re-enable the marked tests once fixed. Blackwell (sm_120, e.g. RTX PRO 6000) with the -# nemo:26.06 / TE 2.16 / CUDA 13 stack hits a flaky, intermittent CUDA "illegal memory access" in -# several tests: When fails, it poisons the process CUDA context and cascades hang across subsequent tests. -# Passes locally on different GPUs. Likely an upstream TE/CUDA-13 kernel bug, not modelopt logic. -skip_flaky_on_blackwell = pytest.mark.skipif( - torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 12, - reason="Flaky CUDA illegal memory access on Blackwell (nemo:26.06 / TE 2.16 / CUDA 13)", -) +# Marked tests intermittently hit a CUDA illegal memory access on Blackwell (sm_120); see conftest.py. +flaky_on_blackwell = pytest.mark.flaky_cuda_crash_on_blackwell def test_convert_megatron_parallel_linear(distributed_setup_size_1): @@ -497,7 +491,7 @@ def _test_sharded_state_dict( ) @pytest.mark.parametrize("meta_device", [False, True]) @pytest.mark.parametrize("transformer_impl", ["local", "modelopt"]) -@skip_flaky_on_blackwell +@flaky_on_blackwell def test_homogeneous_sharded_state_dict( dist_workers, tmp_path, config, meta_device, transformer_impl ): @@ -523,9 +517,7 @@ def test_homogeneous_sharded_state_dict( @pytest.mark.parametrize("transformer_impl", ["local", "modelopt"]) @pytest.mark.timeout(240) # Compressed state dict takes longer due to real quant conversion & saving/loading -# Its real mtq.compress() path exercises the same TE/CUDA-13 kernels that trip the flaky -# Blackwell (sm_120) illegal-memory-access; #1901 skipped the fake-quant sibling but missed this one. -@skip_flaky_on_blackwell +@flaky_on_blackwell def test_homogeneous_compressed_sharded_state_dict( dist_workers, tmp_path, config, meta_device, transformer_impl ): @@ -607,7 +599,7 @@ def test_homogeneous_sharded_state_dict_hybrid(dist_workers, tmp_path, config): mixed_block_size_config, ], ) -@skip_flaky_on_blackwell +@flaky_on_blackwell def test_heterogenous_sharded_state_dict(dist_workers, tmp_path, config): dist_workers.run( partial(_test_sharded_state_dict, tmp_path, config, 256, None, False, False, {}), @@ -618,7 +610,14 @@ def test_heterogenous_sharded_state_dict(dist_workers, tmp_path, config): "hidden_size", [ 256, - pytest.param(320, marks=skip_flaky_on_blackwell), + # Runs in the pytest process, so a Blackwell CUDA crash here would break every later test + pytest.param( + 320, + marks=pytest.mark.skipif( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 12, + reason="Flaky CUDA illegal memory access on Blackwell", + ), + ), ], ) def test_regular_state_dict(distributed_setup_size_1, hidden_size): @@ -691,7 +690,7 @@ def _test_fp8_real_quantize_helper(rank, size): assert real_quant_mem < cur_mem -@skip_flaky_on_blackwell +@flaky_on_blackwell def test_fp8_real_quantize(dist_workers): dist_workers.run(_test_fp8_real_quantize_helper) @@ -1726,7 +1725,7 @@ def forward_loop(_model): assert all(m.input_quantizer.amax is not None for m in calibrated) -@skip_flaky_on_blackwell +@flaky_on_blackwell @pytest.mark.parametrize("drive_wrapper", [False, True]) def test_layerwise_calibrate_vlm_text_forward_loop(dist_workers_size_1, drive_wrapper): dist_workers_size_1.run(partial(_test_layerwise_calibrate_vlm_text_forward_loop, drive_wrapper)) @@ -2306,8 +2305,7 @@ def _test_dsa_kv_cache_quant_helper(tmp_path, rank, size): assert all(m.k_bmm_quantizer.is_enabled and m.v_bmm_quantizer.is_enabled for m in restored) -# Intermittent since 2026-10-27 (5 of ~45 CI runs), always in the DSA forward during calibration. -@skip_flaky_on_blackwell +@flaky_on_blackwell def test_dsa_kv_cache_quant(dist_workers_size_1, tmp_path): """DSAttention gets calibrated K/V KV-cache quantizers that survive a checkpoint round-trip.""" pytest.importorskip("megatron.core.transformer.experimental_attention_variant.dsa") From ef28231aa80b06934459efbbaa423f7ab0e8bab9 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Tue, 27 Oct 2026 11:27:33 -0700 Subject: [PATCH 3/9] Match only the CUDA error message, not worker deaths, in the Blackwell hook A dead worker (e.g. an OOM kill) is no longer reported as a known Blackwell crash, and the hook matches the exception message instead of the whole report, whose traceback source lines can mention these errors. Share one is_blackwell() check, evaluated lazily in the hook. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- tests/_test_utils/torch/misc.py | 7 +++++++ tests/gpu_megatron/conftest.py | 16 +++++++--------- .../torch/peft/plugins/test_megatron_peft.py | 3 ++- .../torch/quantization/plugins/test_megatron.py | 5 ++--- 4 files changed, 18 insertions(+), 13 deletions(-) diff --git a/tests/_test_utils/torch/misc.py b/tests/_test_utils/torch/misc.py index e20c6919fbd..56cc3dc3be1 100644 --- a/tests/_test_utils/torch/misc.py +++ b/tests/_test_utils/torch/misc.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import random import pytest @@ -53,6 +54,12 @@ def set_seed(seed_value=42): random.seed(seed_value) +@functools.cache +def is_blackwell(): + """Whether the GPU is sm_120 or newer (e.g. RTX PRO 6000), as on the GPU CI runners.""" + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 12 + + def minimum_sm(sm): major, minor = sm // 10, sm % 10 return pytest.mark.skipif( diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index 95f6b27497c..a2f389f021e 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -24,6 +24,7 @@ get_live_pool, reset_worker_state, ) +from _test_utils.torch.misc import is_blackwell from _test_utils.torch.transformers_models import get_tiny_tokenizer from megatron.core.parallel_state import destroy_model_parallel @@ -61,13 +62,8 @@ def _prebuild_quant_cuda_extensions(): ext.precompile() -_IS_BLACKWELL = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 12 # Errors the nemo:26.08 TE/CUDA 13 stack intermittently raises on Blackwell (sm_120); not modelopt bugs -_BLACKWELL_CUDA_CRASHES = ( - "illegal memory access", - "CUBLAS_STATUS_INTERNAL_ERROR", - "Worker process(es) died", -) +_BLACKWELL_CUDA_CRASHES = ("illegal memory access", "CUBLAS_STATUS_INTERNAL_ERROR") def pytest_configure(config): @@ -85,11 +81,13 @@ def pytest_runtest_makereport(item, call): """ report = yield if ( - _IS_BLACKWELL - and report.when == "call" + report.when == "call" and report.failed and item.get_closest_marker("flaky_cuda_crash_on_blackwell") - and any(err in report.longreprtext for err in _BLACKWELL_CUDA_CRASHES) + and call.excinfo is not None + # The message, not report.longreprtext: the traceback source lines can mention these errors + and any(err in str(call.excinfo.value) for err in _BLACKWELL_CUDA_CRASHES) + and is_blackwell() ): report.outcome = "skipped" report.wasxfail = "known flaky CUDA crash on Blackwell (nemo:26.08 TE/CUDA 13 stack)" diff --git a/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py b/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py index 9fd6f8d255a..c9f5df7459c 100644 --- a/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py +++ b/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py @@ -25,6 +25,7 @@ load_distributed_checkpoint, save_distributed_checkpoint, ) +from _test_utils.torch.misc import is_blackwell import modelopt.torch.peft as mtpeft import modelopt.torch.quantization as mtq @@ -40,7 +41,7 @@ # process CUDA context and cascades into every test after it. Deterministic across CI runs; passes on # sm_89 with the same container and on nemo:26.06. skip_on_blackwell = pytest.mark.skipif( - torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 12, + is_blackwell(), reason="CUDA illegal memory access on Blackwell (nemo:26.08 TE/CUDA stack)", ) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 50f2b4b7e59..6d829276766 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -42,7 +42,7 @@ save_distributed_checkpoint, sharded_state_dict_test_helper, ) -from _test_utils.torch.misc import set_seed +from _test_utils.torch.misc import is_blackwell, set_seed from _test_utils.torch.quantization.models import RegularQuantModelForTP from _test_utils.torch.quantization.quant_utils import get_model_size from _test_utils.torch.quantization.quantize_common import ( @@ -614,8 +614,7 @@ def test_heterogenous_sharded_state_dict(dist_workers, tmp_path, config): pytest.param( 320, marks=pytest.mark.skipif( - torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 12, - reason="Flaky CUDA illegal memory access on Blackwell", + is_blackwell(), reason="Flaky CUDA illegal memory access on Blackwell" ), ), ], From 20f7554dead64deb876f569d6e0a0b6b174c8683 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:54:44 -0700 Subject: [PATCH 4/9] Rerun known Blackwell CUDA crashes in gpu_megatron instead of xfailing them On Blackwell, every gpu_megatron test gets one rerun on a fresh worker when it fails with a known CUDA crash (illegal memory access, CUBLAS_STATUS_INTERNAL_ERROR, or a worker abort). Any other failure, or the same crash twice, still fails. This replaces the flaky_cuda_crash_on_blackwell marker and xfail hook, since the crash depends on how many jobs a worker has run and can hit any test. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- pyproject.toml | 1 + tests/gpu_megatron/conftest.py | 42 +++++++------------ .../torch/peft/plugins/test_megatron_peft.py | 2 - .../quantization/plugins/test_megatron.py | 9 ---- 4 files changed, 15 insertions(+), 39 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1ca6a7de9d8..ca75164ffb5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,6 +122,7 @@ dev-test = [ "pytest~=9.1.0", "pytest-cov~=7.1.0", "pytest-instafail==0.5.0", + "pytest-rerunfailures~=16.7", # rerun known Blackwell cuBLAS crashes (tests/gpu_megatron) "pytest-timeout~=2.4.0", "pytest-xdist~=3.8.0", "uv", diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index a2f389f021e..802c474c1d5 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -15,6 +15,7 @@ import contextlib import gc import queue +from pathlib import Path import pytest import torch @@ -62,36 +63,21 @@ def _prebuild_quant_cuda_extensions(): ext.precompile() -# Errors the nemo:26.08 TE/CUDA 13 stack intermittently raises on Blackwell (sm_120); not modelopt bugs -_BLACKWELL_CUDA_CRASHES = ("illegal memory access", "CUBLAS_STATUS_INTERNAL_ERROR") +# On Blackwell (sm_120), nemo:26.08's cuBLAS intermittently fails a GEMM once a worker process has +# run about a dozen model-building jobs. A failed job replaces the worker, so rerun it once on a fresh one. +_BLACKWELL_CUDA_CRASHES = [ + "illegal memory access", + "CUBLAS_STATUS_INTERNAL_ERROR", + r"Worker process\(es\) died \(rank, exit code\): .*, -6\)", # the crash can also abort the worker +] -def pytest_configure(config): - config.addinivalue_line( - "markers", - "flaky_cuda_crash_on_blackwell: on Blackwell, report a known CUDA crash as xfail, other failures still fail", - ) - - -@pytest.hookimpl(wrapper=True) -def pytest_runtest_makereport(item, call): - """Report a known Blackwell CUDA crash in a marked test as xfail instead of a failure. - - Only for tests run in a worker pool: a crashed pool is replaced, so later tests are unaffected. - """ - report = yield - if ( - report.when == "call" - and report.failed - and item.get_closest_marker("flaky_cuda_crash_on_blackwell") - and call.excinfo is not None - # The message, not report.longreprtext: the traceback source lines can mention these errors - and any(err in str(call.excinfo.value) for err in _BLACKWELL_CUDA_CRASHES) - and is_blackwell() - ): - report.outcome = "skipped" - report.wasxfail = "known flaky CUDA crash on Blackwell (nemo:26.08 TE/CUDA 13 stack)" - return report +def pytest_collection_modifyitems(config, items): + if not is_blackwell(): + return + for item in items: + if item.path.is_relative_to(Path(__file__).parent): + item.add_marker(pytest.mark.flaky(reruns=1, only_rerun=_BLACKWELL_CUDA_CRASHES)) @pytest.fixture(scope="session", autouse=True) diff --git a/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py b/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py index c9f5df7459c..af8f38edb50 100644 --- a/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py +++ b/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py @@ -484,8 +484,6 @@ def _test_adapter_gradient_flow_freeze_lora_model(lora_config, tmp_path, rank, s LARGE_LORA_CFG_RANDOM_INIT_TEST, # Use random init so gradients flow to both lora_a and lora_b ], ) -# Intermittent CUBLAS_STATUS_INTERNAL_ERROR in backward on Blackwell; see conftest.py -@pytest.mark.flaky_cuda_crash_on_blackwell def test_adapter_gradient_flow_freeze_lora_model(dist_workers, lora_config, tmp_path): dist_workers.run( partial(_test_adapter_gradient_flow_freeze_lora_model, lora_config, str(tmp_path)) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 6d829276766..bb54faf40e0 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -119,9 +119,6 @@ SEED = 1234 -# Marked tests intermittently hit a CUDA illegal memory access on Blackwell (sm_120); see conftest.py. -flaky_on_blackwell = pytest.mark.flaky_cuda_crash_on_blackwell - def test_convert_megatron_parallel_linear(distributed_setup_size_1): initialize_for_megatron(seed=SEED) @@ -491,7 +488,6 @@ def _test_sharded_state_dict( ) @pytest.mark.parametrize("meta_device", [False, True]) @pytest.mark.parametrize("transformer_impl", ["local", "modelopt"]) -@flaky_on_blackwell def test_homogeneous_sharded_state_dict( dist_workers, tmp_path, config, meta_device, transformer_impl ): @@ -517,7 +513,6 @@ def test_homogeneous_sharded_state_dict( @pytest.mark.parametrize("transformer_impl", ["local", "modelopt"]) @pytest.mark.timeout(240) # Compressed state dict takes longer due to real quant conversion & saving/loading -@flaky_on_blackwell def test_homogeneous_compressed_sharded_state_dict( dist_workers, tmp_path, config, meta_device, transformer_impl ): @@ -599,7 +594,6 @@ def test_homogeneous_sharded_state_dict_hybrid(dist_workers, tmp_path, config): mixed_block_size_config, ], ) -@flaky_on_blackwell def test_heterogenous_sharded_state_dict(dist_workers, tmp_path, config): dist_workers.run( partial(_test_sharded_state_dict, tmp_path, config, 256, None, False, False, {}), @@ -689,7 +683,6 @@ def _test_fp8_real_quantize_helper(rank, size): assert real_quant_mem < cur_mem -@flaky_on_blackwell def test_fp8_real_quantize(dist_workers): dist_workers.run(_test_fp8_real_quantize_helper) @@ -1724,7 +1717,6 @@ def forward_loop(_model): assert all(m.input_quantizer.amax is not None for m in calibrated) -@flaky_on_blackwell @pytest.mark.parametrize("drive_wrapper", [False, True]) def test_layerwise_calibrate_vlm_text_forward_loop(dist_workers_size_1, drive_wrapper): dist_workers_size_1.run(partial(_test_layerwise_calibrate_vlm_text_forward_loop, drive_wrapper)) @@ -2304,7 +2296,6 @@ def _test_dsa_kv_cache_quant_helper(tmp_path, rank, size): assert all(m.k_bmm_quantizer.is_enabled and m.v_bmm_quantizer.is_enabled for m in restored) -@flaky_on_blackwell def test_dsa_kv_cache_quant(dist_workers_size_1, tmp_path): """DSAttention gets calibrated K/V KV-cache quantizers that survive a checkpoint round-trip.""" pytest.importorskip("megatron.core.transformer.experimental_attention_variant.dsa") From bc8383bb9704e7581046d6f1ad3dc221baa2bb29 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 28 Oct 2026 00:19:22 -0700 Subject: [PATCH 5/9] Rerun Blackwell CUDA crashes only for tests that use a worker pool A crash in the pytest process poisons the rerun, which runs in the same process, so only tests requesting a dist_workers* fixture get the rerun marker. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- tests/gpu_megatron/conftest.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index 802c474c1d5..5fe34b3c693 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -76,7 +76,10 @@ def pytest_collection_modifyitems(config, items): if not is_blackwell(): return for item in items: - if item.path.is_relative_to(Path(__file__).parent): + # Only pool tests: a crash in the pytest process would poison the rerun, which runs there too + if item.path.is_relative_to(Path(__file__).parent) and any( + name.startswith("dist_workers") for name in item.fixturenames + ): item.add_marker(pytest.mark.flaky(reruns=1, only_rerun=_BLACKWELL_CUDA_CRASHES)) From c7f3067a2a371ddef772a19f96b6d1a8d9bbdbe2 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 28 Oct 2026 00:39:00 -0700 Subject: [PATCH 6/9] Rerun transient HF/network failures and Blackwell crashes with one mechanism Every test now gets one rerun, 10 s later, when it fails with a known transient error (HF Hub/network/port collision, shared in _test_utils.transient_errors), or on Blackwell with the known cuBLAS crash in a gpu_megatron worker-pool test. run_example_command no longer retries internally: on a transient failure it raises TransientError, which the rerun matches. In-process steps are now matched on the exception chain's messages and the step's output instead of the formatted traceback, whose source lines could match a marker. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- tests/_test_utils/examples/run_command.py | 118 ++++++++-------------- tests/_test_utils/transient_errors.py | 52 ++++++++++ tests/conftest.py | 29 +++++- tests/gpu_megatron/conftest.py | 22 ---- tests/unit/test_example_run_command.py | 39 +++---- 5 files changed, 137 insertions(+), 123 deletions(-) create mode 100644 tests/_test_utils/transient_errors.py diff --git a/tests/_test_utils/examples/run_command.py b/tests/_test_utils/examples/run_command.py index a2e6487916c..f0f99e557f2 100644 --- a/tests/_test_utils/examples/run_command.py +++ b/tests/_test_utils/examples/run_command.py @@ -19,45 +19,14 @@ import signal import subprocess import threading -import time -import traceback import warnings from pathlib import Path from _test_utils.torch.distributed.utils import get_free_port +from _test_utils.transient_errors import TransientError, find_transient_marker MODELOPT_ROOT = Path(__file__).parents[3] -# Substrings that identify a *transient* failure a rerun clears: a HuggingFace Hub / dataset access -# error, or a rendezvous port that another process bound between its selection and its use -_TRANSIENT_MARKERS = ( - "HfHubHTTPError", - "Too Many Requests", - "500 Server Error", - "502 Bad Gateway", - "503 Server Error", - "504 Server Error", - "Bad Gateway", - "Service Unavailable", - "Gateway Time-out", - "ConnectionError", - "ReadTimeout", - "ConnectTimeout", - "Max retries exceeded", - "NewConnectionError", - "Connection reset by peer", - "Connection aborted", - "Consistency check failed", # partial / interrupted HF download - "couldn't connect to 'https://huggingface.co'", - "Couldn't reach", # datasets: "Couldn't reach on the Hub" - "Temporary failure in name resolution", # transient DNS - "Name or service not known", # transient DNS - "EADDRINUSE", # torch.distributed's TCPStore: "name: EADDRINUSE, message: address already in use" - "Address already in use", # the same from Python's OSError, gloo and NCCL -) -_HF_MAX_RETRIES = 1 -_HF_RETRY_DELAY_S = 10 - def extend_cmd_parts(cmd_parts: list[str], **kwargs): for key, value in kwargs.items(): @@ -75,7 +44,7 @@ def extend_cmd_parts(cmd_parts: list[str], **kwargs): def _run_capturing(cmd_parts: list[str], cwd: Path, env: dict[str, str]) -> tuple[int, str]: - """Run a command, streaming and capturing combined stdout/stderr to catch transient HF errors. + """Run a command, streaming and capturing combined stdout/stderr to catch transient errors. Drained by a reader thread rather than ``subprocess.run``, which waits for EOF: a descendant outliving the command keeps the pipe open, hanging the read and discarding every log line. @@ -151,15 +120,23 @@ def set_in_process_runner(runner) -> None: _in_process_runner = runner +def _exception_chain_text(exc: BaseException | None) -> str: + """``Type: message`` of ``exc`` and of each exception it was raised from or while handling.""" + lines, seen = [], set() + while exc is not None and id(exc) not in seen: + seen.add(id(exc)) + lines.append(f"{type(exc).__name__}: {exc}") + exc = exc.__cause__ or (None if exc.__suppress_context__ else exc.__context__) + return "\n".join(lines) + + def run_example_command( cmd_parts: list[str], example_path: str, setup_free_port: bool = False, env: dict[str, str] | None = None, - hf_max_retries: int = _HF_MAX_RETRIES, - hf_retry_delay_s: int = _HF_RETRY_DELAY_S, -) -> str | None: - """Run an example command, retrying transient errors (see ``_TRANSIENT_MARKERS``) once.""" +) -> str: + """Run an example command, raising ``TransientError`` on a transient failure so the test reruns.""" print(f"[{example_path}] Running command: {cmd_parts}") in_process = _in_process_runner if env is None else None if _in_process_runner is not None and env is not None: @@ -169,45 +146,34 @@ def run_example_command( env = os.environ.copy() if env is None else env cwd = MODELOPT_ROOT / "examples" / example_path - for attempt in range(hf_max_retries + 1): - if setup_free_port: - # Subprocess steps only: an in-process runner picks its own free port, since env is a - # copy it never sees. - env["MASTER_PORT"] = str(get_free_port()) # fresh port per attempt - if in_process is not None: - # Inside the loop so in-process steps get the same transient-error retries; these - # tests do hit the Hub (e.g. calib_dataset_name="cnn_dailymail"). - try: - result = in_process(cmd_parts, example_path) - except Exception as e: - # Re-raise unless it looks transient, so a real failure keeps its traceback - # instead of being flattened into CalledProcessError. - # format_exc() walks __cause__/__context__: most markers are exception type - # names that only ever appear in a formatted traceback, never in str(e) -- a Hub - # ConnectionError typically surfaces wrapped in a DatasetGenerationError. The - # runner also attaches what the step printed, for launcher-level failures. - captured = getattr(e, "captured_output", "") - text = f"{traceback.format_exc()}\n{captured}" - if attempt == hf_max_retries or not any( - marker in text for marker in _TRANSIENT_MARKERS - ): - raise - returncode, output = 1, text - else: - assert result is not None, "an in-process runner must return output, not fall back" - return result - else: - returncode, output = _run_capturing(cmd_parts, cwd, env) - if returncode == 0: - return output - transient = any(marker in output for marker in _TRANSIENT_MARKERS) - if not transient or attempt == hf_max_retries: - raise subprocess.CalledProcessError(returncode, cmd_parts, output=output) - warnings.warn( - f"[{example_path}] transient error; retrying in " - f"{hf_retry_delay_s}s (attempt {attempt + 1}/{hf_max_retries})" - ) - time.sleep(hf_retry_delay_s) + if in_process is not None: + try: + result = in_process(cmd_parts, example_path) + except Exception as e: + # Match the messages of the exception chain (a Hub ConnectionError typically surfaces + # wrapped in a DatasetGenerationError) and what the step printed, for launcher-level + # failures. Not a formatted traceback: its source lines can contain a marker. + marker = find_transient_marker( + f"{_exception_chain_text(e)}\n{getattr(e, 'captured_output', '')}" + ) + if marker is None: + raise # a real failure keeps its own traceback + raise TransientError(f"[{example_path}] transient error: {marker!r}") from e + assert result is not None, "an in-process runner must return output, not fall back" + return result + + if setup_free_port: + # Subprocess steps only: an in-process runner picks its own free port, since env is a + # copy it never sees. + env["MASTER_PORT"] = str(get_free_port()) + returncode, output = _run_capturing(cmd_parts, cwd, env) + if returncode == 0: + return output + error = subprocess.CalledProcessError(returncode, cmd_parts, output=output) + marker = find_transient_marker(output) + if marker is None: + raise error + raise TransientError(f"[{example_path}] transient error: {marker!r}") from error def run_hf_ptq_command(*, model: str, quant: str | None = None, vlm: bool = False, **kwargs): diff --git a/tests/_test_utils/transient_errors.py b/tests/_test_utils/transient_errors.py new file mode 100644 index 00000000000..146a3e536dd --- /dev/null +++ b/tests/_test_utils/transient_errors.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Failures that a rerun clears. ``tests/conftest.py`` reruns a test once when it fails with one.""" + +# Substrings that identify a *transient* failure: a HuggingFace Hub / dataset access error, or a +# rendezvous port that another process bound between its selection and its use +TRANSIENT_MARKERS = ( + "HfHubHTTPError", + "Too Many Requests", + "500 Server Error", + "502 Bad Gateway", + "503 Server Error", + "504 Server Error", + "Bad Gateway", + "Service Unavailable", + "Gateway Time-out", + "ConnectionError", + "ReadTimeout", + "ConnectTimeout", + "Max retries exceeded", + "NewConnectionError", + "Connection reset by peer", + "Connection aborted", + "Consistency check failed", # partial / interrupted HF download + "couldn't connect to 'https://huggingface.co'", + "Couldn't reach", # datasets: "Couldn't reach on the Hub" + "Temporary failure in name resolution", # transient DNS + "Name or service not known", # transient DNS + "EADDRINUSE", # torch.distributed's TCPStore: "name: EADDRINUSE, message: address already in use" + "Address already in use", # the same from Python's OSError, gloo and NCCL +) + + +class TransientError(RuntimeError): + """A command failed and its output shows one of ``TRANSIENT_MARKERS``.""" + + +def find_transient_marker(text: str) -> str | None: + """Return the first of ``TRANSIENT_MARKERS`` found in ``text``, if any.""" + return next((marker for marker in TRANSIENT_MARKERS if marker in text), None) diff --git a/tests/conftest.py b/tests/conftest.py index 7f606326b8d..7b935cdcc22 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,6 +15,7 @@ import os import platform +import re import sys from pathlib import Path @@ -24,6 +25,8 @@ from _test_utils.fs_utils import assert_unmodified_tree from _test_utils.mlflow import FakeMlflow, pin_tracking_env from _test_utils.torch.distributed.utils import init_process +from _test_utils.torch.misc import is_blackwell +from _test_utils.transient_errors import TRANSIENT_MARKERS, TransientError import modelopt.torch.opt as mto @@ -70,9 +73,18 @@ def pytest_addoption(parser): "unit": 120, } +# On Blackwell (sm_120), nemo:26.08's cuBLAS intermittently fails a GEMM once a gpu_megatron worker +# process has run about a dozen model-building jobs. A failed job replaces the worker, so the rerun +# gets a fresh one. +_BLACKWELL_CUDA_CRASHES = [ + "illegal memory access", + "CUBLAS_STATUS_INTERNAL_ERROR", + r"Worker process\(es\) died \(rank, exit code\): .*, -6\)", # the crash can also abort the worker +] + def pytest_collection_modifyitems(config, items): - """Skip flag-gated tests and apply a default per-test timeout based on the test directory.""" + """Skip flag-gated tests, apply a default per-test timeout, and rerun transient failures once.""" skip_marks = [ ("manual", "--run-manual"), ("nightly", "--run-nightly"), @@ -103,6 +115,21 @@ def pytest_collection_modifyitems(config, items): ) item.add_marker(pytest.mark.timeout(_DEFAULT_TIMEOUT[group])) + if not config.pluginmanager.hasplugin("rerunfailures"): + return + gpu_megatron_root = tests_root / "gpu_megatron" + for item in items: + only_rerun = [TransientError, *map(re.escape, TRANSIENT_MARKERS)] + # Only worker-pool tests: a crash in the pytest process would poison the rerun, which runs + # there too. is_blackwell() last, so other test groups do not initialize CUDA here. + if ( + item.path.is_relative_to(gpu_megatron_root) + and any(name.startswith("dist_workers") for name in item.fixturenames) + and is_blackwell() + ): + only_rerun += _BLACKWELL_CUDA_CRASHES + item.add_marker(pytest.mark.flaky(reruns=1, reruns_delay=10, only_rerun=only_rerun)) + # General Fixtures ################################################################################# @pytest.fixture diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index 5fe34b3c693..1af63e694e5 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -15,7 +15,6 @@ import contextlib import gc import queue -from pathlib import Path import pytest import torch @@ -25,7 +24,6 @@ get_live_pool, reset_worker_state, ) -from _test_utils.torch.misc import is_blackwell from _test_utils.torch.transformers_models import get_tiny_tokenizer from megatron.core.parallel_state import destroy_model_parallel @@ -63,26 +61,6 @@ def _prebuild_quant_cuda_extensions(): ext.precompile() -# On Blackwell (sm_120), nemo:26.08's cuBLAS intermittently fails a GEMM once a worker process has -# run about a dozen model-building jobs. A failed job replaces the worker, so rerun it once on a fresh one. -_BLACKWELL_CUDA_CRASHES = [ - "illegal memory access", - "CUBLAS_STATUS_INTERNAL_ERROR", - r"Worker process\(es\) died \(rank, exit code\): .*, -6\)", # the crash can also abort the worker -] - - -def pytest_collection_modifyitems(config, items): - if not is_blackwell(): - return - for item in items: - # Only pool tests: a crash in the pytest process would poison the rerun, which runs there too - if item.path.is_relative_to(Path(__file__).parent) and any( - name.startswith("dist_workers") for name in item.fixturenames - ): - item.add_marker(pytest.mark.flaky(reruns=1, only_rerun=_BLACKWELL_CUDA_CRASHES)) - - @pytest.fixture(scope="session", autouse=True) def _prestart_preprocess_forkserver(): """Start the forkserver that ``megatron_preprocess_data`` workers fork from, at session start. diff --git a/tests/unit/test_example_run_command.py b/tests/unit/test_example_run_command.py index e04d61b5c09..bdb95a45eb2 100644 --- a/tests/unit/test_example_run_command.py +++ b/tests/unit/test_example_run_command.py @@ -21,37 +21,28 @@ import pytest from _test_utils.examples import run_command +from _test_utils.transient_errors import TransientError @pytest.mark.parametrize( - ("error", "attempts"), + ("error", "transient"), [ - ("HfHubHTTPError: 503 Server Error: Service Unavailable", 2), - ("DistNetworkError: port: 38789, name: EADDRINUSE, message: address already in use", 2), - ("AssertionError: the quantized model diverged", 1), + ("HfHubHTTPError: 503 Server Error: Service Unavailable", True), + ("DistNetworkError: port: 38789, name: EADDRINUSE, message: address already in use", True), + ("AssertionError: the quantized model diverged", False), ], ) -def test_run_example_command_retries_only_transient_errors(monkeypatch, error, attempts): - calls = [] +def test_run_example_command_flags_only_transient_errors(monkeypatch, error, transient): + def failing_step(cmd_parts, example_path): + raise RuntimeError(error) # a marker on a source line must not count: ConnectionError - def step_failing_once(cmd_parts, example_path): - calls.append(cmd_parts) - if len(calls) == 1: - raise RuntimeError(error) - return "output" - - def run(): - return run_command.run_example_command( - ["torchrun", "step.py"], "hf_ptq", hf_retry_delay_s=0 - ) - - monkeypatch.setattr(run_command, "_in_process_runner", step_failing_once) - if attempts == 1: - with pytest.raises(RuntimeError, match="diverged"): - run() - else: - assert run() == "output" - assert len(calls) == attempts + monkeypatch.setattr(run_command, "_in_process_runner", failing_step) + with pytest.raises(RuntimeError) as excinfo: + run_command.run_example_command(["torchrun", "step.py"], "hf_ptq") + # TransientError is what the rerun marker in tests/conftest.py matches + assert isinstance(excinfo.value, TransientError) == transient + original = excinfo.value.__cause__ if transient else excinfo.value + assert str(original) == error def test_run_capturing_does_not_block_on_a_survivor_holding_the_pipe( From 1af7bf62bf75ca4a844904d07017e980c0628907 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 28 Oct 2026 00:44:30 -0700 Subject: [PATCH 7/9] Show the error behind each rerun in the pytest summary A rerun attempt is not reported as a failure, so neither -ra nor --instafail printed why a test was rerun. --rerun-show-tracebacks adds a "rerun test summary info" section with each rerun attempt's traceback. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index ca75164ffb5..6233c3d3c38 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -358,7 +358,7 @@ skips = [ # Default additional options # Show a short test summary info for all except passed tests with -ra flag # print execution time for 50 slowest tests and generate coverage reports -addopts = "-v -ra --instafail --cov-report=term-missing --cov-report=html --cov-report=xml:coverage.xml --cov-config=pyproject.toml --durations=50 --strict-markers" +addopts = "-v -ra --instafail --rerun-show-tracebacks --cov-report=term-missing --cov-report=html --cov-report=xml:coverage.xml --cov-config=pyproject.toml --durations=50 --strict-markers" pythonpath = ["tests/"] # Apply per-test timeouts (see tests/conftest.py) to the test call only, not fixture setup/teardown timeout_func_only = true From ce71e35bf52789206e08839e55f8b739baab7f55 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 28 Oct 2026 00:54:44 -0700 Subject: [PATCH 8/9] Tighten the worker-abort rerun pattern and drop the unused plugin guard - Rerun a worker death only when every dead rank exited with -6 (SIGABRT), so an abort alongside an OOM kill (-9) on another rank still fails. - Remove the hasplugin("rerunfailures") guard: --rerun-show-tracebacks in addopts already makes the plugin required. - Explain the pytest-rerunfailures pin and why its MPL-2.0 license is fine. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- pyproject.toml | 4 +++- tests/conftest.py | 5 ++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6233c3d3c38..049aeba0541 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,7 +122,9 @@ dev-test = [ "pytest~=9.1.0", "pytest-cov~=7.1.0", "pytest-instafail==0.5.0", - "pytest-rerunfailures~=16.7", # rerun known Blackwell cuBLAS crashes (tests/gpu_megatron) + # MPL-2.0 (weak copyleft) is fine: test-only and used unmodified, and its copyleft covers only + # its own files, not code that uses it (tqdm, a core dependency, is MPL-2.0 too) + "pytest-rerunfailures~=16.7", # transient-failure reruns (tests/conftest.py); 16.7 matches the exception chain "pytest-timeout~=2.4.0", "pytest-xdist~=3.8.0", "uv", diff --git a/tests/conftest.py b/tests/conftest.py index 7b935cdcc22..5957ccf4b7e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -79,7 +79,8 @@ def pytest_addoption(parser): _BLACKWELL_CUDA_CRASHES = [ "illegal memory access", "CUBLAS_STATUS_INTERNAL_ERROR", - r"Worker process\(es\) died \(rank, exit code\): .*, -6\)", # the crash can also abort the worker + # The crash can also abort the worker (SIGABRT); not when another rank died otherwise, e.g. OOM + r"Worker process\(es\) died \(rank, exit code\): \[\(\d+, -6\)(, \(\d+, -6\))*\]", ] @@ -115,8 +116,6 @@ def pytest_collection_modifyitems(config, items): ) item.add_marker(pytest.mark.timeout(_DEFAULT_TIMEOUT[group])) - if not config.pluginmanager.hasplugin("rerunfailures"): - return gpu_megatron_root = tests_root / "gpu_megatron" for item in items: only_rerun = [TransientError, *map(re.escape, TRANSIENT_MARKERS)] From 9617ed510e27f827df6c553ddb5a7da8fa1894f0 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 28 Oct 2026 10:09:01 -0700 Subject: [PATCH 9/9] Add a TODO to remove the Blackwell crash rerun after the nemo container update Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- tests/conftest.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/conftest.py b/tests/conftest.py index 5957ccf4b7e..6b15deda80b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -76,6 +76,7 @@ def pytest_addoption(parser): # On Blackwell (sm_120), nemo:26.08's cuBLAS intermittently fails a GEMM once a gpu_megatron worker # process has run about a dozen model-building jobs. A failed job replaces the worker, so the rerun # gets a fresh one. +# TODO: Remove after the CI nemo container is updated (re-check on nemo:26.10). _BLACKWELL_CUDA_CRASHES = [ "illegal memory access", "CUBLAS_STATUS_INTERNAL_ERROR",