Fix node traversal when setting up stacktrace preservation hooks (#118252)

We only want to traverse over each node in the graph exactly once, and we do that by inserting nodes into the "seen" set. The issue is that we forget to check the "seen" set when inserting the root nodes. Typically that is not a problem, because the root nodes are from the different outputs and thus usually correspond to different nodes. With split_with_sizes, though all of the outputs correspond to the same node, ands this leads to the node being iterated over 3 times, and 3 sets of hooks being attached to the same node.

Pull Request resolved: #118252
Approved by: https://github.com/zou3519
ghstack dependencies: #117552, #118234, #118249

Author: soulitzer

test/dynamo/test_aot_autograd.py

@@ -885,6 +885,25 @@ def _prepare_model_args():
             ),
         )
 
+    def test_split_with_sizes_aot_autograd_cleans_up_traceback_meta(self):
+        from torch._functorch.aot_autograd import setup_stacktrace_preservation_hooks
+
+        def fn(result, split_sizes):
+            rs = torch.ops.aten.split_with_sizes(result, split_sizes.tolist())
+            return rs
+
+        example_inputs = (
+            torch.randn(32, requires_grad=True),
+            torch.tensor((7, 16, 9)),
+        )
+        outs = fn(*example_inputs)
+        setup_stacktrace_preservation_hooks([out.grad_fn for out in outs])
+        with fx_traceback.preserve_node_meta():
+            (outs[0].sum() + outs[1].sum() + outs[2].sum()).backward()
+
+        self.assertNotIn("grad_fn_seq_nr", fx_traceback.current_meta)
+        self.assertNotIn("in_grad_fn", fx_traceback.current_meta)
+
     # https://github.com/pytorch/pytorch/issues/110121
     def test_aot_export_joint_simple_repro(self):
         class Mod(torch.nn.Module):

torch/_functorch/_aot_autograd/logging_utils.py

@@ -65,7 +65,7 @@ def iter_graph(roots):
         seen = set()
         q = collections.deque()  # type: ignore[var-annotated]
         for node in roots:
-            if node is not None:
+            if node is not None and node not in seen:
                 seen.add(node)
                 q.append(node)