Fix ONNX export of channels_last_3d pre-decoder models; fail fast on requested export failures

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.

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":

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":

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):

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):

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))

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