Fixes ONNX validation and AutoCast handling for ONNX Runtime legacy operators

modelopt/onnx/autocast/graphsanitizer.py

@@ -51,6 +51,7 @@ def __init__(
         self.min_opset = min_opset
         self.max_ir_version = max_ir_version
         self.standard_ops = {schema.name for schema in onnx.defs.get_all_schemas()}
+        self.ort_legacy_ops = onnx_utils.register_ort_legacy_schemas()
         self.custom_ops = None
         self.custom_ops_low_precision_nodes = []
         self.trt_plugins = trt_plugins
@@ -115,16 +116,25 @@ def find_custom_nodes(self) -> None:
         that are not part of the standard ONNX operator set.
         """
         self.custom_ops = {
-            node.op_type for node in self.model.graph.node if node.op_type not in self.standard_ops
+            node.op_type
+            for node in self.model.graph.node
+            if node.op_type not in self.standard_ops
+            and not (not node.domain and node.op_type in self.ort_legacy_ops)
         }
         if self.custom_ops:
-            from modelopt.onnx.trt_utils import infer_types_shapes_tensorrt, set_trt_plugin_domain
+            from modelopt.onnx import trt_utils
+
+            if not trt_utils.TRT_PYTHON_AVAILABLE:
+                logger.warning(
+                    "TensorRT Python bindings are not available; skipping custom-layer introspection."
+                )
+                return
 
             # Set TensorRT plugin domain info in the graph for ORT compatibility
-            self.model = set_trt_plugin_domain(self.model, self.custom_ops)
+            self.model = trt_utils.set_trt_plugin_domain(self.model, self.custom_ops)
 
             # Infer types and shapes in the graph for ORT compatibility
-            self.model = infer_types_shapes_tensorrt(self.model, self.trt_plugins)
+            self.model = trt_utils.infer_types_shapes_tensorrt(self.model, self.trt_plugins)
 
     def remove_disconnected_outputs(self) -> None:
         """Remove disconnected outputs from the model."""

modelopt/onnx/quantization/quantize.py

@@ -43,7 +43,6 @@
 from typing import Any
 
 import onnx
-import onnx.onnx_cpp2py_export.checker as C
 import onnx_graphsurgeon as gs
 import onnxslim
 from onnxruntime.quantization.calibrate import CalibrationDataReader
@@ -77,6 +76,7 @@
 from modelopt.onnx.utils import (
     BASE_MIN_OPSET,
     QDQ_PRECISION_MIN_OPSET,
+    check_model,
     clear_stale_value_info,
     duplicate_shared_constants,
     get_opset_version,
@@ -955,8 +955,8 @@ def quantize(
     # Check if the quantized model is valid
     try:
         logger.info("Validating quantized model")
-        onnx.checker.check_model(output_path)
-    except C.ValidationError as e:
+        check_model(onnx_model, output_path)
+    except onnx.checker.ValidationError as e:
         logger.warning("ONNX model checker failed, check your deployment status")
         logger.warning(e)
 

modelopt/onnx/trt_utils.py

@@ -27,6 +27,7 @@
 
 from modelopt.onnx.logging_config import logger
 from modelopt.onnx.utils import (
+    check_model,
     get_dynamic_graph_inputs,
     get_tensor_by_name,
     parse_shapes_spec,
@@ -453,10 +454,10 @@ def load_onnx_model(
     if use_external_data_format:
         # For large models, use the file path to avoid protobuf size limitation
         model_path_to_check = ir_version_onnx_path or static_shaped_onnx_path or onnx_path
-        onnx.checker.check_model(model_path_to_check)
+        check_model(onnx_model, model_path_to_check)
     else:
         # For smaller models, checking the model object is fine
-        onnx.checker.check_model(onnx_model)
+        check_model(onnx_model)
 
     return (
         onnx_model,

modelopt/onnx/utils.py

@@ -34,6 +34,89 @@
 # Base minimum opset for quantization (opset 19 is the first to support fp16 scales)
 BASE_MIN_OPSET = 19
 
+# Populated permanently by the first legacy-schema registration.
+_ORT_LEGACY_ONNX_DOMAIN_OPS: frozenset[str] | None = None
+
+
+def _convert_ort_parameter(parameter: Any) -> onnx.defs.OpSchema.FormalParameter:
+    """Rebuild an ORT formal parameter using ONNX's separate pybind type."""
+    schema_type = onnx.defs.OpSchema
+    # The two bindings define equivalent enums but do not share Python types, so
+    # translate by stable enum name instead of relying on their numeric values.
+    return schema_type.FormalParameter(
+        parameter.name,
+        parameter.typeStr,
+        parameter.description or "",
+        param_option=getattr(schema_type.FormalParameterOption, parameter.option.name),
+        is_homogeneous=parameter.isHomogeneous,
+    )
+
+
+def _convert_ort_schema(ort_schema: Any) -> onnx.defs.OpSchema:
+    """Rebuild an ORT operator schema in ONNX's independent schema type.
+
+    ORT vendors its own ONNX C++ library, so its pybind objects and registry are
+    not visible to the separately installed ONNX package.
+    """
+    schema_type = onnx.defs.OpSchema
+    # This is the metadata used by check_model(full_check=False). Native ORT
+    # shape-inference functions cannot cross the binary boundary between wheels.
+    return schema_type(
+        ort_schema.name,
+        ort_schema.domain,
+        ort_schema.since_version,
+        ort_schema.doc or "",
+        inputs=[_convert_ort_parameter(parameter) for parameter in ort_schema.inputs],
+        outputs=[_convert_ort_parameter(parameter) for parameter in ort_schema.outputs],
+        type_constraints=[
+            (
+                constraint.type_param_str,
+                list(constraint.allowed_type_strs),
+                constraint.description or "",
+            )
+            for constraint in ort_schema.type_constraints
+        ],
+        attributes=[
+            schema_type.Attribute(
+                attribute.name,
+                getattr(schema_type.AttrType, attribute.type.name),
+                attribute.description or "",
+                required=attribute.required,
+            )
+            for attribute in ort_schema.attributes.values()
+        ],
+    )
+
+
+def register_ort_legacy_schemas() -> frozenset[str]:
+    """Register ORT-only default-domain schemas and return their operator names."""
+    global _ORT_LEGACY_ONNX_DOMAIN_OPS
+    if _ORT_LEGACY_ONNX_DOMAIN_OPS is not None:
+        return _ORT_LEGACY_ONNX_DOMAIN_OPS
+
+    # ORT exposes its independently compiled schema registry only through this
+    # private binding. Import it lazily with schema registration.
+    from onnxruntime.capi._pybind_state import get_all_operator_schema
+
+    onnx_standard_ops = {
+        schema.name for schema in onnx.defs.get_all_schemas_with_history() if not schema.domain
+    }
+    # Filter by name so ORT cannot add alternate schema histories for operators
+    # that ONNX already owns; only genuinely missing default-domain ops are added.
+    legacy_schemas = sorted(
+        (
+            schema
+            for schema in get_all_operator_schema()
+            if not schema.domain and schema.name not in onnx_standard_ops
+        ),
+        key=lambda schema: (schema.name, schema.since_version),
+    )
+    for schema in legacy_schemas:
+        onnx.defs.register_schema(_convert_ort_schema(schema))
+
+    _ORT_LEGACY_ONNX_DOMAIN_OPS = frozenset(schema.name for schema in legacy_schemas)
+    return _ORT_LEGACY_ONNX_DOMAIN_OPS
+
 
 def get_input_names_from_bytes(model_bytes: bytes, external_inputs_only: bool = True) -> list[str]:
     """This function returns the inputs names of the given onnx model in bytes.
@@ -556,8 +639,38 @@ def _get_unique_name(old_name):
     return onnx_model, is_modified
 
 
-def check_model(model: onnx.ModelProto) -> None:
-    """Checks if the given model is valid."""
+def check_model(model: onnx.ModelProto, model_path: str | None = None) -> None:
+    """Check whether a model is structurally valid.
+
+    ONNX Runtime legacy operators registered in the default ONNX domain are
+    accepted by this initial guard. Whether they can execute remains the
+    responsibility of the selected backend.
+
+    Args:
+        model: Loaded in-memory ONNX model. Used for validation unless
+            model_path is supplied and for detecting legacy operators.
+        model_path: Optional file-backed copy to validate. Use this for models
+            with external data or models too large for protobuf serialization.
+    """
+    ort_legacy_ops = register_ort_legacy_schemas()
+    legacy_ops = sorted(
+        {
+            node.op_type
+            for node in model.graph.node
+            if not node.domain and node.op_type in ort_legacy_ops
+        }
+    )
+    if legacy_ops:
+        logger.warning(
+            "Model uses ONNX Runtime legacy operator(s) in the default ONNX domain: %s. "
+            "Execution support is delegated to the selected backend.",
+            legacy_ops,
+        )
+
+    if model_path is not None:
+        onnx.checker.check_model(model_path)
+        return
+
     save_as_external_data = False
     try:
         model_size = model.ByteSize()

tests/unit/onnx/autocast/test_graphsanitizer.py

@@ -17,9 +17,68 @@
 import pytest
 from onnx import TensorProto, helper, numpy_helper
 
+import modelopt.onnx.trt_utils as trt_utils
 from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer
 
 
+def test_find_custom_nodes_without_tensorrt(monkeypatch):
+    """Custom-op discovery should not require the optional TensorRT Python bindings."""
+    x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4])
+    y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4])
+    node = helper.make_node(
+        "FakeTensorRTPlugin",
+        ["X"],
+        ["Y"],
+        name="plugin",
+        domain="test.plugins",
+    )
+    graph = helper.make_graph([node], "custom_op_test", [x], [y])
+    model = helper.make_model(
+        graph,
+        opset_imports=[
+            helper.make_opsetid("", 18),
+            helper.make_opsetid("test.plugins", 1),
+        ],
+    )
+    sanitizer = GraphSanitizer(model)
+
+    monkeypatch.setattr(trt_utils, "TRT_PYTHON_AVAILABLE", False)
+
+    sanitizer.find_custom_nodes()
+
+    assert sanitizer.custom_ops == {"FakeTensorRTPlugin"}
+    assert sanitizer.model.graph.node[0].domain == "test.plugins"
+    assert all(opset.domain != "trt.plugins" for opset in sanitizer.model.opset_import)
+
+
+def test_find_custom_nodes_treats_ort_legacy_op_as_known(monkeypatch):
+    """ORT legacy operators should not be treated as TensorRT plugins."""
+    x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4])
+    y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4])
+    scale = numpy_helper.from_array(np.ones(4, dtype=np.float32), name="scale")
+    node = helper.make_node(
+        "SimplifiedLayerNormalization",
+        ["X", "scale"],
+        ["Y"],
+        name="layer_norm",
+    )
+    graph = helper.make_graph([node], "legacy_op_test", [x], [y], [scale])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)])
+    sanitizer = GraphSanitizer(model)
+
+    monkeypatch.setattr(trt_utils, "TRT_PYTHON_AVAILABLE", True)
+
+    def fail_if_called(*args, **kwargs):
+        pytest.fail("ORT legacy operator entered TensorRT plugin handling")
+
+    monkeypatch.setattr(trt_utils, "set_trt_plugin_domain", fail_if_called)
+    sanitizer.find_custom_nodes()
+
+    assert sanitizer.custom_ops == set()
+    assert sanitizer.model.graph.node[0].domain == ""
+    assert all(opset.domain != "trt.plugins" for opset in sanitizer.model.opset_import)
+
+
 def create_layernorm_model(input_shape, epsilon=1e-5, axis=-1, add_scale=True, add_bias=True):
     """Helper function to create an ONNX model with a decomposed LayerNorm pattern"""
     x = helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)