address review comments

Signed-off-by: Ali Boubezari <[email protected]>

Author: aboubezari

modelopt/onnx/autocast/graphsanitizer.py

@@ -351,8 +351,9 @@ def sanitize_io_casts(self) -> None:
         """
         model_input_names = {input.name for input in self.model.graph.input}
         model_output_names = {output.name for output in self.model.graph.output}
-        nodes_to_add = []
-        for node in self.model.graph.node:
+        insertions: list[tuple[int, onnx.NodeProto]] = []
+
+        for idx, node in enumerate(self.model.graph.node):
             if (
                 node.op_type == "Cast"
                 and node.input
@@ -363,23 +364,19 @@ def sanitize_io_casts(self) -> None:
                 # Unique per graph output to avoid collisions when multiple outputs are cast from the same input
                 cast_output_name = node.output[0]
                 cast_new_output_name = f"{cast_output_name}__io_cast_src"
-                nodes_to_add.append(
-                    helper.make_node(
-                        "Identity",
-                        inputs=[cast_new_output_name],
-                        outputs=[cast_output_name],
-                        name=f"{node.name}__io_cast_identity",
-                    )
+                identity_node = helper.make_node(
+                    "Identity",
+                    inputs=[cast_new_output_name],
+                    outputs=[cast_output_name],
+                    name=f"{node.name}__io_cast_identity",
                 )
                 # Rewire Cast to produce the new intermediate
                 node.output[0] = cast_new_output_name
+                insertions.append((idx + 1, identity_node))
 
-        for node in nodes_to_add:
-            self.model.graph.node.append(node)
-
-        # Make sure the graph is topologically sorted
-        gs_graph = gs.import_onnx(self.model).cleanup().toposort()
-        self.model = gs.export_onnx(gs_graph)
+        # Insert Identities in-order right after their corresponding Casts
+        for offset, (pos, id_node) in enumerate(insertions):
+            self.model.graph.node.insert(pos + offset, id_node)
 
     def _create_layernorm_node(self, pattern: dict) -> onnx.NodeProto:
         """Create a LayerNormalization node with optional bias."""

tests/unit/onnx/autocast/test_precisionconverter.py

@@ -27,6 +27,7 @@
 
 LATEST_IR_VERSION_SUPPORTED_BY_ORT = 10
 
+
 def low_precision_onnx_type(low_precision_type_str):
     return TensorProto.FLOAT16 if low_precision_type_str == "fp16" else TensorProto.BFLOAT16
 
@@ -1103,6 +1104,7 @@ def test_multiple_output_node_casted_to_output(
     )
     onnx.checker.check_model(converted_model)
 
+
 @pytest.fixture
 def model_with_casted_input_to_output():
     """Create a model with an output produced by a Cast node."""
@@ -1170,4 +1172,4 @@ def test_casted_input_to_output_model(
     converted_model = converter.convert(
         high_precision_nodes=["cast_input"], low_precision_nodes=["add1", "add2"]
     )
-    onnx.checker.check_model(converted_model)
+    onnx.checker.check_model(converted_model)