Fix Inspector crash on multi-method ETRecords

devtools/inspector/_inspector.py

@@ -58,6 +58,7 @@
     gen_etdump_object,
     gen_graphs_from_etrecord,
     get_aot_debug_handle_to_op_name_mapping,
+    get_forward_edge_dialect_program,
     inflate_runtime_output,
     is_debug_output,
     is_inference_output_equal,
@@ -1197,14 +1198,15 @@ def _resolve_reference_graph(
                 debug handle backpropagation fails for "exported_program".
         """
         resolved_graph_name = reference_graph
+        edge_dialect_program = get_forward_edge_dialect_program(self._etrecord)
 
         # Determine the reference graph to use
         if reference_graph is None or reference_graph == "exported_program":
             # Auto-select: try exported_program first, fall back to edge_dialect_exported_program
             if self._etrecord.exported_program and propagate_back_debug_handle(
                 self._etrecord.exported_program,
                 self._etrecord.export_graph_id,
-                self._etrecord.edge_dialect_program,
+                edge_dialect_program,
                 disable_debug_handle_validation,
             ):
                 resolved_graph_name = "exported_program"
@@ -1221,7 +1223,7 @@ def _resolve_reference_graph(
                 )
 
         if resolved_graph_name == "edge_dialect_exported_program":
-            export_program = self._etrecord.edge_dialect_program
+            export_program = edge_dialect_program
             log.info(
                 "Using 'edge_dialect_exported_program' (Edge dialect) as reference graph"
             )

devtools/inspector/_inspector_utils.py

@@ -296,6 +296,18 @@ def is_debug_output(value: Value) -> bool:
     return value.output is not None and value.output.bool_val
 
 
+def get_forward_edge_dialect_program(etrecord: ETRecord) -> Optional[ExportedProgram]:
+    """
+    Return the edge dialect program of the forward method. For multi-method
+    models, ETRecord stores a dict of method name -> edge dialect program, and
+    the Inspector only correlates the forward method.
+    """
+    edge_dialect_program = etrecord.edge_dialect_program
+    if isinstance(edge_dialect_program, dict):
+        return edge_dialect_program.get(FORWARD)
+    return edge_dialect_program
+
+
 def gen_graphs_from_etrecord(
     etrecord: ETRecord, enable_module_hierarchy: bool = False
 ) -> Mapping[str, OperatorGraph]:
@@ -308,9 +320,10 @@ def gen_graphs_from_etrecord(
             )
             for name, exported_program in etrecord.graph_map.items()
         }
-    if etrecord.edge_dialect_program is not None:
+    edge_dialect_program = get_forward_edge_dialect_program(etrecord)
+    if edge_dialect_program is not None:
         op_graph_map[EDGE_DIALECT_GRAPH_KEY] = FXOperatorGraph.gen_operator_graph(
-            etrecord.edge_dialect_program.graph_module,
+            edge_dialect_program.graph_module,
             enable_module_hierarchy=enable_module_hierarchy,
         )
 

devtools/inspector/tests/inspector_test.py

@@ -1293,6 +1293,53 @@ def element_compare(self, a, b) -> float:
                 )
             self.assertIn("Invalid runtime debug handle", str(context.exception))
 
+    def test_inspector_with_multi_method_etrecord(self):
+        class AddOne(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        class TimesTwo(torch.nn.Module):
+            def forward(self, x):
+                return x * 2
+
+        example_inputs = (torch.randn(2, 2),)
+        et_program_manager = to_edge(
+            {
+                "forward": export(AddOne(), example_inputs),
+                "times_two": export(TimesTwo(), example_inputs),
+            },
+            generate_etrecord=True,
+        ).to_executorch()
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            etrecord_path = os.path.join(tmp_dir, "etrecord.bin")
+            et_program_manager.get_etrecord().save(etrecord_path)
+
+            with patch.object(
+                _inspector, "gen_etdump_object", return_value=None
+            ), patch.object(EventBlock, "_gen_from_etdump", return_value=[]):
+                # Used to raise AttributeError: 'dict' object has no attribute 'graph_module'
+                inspector_instance = Inspector(
+                    etdump_path=ETDUMP_PATH,
+                    etrecord=etrecord_path,
+                )
+
+        # Multi-method ETRecords store the edge dialect programs as a dict
+        self.assertIsInstance(inspector_instance._etrecord.edge_dialect_program, dict)
+
+        # The forward method's edge dialect program is used as the reference graph
+        graph_module, graph_name = inspector_instance._resolve_reference_graph(
+            "edge_dialect_exported_program"
+        )
+        self.assertEqual(graph_name, "edge_dialect_exported_program")
+        targets = [
+            str(node.target)
+            for node in graph_module.graph.nodes
+            if node.op == "call_function"
+        ]
+        self.assertTrue(any("add" in target for target in targets))
+        self.assertFalse(any("mul" in target for target in targets))
+
     def test_calculate_numeric_gap_with_reference_graph_name(self):
         """Test calculate_numeric_gap with the reference_graph parameter using a custom graph from graph_map."""
         # Create a context manager to patch functions called by Inspector.__init__