[6771663] Preserve ONNX API output types when wiring casts (#2451)

### What does this PR do?

Type of change: Bug fix

Preserves the public ONNX graph I/O types captured at the API boundary
when `PrecisionConverter` wires output casts. Type inference can change
the working graph's output declaration before conversion; consulting
that mutated declaration caused the required cast back to the original
public type to be discarded and metadata restoration to fail.

The converter now derives its I/O type map from the preserved boundary
metadata and uses that map when deciding whether a cast should become a
public graph output. A regression test covers an FP32 output whose
working declaration is inferred as FP16, and the changelog records the
corrected behavior.

### Usage

```python
# No API changes are required. Existing conversions now preserve the original
# public I/O declarations when keep_io_types=True.
converted = convert_to_f16(model, keep_io_types=True)
```

### Testing

- Ran `pytest tests/unit/onnx/autocast/test_precisionconverter.py` (186
passed).
- Ran `pytest tests/unit/onnx/autocast` (249 passed).
- Ran Ruff check and format validation on the changed Python files.
- Verified the original minimal end-to-end reproduction with Python 3.12
and TensorRT 10.16.1.11; conversion now completes without the
output-metadata restoration error.
- Verified the full CLI quantization path proceeds through the formerly
failing one-Q/DQ scheme and successfully benchmarks the generated
TensorRT engine.

### 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)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- 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)?:
✅
- Did you get Claude approval on this PR?: N/A

### Additional Information

Tracking: [6771663]


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

## Summary by CodeRabbit

- **Bug Fixes**
- Fixed ONNX FP16 conversion to preserve public graph output types when
type inference changes internal declarations.
- Ensured output casts are inserted correctly when preserving
input/output types is enabled.

- **Tests**
- Added regression coverage confirming preserved output types, correct
cast insertion, and valid ONNX model generation.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Ajinkya Rasane <ajinkyaashwin@gmail.com>
Co-authored-by: Ajinkya Rasane <ajinkyaashwin@gmail.com>

CHANGELOG.rst

@@ -38,6 +38,7 @@ Changelog
 
 **Bug Fixes**
 
+- Fix ONNX FP16 conversion failing to preserve public output types when type inference changes a graph output declaration before output casts are inserted.
 - Fix ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` storing the MoE router at Megatron's ``moe_router_dtype``, which is a routing *compute* dtype, not a storage one. The router now exports at the export ``dtype`` like every other unquantized weight, matching what ``hf_ptq.py`` and the released NVFP4 checkpoints contain; pass ``moe_router_dtype`` to ``export_mcore_gpt_to_hf`` explicitly if you want the old fp32 storage.
 - Fix unified Megatron export writing a second, unreferenced copy of the vocab embedding when a model with MTP layers is exported with pipeline parallelism. The duplicate was never loaded but inflated the checkpoint by the size of the embedding (about 1 GB for Qwen3.6-35B-A3B); re-export to reclaim the space.
 - Fail fast on non-finite AutoQuantize output gradients with an actionable error before accumulating sensitivity scores, without changing attention backend settings.

modelopt/onnx/autocast/precisionconverter.py

@@ -142,13 +142,6 @@ class PrecisionConverter:
         self.low_precision_type = PRECISION_MAP[low_precision_type]
         self.high_precision_type = PRECISION_MAP["fp32"]
 
-        # Preserve original network inputs and outputs for sanity checks
-        self.original_network_io = {
-            io.name: io.type.tensor_type.elem_type for io in self.model.graph.input
-        }
-        self.original_network_io.update(
-            {io.name: io.type.tensor_type.elem_type for io in self.model.graph.output}
-        )
         self.original_network_io_metadata = (
             {
                 "input": [deepcopy(io) for io in self.model.graph.input],
@@ -160,6 +153,13 @@ class PrecisionConverter:
                 for field, values in original_network_io_metadata.items()
             }
         )
+        # Preserve the public I/O types captured at the API boundary. Type inference may have
+        # changed the working model's declarations before the converter is initialized.
+        self.original_network_io = {
+            io.name: io.type.tensor_type.elem_type
+            for values in self.original_network_io_metadata.values()
+            for io in values
+        }
         self.min_opset = min_opset
         self.max_ir_version = max_ir_version
         self.trt_plugins = trt_plugins
@@ -1440,7 +1440,10 @@ class PrecisionConverter:
         # Update network output
         for output in self.model.graph.output:
             if output.name == tensor_name and (
-                (self.keep_io_types and cast_to.onnx_type == output.type.tensor_type.elem_type)
+                (
+                    self.keep_io_types
+                    and cast_to.onnx_type == self.original_network_io.get(tensor_name)
+                )
                 or (
                     not self.keep_io_types
                     and cast_to.onnx_type == self.low_precision_type.onnx_type

tests/unit/onnx/autocast/test_precisionconverter.py

@@ -2220,6 +2220,33 @@ def test_convert_to_f16_restores_public_io_metadata_from_entry_boundary():
     onnx.checker.check_model(converted, full_check=True)
 
 
+def test_convert_to_f16_preserves_declared_output_type_after_inference_changes_it():
+    graph_input = helper.make_tensor_value_info("X", TensorProto.FLOAT16, [2, 3])
+    graph_output = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 3])
+    node = helper.make_node("Identity", ["X"], ["Y"], name="Identity_0")
+    graph = helper.make_graph([node], "inferred_output_type", [graph_input], [graph_output])
+    model = helper.make_model(
+        graph,
+        producer_name="inferred_output_type",
+        opset_imports=[helper.make_opsetid("", 19)],
+        ir_version=10,
+    )
+
+    converted = convert_to_f16(
+        model, keep_io_types=True, op_block_list=[], trt_plugins=[], opset=19
+    )
+
+    output = next(vi for vi in converted.graph.output if vi.name == "Y")
+    assert output.type.tensor_type.elem_type == TensorProto.FLOAT
+    output_producers = [node for node in converted.graph.node if "Y" in node.output]
+    assert len(output_producers) == 1
+    assert output_producers[0].op_type == "Cast"
+    assert next(attr.i for attr in output_producers[0].attribute if attr.name == "to") == (
+        TensorProto.FLOAT
+    )
+    onnx.checker.check_model(converted, full_check=True)
+
+
 def test_convert_to_f16_combines_op_and_node_exclusions(simple_model):
     model, *_ = simple_model
     converted = convert_to_f16(