…compute during copy_ (DTensor) (pytorch#170467)

Before this PR, graph capturing a program that performed an input mutation on a DTensor under training would hard error.

The context:

(1) When AOTAutograd traces out a joint graph and detects an input mutation, it makes a call to `old_input.copy_(new_input)`, and relies on the current make_fx call to capture this copy_ as a node in the joint graph ([code](https://github.com/pytorch/pytorch/blob/main/torch/_functorch/_aot_autograd/graph_capture_wrappers.py#L733))

(2) ordinarily, capturing `old_input.copy_(new_input)` doesn't **need** to generate any fresh proxies in the graph, as the output of `copy_` is the self argument, which we expect to already have a proxy. Why does this matter? @IvanKobzarev added some logic to handle the case where a buffer is mutated during both the fw and the bw ([PR](pytorch#155354)), and as part of doing so, tweaked the input mutation handling in AOTAutograd so that these copy_ calls are generated under a [context manager](https://github.com/pytorch/pytorch/blob/main/torch/_functorch/_aot_autograd/graph_capture_wrappers.py#L979C29-L979C72) that prevents proxy_tensor from adding new proxies to the graph. The idea being that we are applying this context manager in a very limited region, where we know no new proxies need to be created

(3) However, this is not true for DTensor. When you call `dtensor.copy_(dtensor)`, DTensor runs fake tensor prop under the hood, which involves constructing fresh FakeTensors for the inputs with which to run the fake prop on ([code](https://github.com/pytorch/pytorch/blob/main/torch/distributed/tensor/_op_schema.py#L510))

The net result is that we end up *not* constructing proxies for these fake tensor inputs, and we get a "proxy not found" error immediately afterwards when attempting to use them when DTensor runs fake prop ([here](https://github.com/pytorch/pytorch/blob/main/torch/distributed/tensor/_sharding_prop.py#L243))

The way I fixed this was just by tweaking the "don't clobber proxies" context manager to be a bit more general: it will still generate proxies for inputs that don't already have proxies, and it simply won't overwrite an existing proxy with a new one when you trace an inplace op.

One alternative would have been to disable proxy tracing when DTensor runs fake prop. Since after all, we don't really care about the ops that DTensor ran during fake prop. I decided not to do this because that code has changed a bunch recently and is pretty fragile, but I'm hoping to it if people prefer that path.

Pull Request resolved: pytorch#170467
Approved by: https://github.com/IvanKobzarev

…compute during copy_ (DTensor) (pytorch#170467)

Before this PR, graph capturing a program that performed an input mutation on a DTensor under training would hard error.

The context:

(1) When AOTAutograd traces out a joint graph and detects an input mutation, it makes a call to `old_input.copy_(new_input)`, and relies on the current make_fx call to capture this copy_ as a node in the joint graph ([code](https://github.com/pytorch/pytorch/blob/main/torch/_functorch/_aot_autograd/graph_capture_wrappers.py#L733))

(2) ordinarily, capturing `old_input.copy_(new_input)` doesn't **need** to generate any fresh proxies in the graph, as the output of `copy_` is the self argument, which we expect to already have a proxy. Why does this matter? @IvanKobzarev added some logic to handle the case where a buffer is mutated during both the fw and the bw ([PR](pytorch#155354)), and as part of doing so, tweaked the input mutation handling in AOTAutograd so that these copy_ calls are generated under a [context manager](https://github.com/pytorch/pytorch/blob/main/torch/_functorch/_aot_autograd/graph_capture_wrappers.py#L979C29-L979C72) that prevents proxy_tensor from adding new proxies to the graph. The idea being that we are applying this context manager in a very limited region, where we know no new proxies need to be created

(3) However, this is not true for DTensor. When you call `dtensor.copy_(dtensor)`, DTensor runs fake tensor prop under the hood, which involves constructing fresh FakeTensors for the inputs with which to run the fake prop on ([code](https://github.com/pytorch/pytorch/blob/main/torch/distributed/tensor/_op_schema.py#L510))

The net result is that we end up *not* constructing proxies for these fake tensor inputs, and we get a "proxy not found" error immediately afterwards when attempting to use them when DTensor runs fake prop ([here](https://github.com/pytorch/pytorch/blob/main/torch/distributed/tensor/_sharding_prop.py#L243))

The way I fixed this was just by tweaking the "don't clobber proxies" context manager to be a bit more general: it will still generate proxies for inputs that don't already have proxies, and it simply won't overwrite an existing proxy with a new one when you trace an inplace op.

One alternative would have been to disable proxy tracing when DTensor runs fake prop. Since after all, we don't really care about the ops that DTensor ran during fake prop. I decided not to do this because that code has changed a bunch recently and is pretty fragile, but I'm hoping to it if people prefer that path.

Pull Request resolved: pytorch#170467
Approved by: https://github.com/IvanKobzarev