[6385267][ONNX][Autocast] Preserve graph-output Cast producers

modelopt/onnx/autocast/precisionconverter.py

@@ -875,9 +875,11 @@ def _remove_preexisting_casts(self) -> None:
                     init.name for init in self.model.graph.initializer
                 }
                 if is_fp_cast and not input_from_initializer:
-                    # Keep cast nodes that are necessary producers of network outputs
-                    if any(node.input[0] == out.name for out in self.model.graph.output) and any(
-                        node.output[0] == out.name for out in self.model.graph.output
+                    # Keep cast nodes that are necessary producers of network outputs.
+                    if any(
+                        node.output[0] == out.name
+                        and cast_to_type == out.type.tensor_type.elem_type
+                        for out in self.model.graph.output
                     ):
                         continue
                     nodes_to_remove.append(node)

tests/unit/onnx/autocast/test_precisionconverter.py

@@ -20,6 +20,7 @@
 
 import modelopt.onnx.autocast.utils as utils
 import modelopt.onnx.utils as onnx_utils
+from modelopt.onnx.autocast.convert import convert_to_mixed_precision
 from modelopt.onnx.autocast.logging_config import configure_logging
 from modelopt.onnx.autocast.precisionconverter import PrecisionConverter
 
@@ -87,6 +88,37 @@ def test_graph_converter_init(simple_model, use_standalone_type_inference):
     assert converter.keep_io_types
 
 
+def test_convert_preserves_cast_chain_graph_output(tmp_path):
+    x = helper.make_tensor_value_info("in0", TensorProto.FLOAT, [2])
+    y = helper.make_tensor_value_info("t2", TensorProto.FLOAT, [2])
+    cast_to_float16 = helper.make_node(
+        "Cast", ["in0"], ["t1"], name="cast_to_float16", to=TensorProto.FLOAT16
+    )
+    cast_to_float = helper.make_node(
+        "Cast", ["t1"], ["t2"], name="cast_to_float", to=TensorProto.FLOAT
+    )
+    graph = helper.make_graph([cast_to_float16, cast_to_float], "g", [x], [y])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 20)])
+    model.ir_version = 10
+    onnx.checker.check_model(model)
+
+    model_path = tmp_path / "cast_chain_output.onnx"
+    onnx.save(model, model_path)
+
+    converted_model = convert_to_mixed_precision(
+        onnx_path=str(model_path), low_precision_type="fp16", providers=["cpu"]
+    )
+
+    onnx.checker.check_model(converted_model)
+    output_producers = [
+        node
+        for node in converted_model.graph.node
+        if converted_model.graph.output[0].name in node.output
+    ]
+    assert len(output_producers) == 1
+    assert output_producers[0].op_type == "Cast"
+
+
 @pytest.mark.parametrize("keep_io_types", [True, False])
 @pytest.mark.parametrize("low_precision_type", ["fp16", "bf16"])
 @pytest.mark.parametrize("use_standalone_type_inference", [True, False])