Repository: pytorch
Issue: [torch.compile] Mutating backward input of autograd function is not supported by ```torch.compile```
Labels: oncall: distributed,triaged,oncall: pt2,module: dynamo,module: pt2-dispatcher,dynamo-autograd-function

### 🐛 Describe the bug

This issue was raised from tracing Megatron/xlformers, where ```torch.distributed.all_reduce``` was called in backward of an ```autograd.Function``` and then it was rewritten by Dynamo into [```torch.distributed._functional_collectives.all_reduce_inplace```](https://github.com/pytorch/pytorch/blob/93e3c91679cbb36bd351d4caa8fdec2fe7388947/torch/distributed/_functional_collectives.py#L1100) which has a inplace ```copy_``` op.

We can use the following simple repro to reproduce this error:
```
import torch
torch.set_default_device('cuda')

class Foo(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return x.clone(), x.clone()

    @staticmethod
    def backward(ctx, grad1, grad2):
        return grad1.copy_(grad2)

@torch.compile(backend="aot_eager", fullgraph=True)
def f(x):
    return Foo.apply(x)

x = torch.randn(3, requires_grad=True)
result = f(x)
print(result)
```
Error stack:
```
Traceback (most recent call last):
  File "/data/users/ybliang/debug/debug7.py", line 19, in <module>
    result = f(x)
  File "/home/ybliang/local/pytorch/torch/_dynamo/eval_frame.py", line 556, in _fn
    return fn(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_dynamo/convert_frame.py", line 1404, in __call__
    return self._torchdynamo_orig_callable(
  File "/home/ybliang/local/pytorch/torch/_dynamo/convert_frame.py", line 549, in __call__
    return _compile(
  File "/home/ybliang/local/pytorch/torch/_dynamo/convert_frame.py", line 985, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/home/ybliang/local/pytorch/torch/_dynamo/convert_frame.py", line 712, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
  File "/home/ybliang/local/pytorch/torch/_utils_internal.py", line 95, in wrapper_function
    return function(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_dynamo/convert_frame.py", line 747, in _compile_inner
    out_code = transform_code_object(code, transform)
  File "/home/ybliang/local/pytorch/torch/_dynamo/bytecode_transformation.py", line 1348, in transform_code_object
    transformations(instructions, code_options)
  File "/home/ybliang/local/pytorch/torch/_dynamo/convert_frame.py", line 233, in _fn
    return fn(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_dynamo/convert_frame.py", line 664, in transform
    tracer.run()
  File "/home/ybliang/local/pytorch/torch/_dynamo/symbolic_convert.py", line 2841, in run
    super().run()
  File "/home/ybliang/local/pytorch/torch/_dynamo/symbolic_convert.py", line 1032, in run
    while self.step():
  File "/home/ybliang/local/pytorch/torch/_dynamo/symbolic_convert.py", line 944, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/ybliang/local/pytorch/torch/_dynamo/symbolic_convert.py", line 3021, in RETURN_VALUE
    self._return(inst)
  File "/home/ybliang/local/pytorch/torch/_dynamo/symbolic_convert.py", line 3006, in _return
    self.output.compile_subgraph(
  File "/home/ybliang/local/pytorch/torch/_dynamo/output_graph.py", line 1110, in compile_subgraph
    self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)
  File "/home/ybliang/local/pytorch/torch/_dynamo/output_graph.py", line 1349, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/home/ybliang/local/pytorch/torch/_dynamo/output_graph.py", line 1399, in call_user_compiler
    return self._call_user_compiler(gm)
  File "/home/ybliang/local/pytorch/torch/_dynamo/output_graph.py", line 1448, in _call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/home/ybliang/local/pytorch/torch/_dynamo/output_graph.py", line 1429, in _call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
  File "/home/ybliang/local/pytorch/torch/_dynamo/repro/after_dynamo.py", line 130, in __call__
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/home/ybliang/local/pytorch/torch/_dynamo/repro/after_dynamo.py", line 130, in __call__
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/home/ybliang/local/pytorch/torch/__init__.py", line 2346, in __call__
    return self.compiler_fn(model_, inputs_, **self.kwargs)
  File "/home/ybliang/local/pytorch/torch/_dynamo/backends/debugging.py", line 153, in aot_eager
    return aot_autograd(
  File "/home/ybliang/local/pytorch/torch/_dynamo/backends/common.py", line 72, in __call__
    cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
  File "/home/ybliang/local/pytorch/torch/_functorch/aot_autograd.py", line 1103, in aot_module_simplified
    compiled_fn = dispatch_and_compile()
  File "/home/ybliang/local/pytorch/torch/_functorch/aot_autograd.py", line 1079, in dispatch_and_compile
    compiled_fn, _ = create_aot_dispatcher_function(
  File "/home/ybliang/local/pytorch/torch/_functorch/aot_autograd.py", line 527, in create_aot_dispatcher_function
    return _create_aot_dispatcher_function(
  File "/home/ybliang/local/pytorch/torch/_functorch/aot_autograd.py", line 778, in _create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
  File "/home/ybliang/local/pytorch/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 373, in aot_dispatch_autograd
    fx_g, joint_inputs, maybe_subclass_meta = aot_dispatch_autograd_graph(
  File "/home/ybliang/local/pytorch/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 318, in aot_dispatch_autograd_graph
    fx_g = _create_graph(joint_fn_to_trace, updated_joint_inputs, aot_config=aot_config)
  File "/home/ybliang/local/pytorch/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 55, in _create_graph
    fx_g = make_fx(
  File "/home/ybliang/local/pytorch/torch/fx/experimental/proxy_tensor.py", line 2188, in wrapped
    return make_fx_tracer.trace(f, *args)
  File "/home/ybliang/local/pytorch/torch/fx/experimental/proxy_tensor.py", line 2126, in trace
    return self._trace_inner(f, *args)
  File "/home/ybliang/local/pytorch/torch/fx/experimental/proxy_tensor.py", line 2097, in _trace_inner
    t = dispatch_trace(
  File "/home/ybliang/local/pytorch/torch/_compile.py", line 32, in inner
    return disable_fn(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_dynamo/eval_frame.py", line 721, in _fn
    return fn(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/fx/experimental/proxy_tensor.py", line 1137, in dispatch_trace
    graph = tracer.trace(root, concrete_args)  # type: ignore[arg-type]
  File "/home/ybliang/local/pytorch/torch/_dynamo/eval_frame.py", line 721, in _fn
    return fn(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/fx/_symbolic_trace.py", line 843, in trace
    (self.create_arg(fn(*args)),),
  File "/home/ybliang/local/pytorch/torch/fx/_symbolic_trace.py", line 700, in flatten_fn
    tree_out = root_fn(*tree_args)
  File "/home/ybliang/local/pytorch/torch/fx/experimental/proxy_tensor.py", line 1192, in wrapped
    out = f(*tensors)
  File "/home/ybliang/local/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 693, in inner_fn
    outs = fn(*args)
  File "/home/ybliang/local/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 644, in joint_helper
    return _functionalized_f_helper(primals, tangents)
  File "/home/ybliang/local/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 480, in _functionalized_f_helper
    assert not has_metadata_mutation(
torch._dynamo.exc.BackendCompilerFailed: backend='aot_eager' raised:
AssertionError: Found an input to the backward that was mutated during the backward pass. This is not supported
```

### Versions

N/A

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @zou3519 @bdhirsh @yf225