[6287717][ONNX][Quantization] Preserve trt.plugins custom-op value_info in clear_stale_value_info (#1697)

### What does this PR do?

Type of change: Bug fix

INT4 quantization upgrades the model to opset >= 21, at which point ONNX
Runtime
runs type inference while building the AWQ calibration
`InferenceSession`. Custom
ops backed by TensorRT plugins (domain `trt.plugins`) have no ORT
type-inference
function, so their output types are only known from the `value_info`
that TensorRT
type/shape inference populated earlier in preprocessing.
`clear_stale_value_info`
cleared `value_info` wholesale, dropping those types, so ORT failed
output type
inference for the custom op at model load, e.g.:

```
Node (Conv-2) Op (IdentityConv) output arg (X2) type inference failed
```

- `modelopt/onnx/utils.py`: in `clear_stale_value_info`, preserve
`value_info`
entries for outputs of `trt.plugins`-domain nodes (which ORT cannot
re-derive);
  clear the rest as before.
- `tests/gpu/onnx/quantization/test_plugin.py`: add a regression test
quantizing a
model with the built-in `CustomSkipLayerNormPluginDynamic` plugin at
INT4 +
awq_clip (the opset >= 21 path), asserting the quantized model is
produced and the
  custom op survives.

### Usage

```python
python -m modelopt.onnx.quantization \
    --onnx_path=model.onnx \
    --quantize_mode=int4 \
    --calibration_method=awq_clip \
    --trt_plugins=/path/to/plugin.so
```

### Testing

- `pytest tests/gpu/onnx/quantization/test_plugin.py -k int4_awq` —
fails before the fix
(ORT type-inference error at calibration-session load) and passes after.
The full
`test_plugin.py` (including the existing INT8 quantization and autocast
cases) passes.
- The example
[here](https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/onnx_ptq/README.md#quantize-an-onnx-model-with-custom-op)
also failed before this fix, now passes.

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors).

- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A

### Additional info
Fixing regression inserted by
https://github.com/NVIDIA/Model-Optimizer/pull/1565

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Preserve metadata for TensorRT plugin outputs during cleanup and
correctly reconcile output data types so custom plugin ops remain intact
after optimization/quantization.
* **Tests**
* Added a GPU ONNX regression test covering int4 quantization with AWQ
calibration to ensure TensorRT plugins are retained.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Gwenaelle Cunha Sergio <gcunhasergio@nvidia.com>

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

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: