Inject identity nodes in sanitizer; revert existing logic; update test

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

Author: aboubezari

modelopt/onnx/autocast/graphsanitizer.py

@@ -332,20 +332,26 @@ def sanitize_io_casts(self) -> None:
         model_output_names = {output.name for output in self.model.graph.output}
         nodes_to_add = []
         for node in self.model.graph.node:
-            if node.op_type == "Cast":
-                if node.input[0] in model_input_names and node.output[0] in model_output_names:
-                    cast_input_name = node.input[0]
-                    cast_output_name = node.output[0]
-                    cast_new_output_name = cast_input_name + "_io_cast_identity"
-                    nodes_to_add.append(
-                        helper.make_node(
-                            "Identity",
-                            inputs=[cast_new_output_name],
-                            outputs=[cast_output_name],
-                            name=node.name + "_io_cast_identity",
-                        )
+            if (
+                node.op_type == "Cast"
+                and node.input
+                and node.output
+                and node.input[0] in model_input_names
+                and node.output[0] in model_output_names
+            ):
+                # 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",
                     )
-                    node.output[0] = cast_new_output_name
+                )
+                # Rewire Cast to produce the new intermediate
+                node.output[0] = cast_new_output_name
 
         for node in nodes_to_add:
             self.model.graph.node.append(node)