From 7bdb9e36dd905cbcbe4d0a94a4bb8a94e5230016 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Mon, 6 Jul 2026 02:56:57 -0700 Subject: [PATCH 1/2] [6403893][ONNX][AutoCast] Fix false-success TRT parse for large ModelProto get_custom_layers silently returned 0 layers/tensors for in-memory ModelProtos at/above the protobuf 2GiB limit. Route such models through a temporary external-data file and parse_from_file(), matching the string-path behavior. Verified on the 8.12GB VLA trunk (28 layers/5468 tensors) and a synthetic 2.2GB model. Adds a regression test forcing the file-backed path. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- modelopt/onnx/trt_utils.py | 55 ++++++++++++++++++++-- tests/gpu/onnx/quantization/test_plugin.py | 36 +++++++++++++- 2 files changed, 86 insertions(+), 5 deletions(-) diff --git a/modelopt/onnx/trt_utils.py b/modelopt/onnx/trt_utils.py index b95dab46df1..b407fdc5411 100644 --- a/modelopt/onnx/trt_utils.py +++ b/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)}") diff --git a/tests/gpu/onnx/quantization/test_plugin.py b/tests/gpu/onnx/quantization/test_plugin.py index 3498150b795..4ddd6b8e6d1 100755 --- a/tests/gpu/onnx/quantization/test_plugin.py +++ b/tests/gpu/onnx/quantization/test_plugin.py @@ -25,7 +25,7 @@ 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 +150,40 @@ 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. + """ + import modelopt.onnx.trt_utils as trt_utils + + 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: From f764597c2fe9b83b01a1d37de4b6ac44239338e8 Mon Sep 17 00:00:00 2001 From: Gal Hubara-Agam <96368689+galagam@users.noreply.github.com> Date: Tue, 7 Jul 2026 15:12:14 +0300 Subject: [PATCH 2/2] fix import of trt_utils Signed-off-by: Gal Hubara-Agam <96368689+galagam@users.noreply.github.com> --- tests/gpu/onnx/quantization/test_plugin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/gpu/onnx/quantization/test_plugin.py b/tests/gpu/onnx/quantization/test_plugin.py index 4ddd6b8e6d1..f15f4ecf4cd 100755 --- a/tests/gpu/onnx/quantization/test_plugin.py +++ b/tests/gpu/onnx/quantization/test_plugin.py @@ -22,6 +22,7 @@ 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 @@ -160,7 +161,6 @@ def test_get_custom_layers_file_backed_matches_in_memory(tmp_path, monkeypatch): 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. """ - import modelopt.onnx.trt_utils as trt_utils model = _create_test_model_trt() onnx_path = os.path.join(tmp_path, "model_with_trt_plugin_file_backed.onnx")