Repository: pytorch
Issue: AOTDispatcher loses stack traces of bw nodes from compiled autograd
Labels: triaged,oncall: pt2,module: aotdispatch,module: compiled autograd,module: pt2-dispatcher

repro:
```
import torch


@torch.compile
def f(x):
    return torch.matmul(x, x).sin()

x = torch.randn(4, 4, requires_grad=True)
with torch._dynamo.utils.maybe_enable_compiled_autograd(True):
    out = f(x)
    out.sum().backward()
```

running with `TORCH_LOGS="compiled_autograd_verbose,aot"` per @xmfan, you can see that the compiled autograd graph shows info about backward nodes in the stack trace, but we get "no stacktrace found" in the AOT inference graph:
```
INFO: TRACED GRAPH
 ===== Joint graph 0 =====
 /home/hirsheybar/local/b/pytorch/torch/fx/_lazy_graph_module.py class joint_helper(torch.nn.Module):
    def forward(self, primals, tangents):
        primals_1: "f32[4, 4][4, 1]cpu"; tangents_1: "f32[4, 4][4, 1]cpu";

        primals_1, tangents_1, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec)
         # File: /home/hirsheybar/local/b/pytorch/tmp5.py:6 in f, code: return torch.matmul(x, x).sin()
        mm: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(primals_1, primals_1)
        sin: "f32[4, 4][4, 1]cpu" = torch.ops.aten.sin.default(mm)
        cos: "f32[4, 4][4, 1]cpu" = torch.ops.aten.cos.default(mm);  mm = None
        mul: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mul.Tensor(tangents_1, cos);  tangents_1 = cos = None
        permute: "f32[4, 4][1, 4]cpu" = torch.ops.aten.permute.default(primals_1, [1, 0])
        mm_1: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(permute, mul);  permute = None
        permute_1: "f32[4, 4][1, 4]cpu" = torch.ops.aten.permute.default(primals_1, [1, 0]);  primals_1 = None
        mm_2: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(mul, permute_1);  mul = permute_1 = None

         # File: /home/hirsheybar/local/b/pytorch/tmp5.py:6 in f, code: return torch.matmul(x, x).sin()
        add: "f32[4, 4][4, 1]cpu" = torch.ops.aten.add.Tensor(mm_2, mm_1);  mm_2 = mm_1 = None
        return pytree.tree_unflatten([sin, add], self._out_spec)


INFO: aot_config id: 0, fw_metadata=ViewAndMutationMeta(input_info=[InputAliasInfo(is_leaf=True, mutates_data=False, mutates_metadata=False, mutations_hidden_from_autograd=True, mutations_under_no_grad_or_inference_mode=False, mutation_inductor_storage_resize=False, mutates_storage_metadata=False, requires_grad=True, keep_input_mutations=True)], output_info=[OutputAliasInfo(output_type=<OutputType.non_alias: 1>, raw_type=<class 'torch._subclasses.functional_tensor.FunctionalTensor'>, base_idx=None, dynamic_dims=set(), requires_grad=True, functional_tensor=None)], num_intermediate_bases=0, keep_input_mutations=True, traced_tangents=[FakeTensor(..., size=(4, 4))], subclass_inp_meta=[0], subclass_fw_graph_out_meta=[0], subclass_tangent_meta=[0], is_train=True, traced_tangent_metas=None, num_symints_saved_for_bw=0, grad_enabled_mutation=None, deterministic=False, static_input_indices=[], tokens={}, indices_of_inputs_that_requires_grad_with_mutations_in_bw=[], bw_donated_idxs=None), inner_meta=ViewAndMutationMeta(input_info=[InputAliasInfo(is_leaf=True, mutates_data=False, mutates_metadata=False, mutations_hidden_from_autograd=True, mutations_under_no_grad_or_inference_mode=False, mutation_inductor_storage_resize=False, mutates_storage_metadata=False, requires_grad=True, keep_input_mutations=True)], output_info=[OutputAliasInfo(output_type=<OutputType.non_alias: 1>, raw_type=<class 'torch._subclasses.functional_tensor.FunctionalTensor'>, base_idx=None, dynamic_dims=set(), requires_grad=True, functional_tensor=None)], num_intermediate_bases=0, keep_input_mutations=True, traced_tangents=[FakeTensor(..., size=(4, 4))], subclass_inp_meta=[0], subclass_fw_graph_out_meta=[0], subclass_tangent_meta=[0], is_train=True, traced_tangent_metas=None, num_symints_saved_for_bw=0, grad_enabled_mutation=None, deterministic=False, static_input_indices=[], tokens={}, indices_of_inputs_that_requires_grad_with_mutations_in_bw=[], bw_donated_idxs=None)
INFO: TRACED GRAPH
 ===== Forward graph 0 =====
 /home/hirsheybar/local/b/pytorch/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, primals_1: "f32[4, 4][4, 1]cpu"):
         # File: /home/hirsheybar/local/b/pytorch/tmp5.py:6 in f, code: return torch.matmul(x, x).sin()
        mm: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(primals_1, primals_1)
        sin: "f32[4, 4][4, 1]cpu" = torch.ops.aten.sin.default(mm)
        permute: "f32[4, 4][1, 4]cpu" = torch.ops.aten.permute.default(primals_1, [1, 0]);  primals_1 = None
        return (sin, mm, permute)


INFO: TRACED GRAPH
 ===== Backward graph 0 =====
 <eval_with_key>.1 class GraphModule(torch.nn.Module):
    def forward(self, mm: "f32[4, 4][4, 1]cpu", permute: "f32[4, 4][1, 4]cpu", tangents_1: "f32[4, 4][4, 1]cpu"):
         # File: /home/hirsheybar/local/b/pytorch/tmp5.py:6 in f, code: return torch.matmul(x, x).sin()
        cos: "f32[4, 4][4, 1]cpu" = torch.ops.aten.cos.default(mm);  mm = None
        mul: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mul.Tensor(tangents_1, cos);  tangents_1 = cos = None
        mm_1: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(permute, mul)
        mm_2: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(mul, permute);  mul = permute = None

         # File: /home/hirsheybar/local/b/pytorch/tmp5.py:6 in f, code: return torch.matmul(x, x).sin()
        add: "f32[4, 4][4, 1]cpu" = torch.ops.aten.add.Tensor(mm_2, mm_1);  mm_2 = mm_1 = None
        return (add,)


INFO: TRACED GRAPH
 ===== Joint graph 1 =====
 /home/hirsheybar/local/b/pytorch/torch/fx/_lazy_graph_module.py class joint_helper(torch.nn.Module):
    def forward(self, primals, tangents):
        primals_1: "f32[4, 4][4, 1]cpu"; tangents_1: "f32[4, 4][4, 1]cpu";

        primals_1, tangents_1, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec)
         # File: /home/hirsheybar/local/b/pytorch/tmp5.py:6 in f, code: return torch.matmul(x, x).sin()
        mm: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(primals_1, primals_1)
        sin: "f32[4, 4][4, 1]cpu" = torch.ops.aten.sin.default(mm)
        cos: "f32[4, 4][4, 1]cpu" = torch.ops.aten.cos.default(mm);  mm = None
        mul: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mul.Tensor(tangents_1, cos);  tangents_1 = cos = None
        permute: "f32[4, 4][1, 4]cpu" = torch.ops.aten.permute.default(primals_1, [1, 0])
        mm_1: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(permute, mul);  permute = None
        permute_1: "f32[4, 4][1, 4]cpu" = torch.ops.aten.permute.default(primals_1, [1, 0]);  primals_1 = None
        mm_2: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(mul, permute_1);  mul = permute_1 = None

         # File: /home/hirsheybar/local/b/pytorch/tmp5.py:6 in f, code: return torch.matmul(x, x).sin()
        add: "f32[4, 4][4, 1]cpu" = torch.ops.aten.add.Tensor(mm_2, mm_1);  mm_2 = mm_1 = None
        return pytree.tree_unflatten([sin, add], self._out_spec)


INFO: aot_config id: 1, fw_metadata=ViewAndMutationMeta(input_info=[InputAliasInfo(is_leaf=True, mutates_data=False, mutates_metadata=False, mutations_hidden_from_autograd=True, mutations_under_no_grad_or_inference_mode=False, mutation_inductor_storage_resize=False, mutates_storage_metadata=False, requires_grad=True, keep_input_mutations=True)], output_info=[OutputAliasInfo(output_type=<OutputType.non_alias: 1>, raw_type=<class 'torch._subclasses.functional_tensor.FunctionalTensor'>, base_idx=None, dynamic_dims=set(), requires_grad=True, functional_tensor=None)], num_intermediate_bases=0, keep_input_mutations=True, traced_tangents=[FakeTensor(..., size=(4, 4))], subclass_inp_meta=[0], subclass_fw_graph_out_meta=[0], subclass_tangent_meta=[0], is_train=True, traced_tangent_metas=None, num_symints_saved_for_bw=0, grad_enabled_mutation=None, deterministic=False, static_input_indices=[], tokens={}, indices_of_inputs_that_requires_grad_with_mutations_in_bw=[], bw_donated_idxs=None), inner_meta=ViewAndMutationMeta(input_info=[InputAliasInfo(is_leaf=True, mutates_data=False, mutates_metadata=False, mutations_hidden_from_autograd=True, mutations_under_no_grad_or_inference_mode=False, mutation_inductor_storage_resize=False, mutates_storage_metadata=False, requires_grad=True, keep_input_mutations=True)], output_info=[OutputAliasInfo(output_type=<OutputType.non_alias: 1>, raw_type=<class 'torch._subclasses.functional_tensor.FunctionalTensor'>, base_idx=None, dynamic_dims=set(), requires_grad=True, functional_tensor=None)], num_intermediate_bases=0, keep_input_mutations=True, traced_tangents=[FakeTensor(..., size=(4, 4))], subclass_inp_meta=[0], subclass_fw_graph_out_meta=[0], subclass_tangent_meta=[0], is_train=True, traced_tangent_metas=None, num_symints_saved_for_bw=0, grad_enabled_mutation=None, deterministic=False, static_input_indices=[], tokens={}, indices_of_inputs_that_requires_grad_with_mutations_in_bw=[], bw_donated_idxs=None)
INFO: TRACED GRAPH
 ===== Forward graph 1 =====
 /home/hirsheybar/local/b/pytorch/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, primals_1: "f32[4, 4][4, 1]cpu"):
         # File: /home/hirsheybar/local/b/pytorch/tmp5.py:6 in f, code: return torch.matmul(x, x).sin()
        mm: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(primals_1, primals_1)
        sin: "f32[4, 4][4, 1]cpu" = torch.ops.aten.sin.default(mm)
        permute: "f32[4, 4][1, 4]cpu" = torch.ops.aten.permute.default(primals_1, [1, 0]);  primals_1 = None
        return (sin, mm, permute)


INFO: TRACED GRAPH
 ===== Backward graph 1 =====
 <eval_with_key>.5 class GraphModule(torch.nn.Module):
    def forward(self, mm: "f32[4, 4][4, 1]cpu", permute: "f32[4, 4][1, 4]cpu", tangents_1: "f32[4, 4][4, 1]cpu"):
         # File: /home/hirsheybar/local/b/pytorch/tmp5.py:6 in f, code: return torch.matmul(x, x).sin()
        cos: "f32[4, 4][4, 1]cpu" = torch.ops.aten.cos.default(mm);  mm = None
        mul: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mul.Tensor(tangents_1, cos);  tangents_1 = cos = None
        mm_1: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(permute, mul)
        mm_2: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(mul, permute);  mul = permute = None

         # File: /home/hirsheybar/local/b/pytorch/tmp5.py:6 in f, code: return torch.matmul(x, x).sin()
        add: "f32[4, 4][4, 1]cpu" = torch.ops.aten.add.Tensor(mm_2, mm_1);  mm_2 = mm_1 = None
        return (add,)


DEBUG: Cache miss due to new autograd node: torch::autograd::GraphRoot (NodeCall 0) with key size 39, previous key sizes=[]
DEBUG: TRACED GRAPH
 ===== Compiled autograd graph =====
 <eval_with_key>.8 class CompiledAutograd(torch.nn.Module):
    def forward(self, inputs, sizes, scalars, hooks):
        # No stacktrace found for following nodes
        getitem: "f32[]cpu" = inputs[0]
        getitem_1: "f32[4, 4]cpu" = inputs[1]
        getitem_2: "f32[4, 4]cpu" = inputs[2]
        getitem_3: "f32[4, 4]cpu" = inputs[3];  inputs = None

         # File: /home/hirsheybar/local/b/pytorch/torch/_dynamo/compiled_autograd.py:379 in set_node_origin, code: SumBackward0 (NodeCall 1)
        expand: "f32[4, 4]cpu" = torch.ops.aten.expand.default(getitem, [4, 4]);  getitem = None

         # File: /home/hirsheybar/local/b/pytorch/torch/_dynamo/compiled_autograd.py:379 in set_node_origin, code: CompiledFunctionBackward (NodeCall 2)
        clone: "f32[4, 4]cpu" = torch.ops.aten.clone.default(expand, memory_format = torch.contiguous_format);  expand = None
        cos: "f32[4, 4]cpu" = torch.ops.aten.cos.default(getitem_1);  getitem_1 = None
        mul: "f32[4, 4]cpu" = torch.ops.aten.mul.Tensor(clone, cos);  clone = cos = None
        mm: "f32[4, 4]cpu" = torch.ops.aten.mm.default(getitem_2, mul)
        mm_1: "f32[4, 4]cpu" = torch.ops.aten.mm.default(mul, getitem_2);  mul = getitem_2 = None
        add: "f32[4, 4]cpu" = torch.ops.aten.add.Tensor(mm_1, mm);  mm_1 = mm = None

         # File: /home/hirsheybar/local/b/pytorch/torch/_dynamo/compiled_autograd.py:379 in set_node_origin, code: torch::autograd::AccumulateGrad (NodeCall 3)
        accumulate_grad_ = torch.ops.inductor.accumulate_grad_.default(getitem_3, add);  getitem_3 = add = accumulate_grad_ = None
        _exec_final_callbacks_stub = torch__dynamo_external_utils__exec_final_callbacks_stub();  _exec_final_callbacks_stub = None
        return []


INFO: TRACED GRAPH
 ===== Forward graph 2 =====
 /home/hirsheybar/local/b/pytorch/torch/fx/_lazy_graph_module.py class <lambda>(torch.nn.Module):
    def forward(self, arg0_1: "f32[][]cpu", arg1_1: "f32[4, 4][4, 1]cpu", arg2_1: "f32[4, 4][1, 4]cpu", arg3_1: "f32[4, 4][4, 1]cpu"):
        # No stacktrace found for following nodes
        expand: "f32[4, 4][0, 0]cpu" = torch.ops.aten.expand.default(arg0_1, [4, 4]);  arg0_1 = None
        clone: "f32[4, 4][4, 1]cpu" = torch.ops.aten.clone.default(expand, memory_format = torch.contiguous_format);  expand = None
        cos: "f32[4, 4][4, 1]cpu" = torch.ops.aten.cos.default(arg1_1);  arg1_1 = None
        mul: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mul.Tensor(clone, cos);  clone = cos = None
        mm: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(arg2_1, mul)
        mm_1: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(mul, arg2_1);  mul = arg2_1 = None
        add: "f32[4, 4][4, 1]cpu" = torch.ops.aten.add.Tensor(mm_1, mm);  mm_1 = mm = None
        clone_1: "f32[4, 4][4, 1]cpu" = torch.ops.aten.clone.default(add);  add = None
        return (clone_1,)


INFO: TRACED GRAPH
 ===== Forward graph 3 =====
 /home/hirsheybar/local/b/pytorch/torch/fx/_lazy_graph_module.py class <lambda>(torch.nn.Module):
    def forward(self, arg0_1: "f32[][]cpu", arg1_1: "f32[4, 4][4, 1]cpu", arg2_1: "f32[4, 4][1, 4]cpu", arg3_1: "f32[4, 4][4, 1]cpu"):
        # No stacktrace found for following nodes
        expand: "f32[4, 4][0, 0]cpu" = torch.ops.aten.expand.default(arg0_1, [4, 4]);  arg0_1 = None
        clone: "f32[4, 4][4, 1]cpu" = torch.ops.aten.clone.default(expand, memory_format = torch.contiguous_format);  expand = None
        cos: "f32[4, 4][4, 1]cpu" = torch.ops.aten.cos.default(arg1_1);  arg1_1 = None
        mul: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mul.Tensor(clone, cos);  clone = cos = None
        mm: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(arg2_1, mul)
        mm_1: "f32[4, 4][4, 1]cpu" = torch.ops.aten.mm.default(mul, arg2_1);  mul = arg2_1 = None
        add: "f32[4, 4][4, 1]cpu" = torch.ops.aten.add.Tensor(mm_1, mm);  mm_1 = mm = None
        clone_1: "f32[4, 4][4, 1]cpu" = torch.ops.aten.clone.default(add);  add = None
        return (clone_1,)

```

cc @ezyang @chauhang @penguinwu @zou3519