From 1af3aceda94e94187feec9562c3dd87ef83b1d8f Mon Sep 17 00:00:00 2001 From: Gwenaelle Cunha Sergio Date: Fri, 12 Jun 2026 14:15:24 +0000 Subject: [PATCH] [6287717] Preserve trt.plugins custom-op value_info in clear_stale_value_info int4 quantization forces an opset upgrade to >= 21, where ORT performs type inference at InferenceSession load. Custom (trt.plugins) ops have no ORT type-inference function, so their output types only exist via the value_info populated by TensorRT inference. clear_stale_value_info cleared value_info wholesale, dropping those types, which made ORT fail output type inference for the custom op (e.g. "Op (IdentityConv) output arg (X2) type inference failed"). Preserve value_info entries for trt.plugins-domain node outputs; clear the rest as before. Add a GPU regression test that quantizes a model with the built-in CustomSkipLayerNormPluginDynamic plugin at int4 + awq_clip (no plugin .so compilation needed). Signed-off-by: Gwenaelle Cunha Sergio --- modelopt/onnx/utils.py | 20 +++++++++++++------ tests/gpu/onnx/quantization/test_plugin.py | 23 ++++++++++++++++++++++ 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index 951ce2cc98c..64b1fd1c414 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -1868,14 +1868,15 @@ def clear_stale_value_info(model: onnx.ModelProto) -> int: Walks every ``Cast`` node and forces the ``elem_type`` of any ``graph.output`` entry produced by that Cast to match the Cast's ``to`` attribute (the spec-defined contract for a Cast's output dtype). Then - clears ``value_info`` wholesale so ORT/shape-inference re-derives - intermediate-tensor types from the operator graph during session setup. + clears ``value_info`` so ORT/shape-inference re-derives intermediate-tensor + types from the operator graph during session setup -- except entries for + outputs of ``trt.plugins`` custom-op nodes, whose types ORT cannot infer. Args: model: Loaded in-memory onnx ModelProto. Returns: - Total number of entries reconciled or cleared. + Number of Cast outputs reconciled plus value_info entries cleared. """ cast_to_by_output = { node.output[0]: get_cast_to_type(node) @@ -1890,7 +1891,14 @@ def clear_stale_value_info(model: onnx.ModelProto) -> int: o.type.tensor_type.elem_type = to_attr fixed_outputs += 1 - n_vi = len(model.graph.value_info) - if n_vi: + # Outputs of TensorRT-plugin nodes carry types ORT cannot infer so they must survive the + # value_info clear, otherwise ORT fails output type inference for the custom op. + preserve_names = { + out for node in model.graph.node if node.domain == "trt.plugins" for out in node.output + } + preserved = [vi for vi in model.graph.value_info if vi.name in preserve_names] + n_cleared = len(model.graph.value_info) - len(preserved) + if n_cleared: del model.graph.value_info[:] - return fixed_outputs + n_vi + model.graph.value_info.extend(preserved) + return fixed_outputs + n_cleared diff --git a/tests/gpu/onnx/quantization/test_plugin.py b/tests/gpu/onnx/quantization/test_plugin.py index 336e3c7da67..3498150b795 100755 --- a/tests/gpu/onnx/quantization/test_plugin.py +++ b/tests/gpu/onnx/quantization/test_plugin.py @@ -127,6 +127,29 @@ def test_trt_plugin_quantization(tmp_path): assert assert_nodes_are_quantized(quantizable_nodes) +def test_trt_plugin_quantization_int4_awq(tmp_path): + model = _create_test_model_trt() + with open(os.path.join(tmp_path, "model_with_trt_plugin_int4.onnx"), "w") as f: + onnx.save_model(model, f.name) + + # Quantize at int4 with awq_clip (the path that forces opset >= 21). + quantize( + f.name, + quantize_mode="int4", + calibration_method="awq_clip", + calibration_eps=["trt", "cuda:0", "cpu"], + ) + + # The regression was a hard failure at calibration-session load; reaching a + # written output model means the custom op's type survived the value_info clear. + output_onnx_path = f.name.replace(".onnx", ".quant.onnx") + assert os.path.isfile(output_onnx_path) + + # The custom op must still be present (not dropped) in the quantized model. + graph = gs.import_onnx(onnx.load(output_onnx_path)) + assert any(n.op == "CustomSkipLayerNormPluginDynamic" for n in graph.nodes) + + 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: