From a7ba11c0c4e48c1824969bc0e5593bbd7745405b Mon Sep 17 00:00:00 2001 From: Ivan Basov Date: Tue, 21 Jul 2026 19:37:29 +0000 Subject: [PATCH] Fix ONNX export of channels_last_3d models; fail fast on requested export failure The eval-input layout match introduced with the AMP training work calls contiguous(memory_format=channels_last_3d) inside PreDecoderMemoryEvalModule.forward. The surface inference path converts every model to channels_last_3d before export, so the traced graph contains a memory-format op the legacy exporter cannot lower and torch.onnx.export(..., dynamo=False) fails with "onnx memory_format support is not implemented". Skip the layout match during export: ONNX has no memory-format concept, so the artifact is value-identical. Also stop silently downgrading to PyTorch when ONNX_WORKFLOW=1/2 was explicitly requested and the export fails: broadcast the failure to all ranks (so multi-GPU runs exit promptly instead of blocking in the next collective) and raise, making local_run.sh exit nonzero. INT8 quantization failure in the ablation path now falls back to the FP32 ONNX like the LER path instead of aborting the run; FP8 stays fail-fast. Add onnx to the public inference requirements (the legacy exporter needs it to serialize the requested artifact) and add a CI-collected regression test that exports a channels_last_3d wrapper and checks the ONNX output against eager. Co-Authored-By: Claude Fable 5 --- README.md | 3 ++ code/evaluation/failure_analysis.py | 68 +++++++++++++++----------- code/evaluation/logical_error_rate.py | 33 +++++++++++-- code/requirements_public_inference.txt | 3 ++ code/tests/test_failure_analysis.py | 48 ++++++++++++++---- code/tests/test_precision.py | 47 ++++++++++++++++++ code/training/precision.py | 6 +++ 7 files changed, 164 insertions(+), 44 deletions(-) diff --git a/README.md b/README.md index 54af744..ce57a6d 100644 --- a/README.md +++ b/README.md @@ -362,6 +362,9 @@ ONNX_WORKFLOW=3 WORKFLOW=inference bash code/scripts/local_run.sh Notes: - TensorRT workflows (`ONNX_WORKFLOW=2` or `3`) require `tensorrt` and `modelopt`. +- A failed ONNX export (`ONNX_WORKFLOW=1` or `2`) is fatal (nonzero exit) instead of silently + falling back to PyTorch. A TensorRT build/load failure after a successful export still falls + back to PyTorch. - FP8 quantization failure is fatal. INT8 failure falls back to the FP32 ONNX model silently. - ONNX and engine files are written to the current working directory. - `ONNX_WORKFLOW` is also honoured by the `decoder_ablation` workflow — see below. diff --git a/code/evaluation/failure_analysis.py b/code/evaluation/failure_analysis.py index 56e476a..0c16e4e 100644 --- a/code/evaluation/failure_analysis.py +++ b/code/evaluation/failure_analysis.py @@ -29,6 +29,7 @@ _build_stab_maps, _decode_batch, _parse_quant_format, + _sync_and_raise_on_export_failure, map_grid_to_stabilizer_tensor, sample_predictions, ) @@ -750,6 +751,7 @@ def _setup_trt_for_ablation(model, cfg, dist, device, basis, D, half, stim_dets) ) elif onnx_workflow in (OnnxWorkflow.EXPORT_ONNX_ONLY, OnnxWorkflow.EXPORT_AND_USE_TRT): + export_error = None if dist.rank == 0: try: fp32_onnx_path = ( @@ -781,39 +783,47 @@ def _setup_trt_for_ablation(model, cfg, dist, device, basis, D, half, stim_dets) print(f"[Ablation] Exported FP32 ONNX: {fp32_onnx_path}") if quant_format: - calib_samples = int(os.environ.get("QUANT_CALIB_SAMPLES", "256")) - calib_dets = stim_dets[:calib_samples].astype(np.uint8) + # Quantization failure semantics mirror the LER path: FP8 is + # fail-fast, INT8 falls back to the FP32 ONNX (README notes). try: - import modelopt.onnx.quantization as mq - quant_kwargs = {} - if quant_format == "fp8": - quant_kwargs["op_types_to_quantize"] = ["Conv"] - quant_kwargs["high_precision_dtype"] = "fp16" - mq.quantize( - onnx_path=fp32_onnx_path, - quantize_mode=quant_format, - calibration_data={"dets": calib_dets.astype("float32")}, - output_path=onnx_path, - **quant_kwargs, - ) - except ImportError: + calib_samples = int(os.environ.get("QUANT_CALIB_SAMPLES", "256")) + calib_dets = stim_dets[:calib_samples].astype(np.uint8) + try: + import modelopt.onnx.quantization as mq + quant_kwargs = {} + if quant_format == "fp8": + quant_kwargs["op_types_to_quantize"] = ["Conv"] + quant_kwargs["high_precision_dtype"] = "fp16" + mq.quantize( + onnx_path=fp32_onnx_path, + quantize_mode=quant_format, + calibration_data={"dets": calib_dets.astype("float32")}, + output_path=onnx_path, + **quant_kwargs, + ) + except ImportError: + if quant_format == "fp8": + raise RuntimeError( + "[Ablation] FP8 quantization requires nvidia-modelopt." + ) + from evaluation.logical_error_rate import _ort_quantize_int8 + _ort_quantize_int8(fp32_onnx_path, onnx_path, calib_dets) + print(f"[Ablation] Exported quantized ONNX: {onnx_path}") + except Exception as e: if quant_format == "fp8": raise RuntimeError( - "[Ablation] FP8 quantization requires nvidia-modelopt." - ) - from evaluation.logical_error_rate import _ort_quantize_int8 - _ort_quantize_int8(fp32_onnx_path, onnx_path, calib_dets) - print(f"[Ablation] Exported quantized ONNX: {onnx_path}") + f"[Ablation] FP8 ONNX quantization failed (fail-fast): {e}" + ) from e + print(f"[Ablation] ONNX quantization failed: {e}; using FP32 ONNX.") + onnx_path = fp32_onnx_path except Exception as e: - print(f"[Ablation] ONNX export failed: {e}; using PyTorch.") - onnx_workflow = OnnxWorkflow.TORCH_ONLY - - if dist.world_size > 1: - # Broadcast rank 0's onnx_workflow (may have been set to TORCH_ONLY on - # export failure) so non-zero ranks skip the TRT build when rank 0 failed. - wf_list = [onnx_workflow] - torch.distributed.broadcast_object_list(wf_list, src=0) - onnx_workflow = wf_list[0] + # Stash instead of raising here: non-zero ranks must first learn + # about the failure through the collective below, or they would + # block in it after rank 0 dies. + export_error = e + # Doubles as the post-export sync so non-zero ranks don't race ahead to + # the TRT build. + _sync_and_raise_on_export_failure(export_error, onnx_workflow, dist, "[Ablation]") engine_path = onnx_path.replace(".onnx", ".engine") if onnx_workflow == OnnxWorkflow.EXPORT_AND_USE_TRT and device.type == "cuda": diff --git a/code/evaluation/logical_error_rate.py b/code/evaluation/logical_error_rate.py index 99a9359..b39fbcf 100644 --- a/code/evaluation/logical_error_rate.py +++ b/code/evaluation/logical_error_rate.py @@ -254,6 +254,28 @@ def rewind(self): ) +def _sync_and_raise_on_export_failure(export_error, onnx_workflow, dist, tag: str) -> None: + """Fail fast (on every rank) when a requested ONNX export failed on rank 0. + + ONNX_WORKFLOW defaults to 0, so a failed export here was explicitly + requested: raising instead of silently falling back to PyTorch makes the + run exit nonzero, so automation cannot mistake a fallback run for a + produced artifact. The broadcast doubles as the post-export barrier and + lets non-zero ranks exit promptly instead of blocking in a later + collective once rank 0 dies. + """ + if dist.world_size > 1: + msg_list = [None if export_error is None else str(export_error)] + torch.distributed.broadcast_object_list(msg_list, src=0) + if export_error is None and msg_list[0] is not None: + export_error = RuntimeError(msg_list[0]) + if export_error is not None: + raise RuntimeError( + f"{tag} ONNX export failed with ONNX_WORKFLOW={onnx_workflow.value} " + f"({onnx_workflow.name}) explicitly requested: {export_error}" + ) from export_error + + def _time_single_shot_latency_stim( matcher, baseline_syndromes: np.ndarray, @@ -1168,6 +1190,7 @@ def run_inference_and_decode_pre_decoder_memory(model, device, dist, cfg) -> dic ) elif onnx_workflow in (OnnxWorkflow.EXPORT_ONNX_ONLY, OnnxWorkflow.EXPORT_AND_USE_TRT): + export_error = None if dist.rank == 0: try: example_dets = torch.randint(0, 2, example_shape, dtype=torch.uint8, device=device) @@ -1250,11 +1273,11 @@ def run_inference_and_decode_pre_decoder_memory(model, device, dist, cfg) -> dic print(f"[LER] ONNX quantization failed: {e}; using FP32 ONNX.") onnx_path = fp32_onnx_path except Exception as e: - if dist.rank == 0: - print(f"[LER] ONNX export failed: {e}; falling back to PyTorch.") - onnx_workflow = OnnxWorkflow.TORCH_ONLY - if dist.world_size > 1: - torch.distributed.barrier() + # Stash instead of raising here: non-zero ranks must first learn + # about the failure through the collective below, or they would + # block in it after rank 0 dies. + export_error = e + _sync_and_raise_on_export_failure(export_error, onnx_workflow, dist, "[LER]") # Re-derive engine_path from the final onnx_path (may have changed on quant fallback) engine_path = str(Path(onnx_path).with_suffix(".engine")) if onnx_workflow == OnnxWorkflow.EXPORT_AND_USE_TRT and device.type == "cuda": diff --git a/code/requirements_public_inference.txt b/code/requirements_public_inference.txt index 13989f9..0b64721 100644 --- a/code/requirements_public_inference.txt +++ b/code/requirements_public_inference.txt @@ -27,6 +27,9 @@ safetensors>=0.4.0 scipy ldpc beliefmatching +# Required by the legacy TorchScript ONNX exporter at serialization time +# (ONNX_WORKFLOW=1/2 cannot produce an artifact without it). +onnx # Color-code support (Torch + cuStabilizer runtime). chromobius # Optional GPU-only prerequisites (not pip-installed here due to size and CUDA dependency): diff --git a/code/tests/test_failure_analysis.py b/code/tests/test_failure_analysis.py index a29e6d5..fcb40e6 100644 --- a/code/tests/test_failure_analysis.py +++ b/code/tests/test_failure_analysis.py @@ -1161,6 +1161,32 @@ def test_invalid_value_raises_valueerror(self): OnnxWorkflow(99) +class TestSyncAndRaiseOnExportFailure(unittest.TestCase): + """A failed explicitly-requested export must fail fast on every rank. + + Covers the shared helper used by both the LER inference path + (run_inference_and_decode_pre_decoder_memory) and decoder_ablation_study. + """ + + def test_no_error_is_a_noop(self): + from evaluation.logical_error_rate import (OnnxWorkflow, _sync_and_raise_on_export_failure) + _sync_and_raise_on_export_failure( + None, OnnxWorkflow.EXPORT_ONNX_ONLY, _DummyDist(), "[LER]" + ) + + def test_export_error_raises_with_context(self): + from evaluation.logical_error_rate import (OnnxWorkflow, _sync_and_raise_on_export_failure) + cause = ValueError("onnx memory_format support is not implemented") + with self.assertRaisesRegex( + RuntimeError, r"\[LER\] ONNX export failed with ONNX_WORKFLOW=1 " + r"\(EXPORT_ONNX_ONLY\) explicitly requested" + ) as ctx: + _sync_and_raise_on_export_failure( + cause, OnnxWorkflow.EXPORT_ONNX_ONLY, _DummyDist(), "[LER]" + ) + self.assertIs(ctx.exception.__cause__, cause) + + class TestDecoderAblationStudyTRTFallback(unittest.TestCase): """ ONNX_WORKFLOW=3 with a missing engine file must fall back to PyTorch silently @@ -1220,8 +1246,9 @@ def test_missing_engine_sample_count_correct(self): class TestDecoderAblationStudyOnnxExport(unittest.TestCase): """ - ONNX_WORKFLOW=1 must attempt ONNX export (rank 0) then fall back to PyTorch for inference. - Results must be identical in structure to the default PyTorch path. + ONNX_WORKFLOW=1 must attempt ONNX export (rank 0) and run inference with + PyTorch; results must be identical in structure to the default PyTorch + path. A failed export must raise instead of silently falling back. """ _D = 3 @@ -1265,9 +1292,13 @@ def _fake_onnx_export(module, *args, **kwargs): self.assertEqual(result["total_samples"], self._N) self.assertTrue(set(DECODER_NAMES).issubset(set(result["decoder_errors"].keys()))) - def test_workflow1_export_failure_falls_back_gracefully(self): - """If ONNX export raises, results must still be valid (PyTorch fallback).""" - from evaluation.failure_analysis import decoder_ablation_study, DECODER_NAMES + def test_workflow1_export_failure_raises(self): + """If an explicitly requested ONNX export raises, the run must fail fast. + + Silently falling back to PyTorch (and exiting 0) would let automation + believe the requested ONNX artifact was generated when it was not. + """ + from evaluation.failure_analysis import decoder_ablation_study from data.datapipe_stim import QCDataPipePreDecoder_Memory_inference real_ds = QCDataPipePreDecoder_Memory_inference( distance=self._D, @@ -1288,11 +1319,8 @@ def test_workflow1_export_failure_falls_back_gracefully(self): patch("torch.onnx.export", side_effect=RuntimeError("export broken")), \ patch("os.getcwd", return_value=tmpdir): mf.create_datapipe_inference.return_value = real_ds - result = decoder_ablation_study( - _ZeroModel(), torch.device("cpu"), _DummyDist(), cfg - ) - self.assertEqual(result["total_samples"], self._N) - self.assertTrue(set(DECODER_NAMES).issubset(set(result["decoder_errors"].keys()))) + with self.assertRaisesRegex(RuntimeError, "ONNX export failed"): + decoder_ablation_study(_ZeroModel(), torch.device("cpu"), _DummyDist(), cfg) class TestDecoderAblationStudyTRTExecution(unittest.TestCase): diff --git a/code/tests/test_precision.py b/code/tests/test_precision.py index 9d9f7bc..4f67fcc 100644 --- a/code/tests/test_precision.py +++ b/code/tests/test_precision.py @@ -10,6 +10,8 @@ from __future__ import annotations import sys +import tempfile +import unittest from pathlib import Path import pytest @@ -110,3 +112,48 @@ def test_match_input_noop_for_contiguous_model(): conv = torch.nn.Conv3d(4, 8, kernel_size=3) # default contiguous x = torch.randn(1, 4, 3, 5, 5) assert match_input_to_model_memory_format(x, conv) is x + + +# unittest.TestCase so CI's `unittest discover` collects it (the pytest-style +# tests above are only run by pytest). +class TestMatchInputDuringOnnxExport(unittest.TestCase): + """The layout-match helper must be a no-op inside torch.onnx.export. + + Regression test: an unguarded contiguous(memory_format=channels_last_3d) + in a traced forward makes the legacy exporter fail with + "onnx memory_format support is not implemented". + """ + + def test_match_input_skipped_during_onnx_export(self): + # The legacy TorchScript exporter needs the onnx package to serialize. + try: + from onnx.reference import ReferenceEvaluator + except ImportError: + self.skipTest("onnx is not installed") + + class Wrap(torch.nn.Module): + + def __init__(self): + super().__init__() + self.model = module_to_channels_last_3d(torch.nn.Conv3d(4, 8, kernel_size=3), True) + + def forward(self, x): + x = match_input_to_model_memory_format(x, self.model) + return self.model(x) + + wrap = Wrap().eval() + x = torch.randn(1, 4, 3, 8, 8) + with tempfile.TemporaryDirectory() as tmpdir: + out_path = Path(tmpdir) / "wrap.onnx" + # Without the in-export guard the legacy exporter raises + # SymbolicValueError: "onnx memory_format support is not implemented". + torch.onnx.export( + wrap, (x,), str(out_path), opset_version=18, input_names=["x"], dynamo=False + ) + self.assertTrue(out_path.exists()) + # Skipping the layout conversion is value-neutral: the exported + # graph must reproduce the eager output. + (onnx_out,) = ReferenceEvaluator(str(out_path)).run(None, {"x": x.numpy()}) + with torch.no_grad(): + eager = wrap(x) + self.assertTrue(torch.allclose(eager, torch.from_numpy(onnx_out), atol=1e-5)) diff --git a/code/training/precision.py b/code/training/precision.py index 28fb8d0..bcc048c 100644 --- a/code/training/precision.py +++ b/code/training/precision.py @@ -119,7 +119,13 @@ def match_input_to_model_memory_format(tensor: torch.Tensor, model) -> torch.Ten If the model runs in channels_last_3d, a contiguous half-precision input forces the slow Conv3D kernel; converting the input keeps eval on the fast Tensor-Core path and consistent with training. No-op for contiguous models. + + Skipped during ONNX export: the legacy exporter cannot lower + ``contiguous(memory_format=channels_last_3d)``, and ONNX has no + memory-format concept — layout only affects kernel dispatch, not values. """ + if getattr(torch.onnx, "is_in_onnx_export", lambda: False)(): + return tensor if tensor.dim() == 5 and model_is_channels_last_3d(model): return tensor.contiguous(memory_format=torch.channels_last_3d) return tensor