[6403893][ONNX][AutoCast] Fix false-success TRT parse for large ModelProto

modelopt/onnx/trt_utils.py

@@ -15,8 +15,11 @@
 
 """This module contains TensorRT utils."""
 
+import copy
 import ctypes
+import os
 import platform
+import tempfile
 
 import lief
 import onnx
@@ -84,6 +87,28 @@ def _load_trt_plugin(plugin_path: str, registry) -> None:
             ctypes.CDLL(plugin_path)
 
 
+def _requires_file_backed_parse(model: onnx.ModelProto) -> bool:
+    """Return True if the model must be parsed from a file rather than in-memory bytes.
+
+    TensorRT's in-memory ``OnnxParser.parse(bytes)`` can silently return an empty network
+    (0 layers/0 tensors) for models at or above the protobuf 2 GiB limit, even though
+    ``SerializeToString()`` succeeds. In that regime we must serialize to an external-data
+    file and use ``parse_from_file()`` instead. If the size cannot even be computed
+    (e.g. ``ByteSize()`` overflows), we conservatively route through the file path as well.
+
+    Args:
+        model: In-memory ONNX model.
+
+    Returns:
+        True if the model should be parsed via a temporary external-data file.
+    """
+    try:
+        return model.ByteSize() >= onnx.checker.MAXIMUM_PROTOBUF
+    except Exception as e:
+        logger.debug(f"Could not compute model ByteSize ({e}); using file-backed TRT parsing.")
+        return True
+
+
 def get_custom_layers(
     onnx_path: str | onnx.ModelProto,
     trt_plugins: list[str] | None,
@@ -123,11 +148,33 @@ def get_custom_layers(
     )
     logger.debug("Created TensorRT builder and network")
 
-    # Parse ONNX file
+    # Parse ONNX file.
+    #
+    # For an in-memory ModelProto at or above the protobuf 2 GiB limit, TensorRT's
+    # `parser.parse(bytes)` can silently build an empty network (0 layers/0 tensors) even
+    # though `SerializeToString()` succeeds. Route such models through a temporary
+    # external-data file and `parse_from_file()` (the same branch used for string paths, and
+    # matching ModelOpt's file-backed policy in `save_onnx()`/`load_onnx_model()`).
     parser = trt.OnnxParser(network, trt_logger)
-    parser_func = parser.parse_from_file if isinstance(onnx_path, str) else parser.parse
-    onnx_path = onnx_path if isinstance(onnx_path, str) else onnx_path.SerializeToString()
-    if not parser_func(onnx_path):
+    if isinstance(onnx_path, str):
+        parse_ok = parser.parse_from_file(onnx_path)
+    elif _requires_file_backed_parse(onnx_path):
+        with tempfile.TemporaryDirectory(prefix="modelopt_trt_") as tmpdir:
+            model_path = os.path.join(tmpdir, "model.onnx")
+            # Deep-copy so externalization doesn't strip weights from the caller's model.
+            model_copy = copy.deepcopy(onnx_path)
+            onnx.save(
+                model_copy,
+                model_path,
+                save_as_external_data=True,
+                all_tensors_to_one_file=True,
+                location="model.data",
+                size_threshold=0,
+            )
+            parse_ok = parser.parse_from_file(model_path)
+    else:
+        parse_ok = parser.parse(onnx_path.SerializeToString())
+    if not parse_ok:
         error_str = [str(parser.get_error(error)) for error in range(parser.num_errors)]
         raise Exception(f"Failed to parse ONNX file: {''.join(error_str)}")
 

tests/gpu/onnx/quantization/test_plugin.py

@@ -22,10 +22,11 @@
 from _test_utils.onnx.autocast.utils import _assert_tensors_are_fp16
 from _test_utils.onnx.quantization.utils import assert_nodes_are_quantized
 
+import modelopt.onnx.trt_utils as trt_utils
 from modelopt.onnx.autocast import convert_to_mixed_precision
 from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer
 from modelopt.onnx.quantization.quantize import quantize
-from modelopt.onnx.trt_utils import load_onnx_model
+from modelopt.onnx.trt_utils import get_custom_layers, load_onnx_model
 
 skip_if_no_libcudnn()
 skip_if_no_tensorrt()
@@ -150,6 +151,39 @@ def test_trt_plugin_quantization_int4_awq(tmp_path):
         assert any(n.op == "CustomSkipLayerNormPluginDynamic" for n in graph.nodes)
 
 
+def test_get_custom_layers_file_backed_matches_in_memory(tmp_path, monkeypatch):
+    """Over-limit in-memory ModelProto must yield the same metadata as the file path.
+
+    Regression for the large-ModelProto false-success bug: TensorRT's in-memory
+    ``parser.parse(bytes)`` silently returns an empty network (0 layers/0 tensors) for
+    models at/above the protobuf 2 GiB limit. ``get_custom_layers`` now routes over-limit
+    models through a temporary external-data file and ``parse_from_file()``. This test
+    forces the file-backed path on a small custom-op model and asserts the result matches
+    both the in-memory fast path and the file-path baseline.
+    """
+
+    model = _create_test_model_trt()
+    onnx_path = os.path.join(tmp_path, "model_with_trt_plugin_file_backed.onnx")
+    onnx.save_model(model, onnx_path)
+
+    # Baseline: string path -> parse_from_file.
+    path_layers, path_tensors = get_custom_layers(onnx_path, trt_plugins=None)
+    assert path_layers == ["custom_skip_ln_0"]
+    assert path_tensors
+
+    # In-memory fast path (small model, below the protobuf limit).
+    mem_layers, mem_tensors = get_custom_layers(model, trt_plugins=None)
+    assert mem_layers == path_layers
+    assert set(mem_tensors) == set(path_tensors)
+
+    # Force the over-limit routing: the ModelProto is serialized to a temporary
+    # external-data file and parsed via parse_from_file, matching the baseline.
+    monkeypatch.setattr(trt_utils, "_requires_file_backed_parse", lambda _model: True)
+    fb_layers, fb_tensors = get_custom_layers(model, trt_plugins=None)
+    assert fb_layers == path_layers
+    assert set(fb_tensors) == set(path_tensors)
+
+
 def test_trt_plugin_autocast(tmp_path):
     model = _create_test_model_trt()
     with open(os.path.join(tmp_path, "model_with_trt_plugin_autocast.onnx"), "w") as f: