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..f15f4ecf4cd 100755 --- a/tests/gpu/onnx/quantization/test_plugin.py +++ b/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: