increase torch dynamo cache size limit to support all tests (#2470)

increase torch dynamo cache size limit to support all tests (#2470)

Summary:

Some tests was suppressed due to cache size limitation. This diff increase the cache size to unblock the issue.

Reviewed By: jerryzh168

Differential Revision: D76794879

Author: Gasoonjia

test/quantization/pt2e/test_numeric_debugger.py

@@ -23,15 +23,17 @@
 if TORCH_VERSION_AT_LEAST_2_8:
     from torch.export import export_for_training
 
+# Increase cache size limit to avoid FailOnRecompileLimitHit error when running multiple tests
+# that use export_for_training, which causes many dynamo recompilations
+if TORCH_VERSION_AT_LEAST_2_8:
+    torch._dynamo.config.cache_size_limit = 128
+
 
 @unittest.skipIf(
     not TORCH_VERSION_AT_LEAST_2_8, "Requires torch 2.8 and above, including nightly"
 )
 @unittest.skipIf(IS_WINDOWS, "Windows not yet supported for torch.compile")
 class TestNumericDebuggerInfra(PT2ENumericDebuggerTestCase):
-    @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..."
-    )
     def test_simple(self):
         m = TestHelperModules.Conv2dThenConv1d()
         example_inputs = m.example_inputs()
@@ -88,9 +90,6 @@ def test_deepcopy_preserve_handle(self):
             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..."
-    )
     def test_re_export_preserve_handle(self):
         m = TestHelperModules.Conv2dThenConv1d()
         example_inputs = m.example_inputs()
@@ -108,9 +107,6 @@ def test_re_export_preserve_handle(self):
 
         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..."
-    )
     def test_run_decompositions_same_handle_id(self):
         m = TestHelperModules.Conv2dThenConv1d()
         example_inputs = m.example_inputs()
@@ -132,9 +128,6 @@ def test_run_decompositions_same_handle_id(self):
             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..."
-    )
     def test_run_decompositions_map_handle_to_new_nodes(self):
         test_models = [
             TestHelperModules.TwoLinearModule(),

torchao/testing/pt2e/utils.py

@@ -180,7 +180,7 @@ def _extract_from_node_source_with_prev_decomp_op_from_node(node):
             nonlocal prev_decomp_op_to_from_node_source_map
             if FROM_NODE_KEY in node.meta and node.meta[FROM_NODE_KEY] is not None:
                 prev_decomp_op = str(node.meta.get("nn_module_stack"))
-                from_node_source = node.meta[FROM_NODE_KEY]
+                from_node_source = _extract_node_source_debug_info(node)
                 if prev_decomp_op not in prev_decomp_op_to_from_node_source_map:
                     prev_decomp_op_to_from_node_source_map[prev_decomp_op] = (
                         from_node_source
@@ -189,8 +189,10 @@ def _extract_from_node_source_with_prev_decomp_op_from_node(node):
                     assert (
                         prev_decomp_op_to_from_node_source_map[prev_decomp_op]
                         == from_node_source
-                    ), f"Node {node} has different from_node info {from_node_source}"
-                    "than previous node sharing the same decomp op {prev_decomp_op}"
+                    ), (
+                        f"Node {node} has different from_node info {from_node_source}"
+                        f"than previous node sharing the same decomp op {prev_decomp_op}"
+                    )
 
         bfs_trace_with_node_process(
             model, _extract_from_node_source_with_prev_decomp_op_from_node