Graduate debug handle in torchao (#2452)

Summary:

This diff maket torchao's debugging infra fully leverage node["from_node"] info and get rid of debug handle.

debug handle, we will miss you 🫡

Reviewed By: jerryzh168

Differential Revision: D76628702

Author: Gasoonjia

test/quantization/pt2e/test_numeric_debugger.py

@@ -37,10 +37,12 @@ def test_simple(self):
         example_inputs = m.example_inputs()
         ep = export_for_training(m, example_inputs, strict=True)
         m = ep.module()
-        self._assert_each_node_has_debug_handle(m)
-        debug_handle_map = self._extract_debug_handles(m)
+        self._assert_each_node_has_from_node_source(m)
+        from_node_source_map = self._extract_from_node_source(m)
 
-        self.assertEqual(len(set(debug_handle_map.values())), len(debug_handle_map))
+        self.assertEqual(
+            len(set(from_node_source_map.values())), len(from_node_source_map)
+        )
 
     @unittest.skip("debug flow not working on model with conditional control flow")
     def test_control_flow(self):
@@ -49,37 +51,42 @@ def test_control_flow(self):
         ep = export_for_training(m, example_inputs, strict=True)
         m = ep.module()
 
-        self._assert_each_node_has_debug_handle(m)
-        debug_handle_map = self._extract_debug_handles(m)
+        self._assert_each_node_has_from_node_source(m)
+        from_node_source_map = self._extract_from_node_source(m)
 
-        self.assertEqual(len(set(debug_handle_map.values())), len(debug_handle_map))
+        self.assertEqual(
+            len(set(from_node_source_map.values())), len(from_node_source_map)
+        )
 
     def test_copy_preserve_handle(self):
         m = TestHelperModules.Conv2dThenConv1d()
         example_inputs = m.example_inputs()
         ep = torch.export.export(m, example_inputs, strict=True)
         m = ep.module()
 
-        self._assert_each_node_has_debug_handle(m)
-        debug_handle_map_ref = self._extract_debug_handles(m)
+        self._assert_each_node_has_from_node_source(m)
+        from_node_source_map_ref = self._extract_from_node_source(m)
 
         ep_copy = copy.copy(ep)
-        debug_handle_map = self._extract_debug_handles(ep_copy.module())
+        from_node_source_map = self._extract_from_node_source(ep_copy.module())
 
-        self._assert_each_node_has_debug_handle(ep)
-        self.assertEqual(debug_handle_map, debug_handle_map_ref)
+        self._assert_each_node_has_from_node_source(ep)
+        self.assertEqual(from_node_source_map, from_node_source_map_ref)
 
     def test_deepcopy_preserve_handle(self):
         m = TestHelperModules.Conv2dThenConv1d()
         example_inputs = m.example_inputs()
         ep = torch.export.export(m, example_inputs, strict=True)
 
-        debug_handle_map_ref = self._extract_debug_handles(ep.module())
+        from_node_source_map_ref = self._extract_from_node_source(ep.module())
         ep_copy = copy.deepcopy(ep)
-        debug_handle_map = self._extract_debug_handles(ep_copy.module())
+        from_node_source_map = self._extract_from_node_source(ep_copy.module())
 
-        self._assert_each_node_has_debug_handle(ep.module())
-        self.assertEqual(debug_handle_map, debug_handle_map_ref)
+        self._assert_each_node_has_from_node_source(ep.module())
+        self.assertEqual(from_node_source_map, from_node_source_map_ref)
+        self.assertEqual(
+            set(from_node_source_map.values()), set(from_node_source_map_ref.values())
+        )
 
     @unittest.skip(
         "torch._dynamo.exc.FailOnRecompileLimitHit: recompile_limit reached with one_graph=True. Excessive recompilations can degrade performance due to the compilation overhead of each recompilation. To monitor recom..."
@@ -90,16 +97,16 @@ def test_re_export_preserve_handle(self):
         ep = export_for_training(m, example_inputs, strict=True)
         m = ep.module()
 
-        self._assert_each_node_has_debug_handle(m)
-        debug_handle_map_ref = self._extract_debug_handles(m)
+        self._assert_each_node_has_from_node_source(m)
+        from_node_source_map_ref = self._extract_from_node_source(m)
 
         ep_reexport = export_for_training(m, example_inputs, strict=True)
         m_reexport = ep_reexport.module()
 
-        self._assert_each_node_has_debug_handle(m_reexport)
-        debug_handle_map = self._extract_debug_handles(m_reexport)
+        self._assert_each_node_has_from_node_source(m_reexport)
+        from_node_source_map = self._extract_from_node_source(m_reexport)
 
-        self.assertEqual(debug_handle_map, debug_handle_map_ref)
+        self.assertEqual(from_node_source_map, from_node_source_map_ref)
 
     @unittest.skip(
         "torch._dynamo.exc.FailOnRecompileLimitHit: recompile_limit reached with one_graph=True. Excessive recompilations can degrade performance due to the compilation overhead of each recompilation. To monitor recom..."
@@ -110,19 +117,19 @@ def test_run_decompositions_same_handle_id(self):
         ep = export_for_training(m, example_inputs, strict=True)
         m = ep.module()
 
-        self._assert_each_node_has_debug_handle(m)
-        debug_handle_map_ref = self._extract_debug_handles(m)
+        self._assert_each_node_has_from_node_source(m)
+        from_node_source_map_ref = self._extract_from_node_source(m)
 
         ep_copy = copy.copy(ep)
         ep_copy = ep_copy.run_decompositions()
         m_decomposed = ep_copy.module()
 
-        self._assert_each_node_has_debug_handle(m_decomposed)
-        debug_handle_map = self._extract_debug_handles(m_decomposed)
+        self._assert_each_node_has_from_node_source(m_decomposed)
+        from_node_source_map = self._extract_from_node_source(m_decomposed)
 
         # checking the map still has the same ids, the node may change
         self.assertEqual(
-            set(debug_handle_map.values()), set(debug_handle_map_ref.values())
+            set(from_node_source_map.values()), set(from_node_source_map_ref.values())
         )
 
     @unittest.skip(
@@ -139,22 +146,23 @@ def test_run_decompositions_map_handle_to_new_nodes(self):
             ep = export_for_training(m, example_inputs, strict=True)
             m = ep.module()
 
-            self._assert_each_node_has_debug_handle(m)
-            pre_decomp_to_debug_handle_map_ref = (
-                self._extract_debug_handles_with_prev_decomp_op(m)
+            self._assert_each_node_has_from_node_source(m)
+            pre_decomp_to_from_node_source_map_ref = (
+                self._extract_from_node_source_with_prev_decomp_op(m)
             )
 
             ep_copy = copy.copy(ep)
             ep_copy = ep_copy.run_decompositions()
             m_decomposed = ep_copy.module()
-            self._assert_each_node_has_debug_handle(m_decomposed)
-            pre_decomp_to_debug_handle_map = (
-                self._extract_debug_handles_with_prev_decomp_op(m_decomposed)
+            self._assert_each_node_has_from_node_source(m_decomposed)
+            pre_decomp_to_from_node_source_map = (
+                self._extract_from_node_source_with_prev_decomp_op(m_decomposed)
             )
 
-            # checking the map still has the same ids, the node may change
+            # checking the map still has the same infos, the node may change
             self.assertEqual(
-                pre_decomp_to_debug_handle_map, pre_decomp_to_debug_handle_map_ref
+                pre_decomp_to_from_node_source_map,
+                pre_decomp_to_from_node_source_map_ref,
             )
 
     def test_prepare_for_propagation_comparison(self):
@@ -178,18 +186,18 @@ def test_added_node_gets_unique_id(self) -> None:
         example_inputs = m.example_inputs()
         ep = export_for_training(m, example_inputs, strict=True)
 
-        ref_handles = self._extract_debug_handles(ep.module())
-        ref_counter = Counter(ref_handles.values())
+        ref_from_node_source = self._extract_from_node_source(ep.module())
+        ref_counter = Counter(ref_from_node_source.values())
 
         for k, v in ref_counter.items():
             self.assertEqual(
                 v,
                 1,
-                msg=f"For handle {k}, there were {v} nodes with that handle, but expected only 1",
+                msg=f"For from_node info {k}, there were {v} nodes with that info, but expected only 1",
             )
 
-        # Now that we have unique ids, add a new node into the graph and re-generate
-        # to make sure that the new node gets a unique id.
+        # Now that we have unique infos, add a new node into the graph and re-generate
+        # to make sure that the new node gets a unique info.
         last_node = next(iter(reversed(ep.graph.nodes)))
         with ep.graph.inserting_before(last_node):
             arg = last_node.args[0]
@@ -200,30 +208,39 @@ def test_added_node_gets_unique_id(self) -> None:
             arg.replace_all_uses_with(n, lambda x: x != n)
         ep.graph_module.recompile()
 
-        # Regenerate handles, make sure only the new relu node has a new id, and
-        # it doesn't clash with any of the existing ids.
+        # Regenerate from_node info, make sure only the new relu node has a new info, and
+        # it doesn't clash with any of the existing infos.
 
         m = ep.module()
-        self._assert_each_node_has_debug_handle(m)
-        handles_after_modification = self._extract_debug_handles(m)
-        handles_counter = Counter(handles_after_modification.values())
-        for name, handle in ref_handles.items():
-            self.assertIn(name, handles_after_modification)
-            # Check that handle was unchanged.
-            self.assertEqual(handles_after_modification[name], handle)
+        self._assert_each_node_has_from_node_source(m)
+        from_node_source_after_modification = self._extract_from_node_source(m)
+        from_node_source_counter = Counter(from_node_source_after_modification.values())
+        for name, from_node_source in ref_from_node_source.items():
+            self.assertIn(name, from_node_source_after_modification)
+            # Check that from_node info was unchanged.
+            self.assertEqual(
+                from_node_source_after_modification[name], from_node_source
+            )
             # Check that total count was unchanged.
-            ref_count = ref_counter[handle]
-            after_count = handles_counter[handle]
+            ref_count = ref_counter[from_node_source]
+            after_count = from_node_source_counter[from_node_source]
             self.assertEqual(
                 after_count,
                 ref_count,
-                msg=f"For handle {handle}, there were {after_count} nodes with that handle, but expected only {ref_count}",
+                msg=f"For from_node info {from_node_source}, there were {after_count} nodes with that info, but expected only {ref_count}",
             )
 
-        # Check for relu specifically. Avoid hardcoding the handle id since it
+        # Check for relu specifically. Avoid hardcoding the from_node info since it
         # may change with future node ordering changes.
-        self.assertNotIn(handles_after_modification["relu_default"], ref_counter)
-        self.assertEqual(handles_counter[handles_after_modification["relu_default"]], 1)
+        self.assertNotIn(
+            from_node_source_after_modification["relu_default"], ref_counter
+        )
+        self.assertEqual(
+            from_node_source_counter[
+                from_node_source_after_modification["relu_default"]
+            ],
+            1,
+        )
 
 
 if __name__ == "__main__":