[aotd] Support mutations of the same input in fw and bw by IvanKobzarev · Pull Request #155354 · pytorch/pytorch

IvanKobzarev · 2025-06-06T18:36:13Z

Stack from ghstack (oldest at bottom):

-> [aotd] Support mutations of the same input in fw and bw #155354

Original issue: #154820

The issue happens when there is a mutation for the same input in forward AND in backward.

AOTD emited copy_ after joint_function tracing. This made this fx-node to correspond to the side effects of both mutations (in forward and in backward).
After that partitioner can put it either in forward or in backward.

The fix:

1/ Introduce joint_function.handle that allows to set "post_forward" callback, to be able to check inputs state after forward

We do not want to apply the mutation after joint, if we already applied it in forward. For that we need "mutation_counter" and memorize the version of mutation that we applied for forward mutation.

2/ Exposing mutation_counter to python

We want to keep invariant that copy_ exist only in the end of joint graph.

3/ We memorize mutation_counter and state of the inputs after forward, using the handle post_forward.
Emit post_forward mutations after joint graph fully traced.

add for post_forward mutations "must_be_in_forward" tag (similar to existing "must_be_in_backward") to keep them in forward.

4/ Ban recompute of the source of mutation. Recompute can apply the same op (e.g. add) in forward and backward.
For this set MUST_SAVE for the source of mutation in forward.

proxy_tensor changes:

By default proxy tensor updates tensor_tracker. In this case applied mutations will be chained.
But we want that this copy_ will be independent and applied just to primals.
For this introducing a contextmanager to be able to disable update of tensor_tracker for adding forward mutations.

cc @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv @voznesenskym @penguinwu @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @jiayisunx @chenyang78 @kadeng @chauhang @amjames

[ghstack-poisoned]

pytorch-bot · 2025-06-06T18:36:17Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/155354

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 1 Pending

As of commit a60f113 with merge base bbf1a6f ():
💚 Looks good so far! There are no failures yet. 💚

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

⏳ inductor-rocm / rocm-py3.10-inductor / test (inductor, 1, 2, linux.rocm.gpu.2, unstable) (gh) (#155917)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

torch/fx/experimental/proxy_tensor.py

Original issue: #154820 The issue happens when there is a mutation for the same input in forward AND in backward. AOTD emited copy_ after joint_function tracing. This made this fx-node to correspond to the side effects of both mutations (in forward and in backward). After that partitioner can put it either in forward or in backward. The fix: 1/ Introduce joint_function.handle that allows to set "post_forward" callback, to be able to check inputs state after forward We do not want to apply the mutation after joint, if we already applied it in forward. For that we need "mutation_counter" and memorize the version of mutation that we applied for forward mutation. 2/ Exposing mutation_counter to python We want to keep invariant that copy_ exist only in the end of joint graph. 3/ We memorize mutation_counter and state of the inputs after forward, using the handle post_forward. Emit post_forward mutations after joint graph fully traced. add for post_forward mutations "must_be_in_forward" tag (similar to existing "must_be_in_backward") to keep them in forward. 4/ Ban recompute of the source of mutation. Recompute can apply the same op (e.g. add) in forward and backward. For this set MUST_SAVE for the source of mutation in forward. proxy_tensor changes: By default proxy tensor updates tensor_tracker. In this case applied mutations will be chained. But we want that this copy_ will be independent and applied just to primals. For this introducing a contextmanager to be able to disable update of tensor_tracker for adding forward mutations. cc ezyang SherlockNoMad EikanWang jgong5 wenzhe-nrv [ghstack-poisoned]

ghstack-source-id: 95a4053 Pull Request resolved: #155354

Original issue: #154820 The issue happens when there is a mutation for the same input in forward AND in backward. AOTD emited copy_ after joint_function tracing. This made this fx-node to correspond to the side effects of both mutations (in forward and in backward). After that partitioner can put it either in forward or in backward. The fix: 1/ Introduce joint_function.handle that allows to set "post_forward" callback, to be able to check inputs state after forward We do not want to apply the mutation after joint, if we already applied it in forward. For that we need "mutation_counter" and memorize the version of mutation that we applied for forward mutation. 2/ Exposing mutation_counter to python We want to keep invariant that copy_ exist only in the end of joint graph. 3/ We memorize mutation_counter and state of the inputs after forward, using the handle post_forward. Emit post_forward mutations after joint graph fully traced. add for post_forward mutations "must_be_in_forward" tag (similar to existing "must_be_in_backward") to keep them in forward. 4/ Ban recompute of the source of mutation. Recompute can apply the same op (e.g. add) in forward and backward. For this set MUST_SAVE for the source of mutation in forward. proxy_tensor changes: By default proxy tensor updates tensor_tracker. In this case applied mutations will be chained. But we want that this copy_ will be independent and applied just to primals. For this introducing a contextmanager to be able to disable update of tensor_tracker for adding forward mutations. cc ezyang SherlockNoMad EikanWang jgong5 wenzhe-nrv [ghstack-poisoned]

ghstack-source-id: 9fb4349 Pull Request resolved: #155354

Original issue: #154820 The issue happens when there is a mutation for the same input in forward AND in backward. AOTD emited copy_ after joint_function tracing. This made this fx-node to correspond to the side effects of both mutations (in forward and in backward). After that partitioner can put it either in forward or in backward. The fix: 1/ Introduce joint_function.handle that allows to set "post_forward" callback, to be able to check inputs state after forward We do not want to apply the mutation after joint, if we already applied it in forward. For that we need "mutation_counter" and memorize the version of mutation that we applied for forward mutation. 2/ Exposing mutation_counter to python We want to keep invariant that copy_ exist only in the end of joint graph. 3/ We memorize mutation_counter and state of the inputs after forward, using the handle post_forward. Emit post_forward mutations after joint graph fully traced. add for post_forward mutations "must_be_in_forward" tag (similar to existing "must_be_in_backward") to keep them in forward. 4/ Ban recompute of the source of mutation. Recompute can apply the same op (e.g. add) in forward and backward. For this set MUST_SAVE for the source of mutation in forward. proxy_tensor changes: By default proxy tensor updates tensor_tracker. In this case applied mutations will be chained. But we want that this copy_ will be independent and applied just to primals. For this introducing a contextmanager to be able to disable update of tensor_tracker for adding forward mutations. cc ezyang SherlockNoMad EikanWang jgong5 wenzhe-nrv [ghstack-poisoned]

ghstack-source-id: 73ebf15 Pull Request resolved: #155354

Original issue: #154820 The issue happens when there is a mutation for the same input in forward AND in backward. AOTD emited copy_ after joint_function tracing. This made this fx-node to correspond to the side effects of both mutations (in forward and in backward). After that partitioner can put it either in forward or in backward. The fix: 1/ Introduce joint_function.handle that allows to set "post_forward" callback, to be able to check inputs state after forward We do not want to apply the mutation after joint, if we already applied it in forward. For that we need "mutation_counter" and memorize the version of mutation that we applied for forward mutation. 2/ Exposing mutation_counter to python We want to keep invariant that copy_ exist only in the end of joint graph. 3/ We memorize mutation_counter and state of the inputs after forward, using the handle post_forward. Emit post_forward mutations after joint graph fully traced. add for post_forward mutations "must_be_in_forward" tag (similar to existing "must_be_in_backward") to keep them in forward. 4/ Ban recompute of the source of mutation. Recompute can apply the same op (e.g. add) in forward and backward. For this set MUST_SAVE for the source of mutation in forward. proxy_tensor changes: By default proxy tensor updates tensor_tracker. In this case applied mutations will be chained. But we want that this copy_ will be independent and applied just to primals. For this introducing a contextmanager to be able to disable update of tensor_tracker for adding forward mutations. cc ezyang SherlockNoMad EikanWang jgong5 wenzhe-nrv voznesenskym penguinwu Guobing-Chen XiaobingSuper zhuhaozhe blzheng jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

ghstack-source-id: 4db1990 Pull Request resolved: #155354

IvanKobzarev · 2025-06-25T19:18:13Z

@pytorchbot merge

pytorchmergebot · 2025-06-25T19:20:00Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2025-06-25T19:25:37Z

Merge failed

Reason: Command git -C /home/runner/work/pytorch/pytorch cherry-pick -x 8a39bd269e67e74d9d8b3c6ce73c7b41442f0747 returned non-zero exit code 1

Auto-merging benchmarks/dynamo/pr_time_benchmarks/expected_results.csv
CONFLICT (content): Merge conflict in benchmarks/dynamo/pr_time_benchmarks/expected_results.csv
Auto-merging torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py
Auto-merging torch/_functorch/_aot_autograd/traced_function_transforms.py
CONFLICT (content): Merge conflict in torch/_functorch/_aot_autograd/traced_function_transforms.py
Auto-merging torch/_functorch/partitioners.py
error: could not apply 8a39bd269e6... [aotd] Support mutations of the same input in fw and bw
hint: After resolving the conflicts, mark them with
hint: "git add/rm <pathspec>", then run
hint: "git cherry-pick --continue".
hint: You can instead skip this commit with "git cherry-pick --skip".
hint: To abort and get back to the state before "git cherry-pick",
hint: run "git cherry-pick --abort".
hint: Disable this message with "git config set advice.mergeConflict false"

Details for Dev Infra team

Raised by workflow job

Original issue: #154820 The issue happens when there is a mutation for the same input in forward AND in backward. AOTD emited copy_ after joint_function tracing. This made this fx-node to correspond to the side effects of both mutations (in forward and in backward). After that partitioner can put it either in forward or in backward. The fix: 1/ Introduce joint_function.handle that allows to set "post_forward" callback, to be able to check inputs state after forward We do not want to apply the mutation after joint, if we already applied it in forward. For that we need "mutation_counter" and memorize the version of mutation that we applied for forward mutation. 2/ Exposing mutation_counter to python We want to keep invariant that copy_ exist only in the end of joint graph. 3/ We memorize mutation_counter and state of the inputs after forward, using the handle post_forward. Emit post_forward mutations after joint graph fully traced. add for post_forward mutations "must_be_in_forward" tag (similar to existing "must_be_in_backward") to keep them in forward. 4/ Ban recompute of the source of mutation. Recompute can apply the same op (e.g. add) in forward and backward. For this set MUST_SAVE for the source of mutation in forward. proxy_tensor changes: By default proxy tensor updates tensor_tracker. In this case applied mutations will be chained. But we want that this copy_ will be independent and applied just to primals. For this introducing a contextmanager to be able to disable update of tensor_tracker for adding forward mutations. cc ezyang SherlockNoMad EikanWang jgong5 wenzhe-nrv voznesenskym penguinwu Guobing-Chen XiaobingSuper zhuhaozhe blzheng jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

ghstack-source-id: 485f601 Pull Request resolved: #155354

Original issue: #154820 The issue happens when there is a mutation for the same input in forward AND in backward. AOTD emited copy_ after joint_function tracing. This made this fx-node to correspond to the side effects of both mutations (in forward and in backward). After that partitioner can put it either in forward or in backward. The fix: 1/ Introduce joint_function.handle that allows to set "post_forward" callback, to be able to check inputs state after forward We do not want to apply the mutation after joint, if we already applied it in forward. For that we need "mutation_counter" and memorize the version of mutation that we applied for forward mutation. 2/ Exposing mutation_counter to python We want to keep invariant that copy_ exist only in the end of joint graph. 3/ We memorize mutation_counter and state of the inputs after forward, using the handle post_forward. Emit post_forward mutations after joint graph fully traced. add for post_forward mutations "must_be_in_forward" tag (similar to existing "must_be_in_backward") to keep them in forward. 4/ Ban recompute of the source of mutation. Recompute can apply the same op (e.g. add) in forward and backward. For this set MUST_SAVE for the source of mutation in forward. proxy_tensor changes: By default proxy tensor updates tensor_tracker. In this case applied mutations will be chained. But we want that this copy_ will be independent and applied just to primals. For this introducing a contextmanager to be able to disable update of tensor_tracker for adding forward mutations. cc ezyang SherlockNoMad EikanWang jgong5 wenzhe-nrv voznesenskym penguinwu Guobing-Chen XiaobingSuper zhuhaozhe blzheng jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

ghstack-source-id: 9d633c6 Pull Request resolved: #155354

IvanKobzarev · 2025-06-26T13:58:29Z

@pytorchbot merge

pytorchmergebot · 2025-06-26T14:00:13Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

This PR adds a new config `backward_pass_autocast`, to set the backward autocast behavior. It does not change the existing behavior. The reason why we need this is that torch.compile acquires a forward and backward graph at the time of the forward pass. This means that implemented naively, if there are any context managers active outside the call to torch.compile, the backward graph will also get the behaviors from those context managers. This PR gives users a way to tweak the autocast behavior of the backward pass. Please see torch._functorch.config for the options to the `backward_pass_autocast` config. Pull Request resolved: #156356 Approved by: https://github.com/bdhirsh ghstack dependencies: #155354

… that perform intermediate compute during copy_ (DTensor)" 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](#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. cc ezyang EikanWang jgong5 wenzhe-nrv [ghstack-poisoned]

…termediate compute during copy_ (DTensor)" 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](#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. cc ezyang EikanWang jgong5 wenzhe-nrv [ghstack-poisoned]

@IvanKobzarev

…compute during copy_ (DTensor) (#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](#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: #170467 Approved by: https://github.com/IvanKobzarev

@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

@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

[aotd] Support mutations of the same input in fw and bw

997fec9

[ghstack-poisoned]

IvanKobzarev requested review from Chillee, albanD, bdhirsh, ezyang and soulitzer as code owners June 6, 2025 18:36

IvanKobzarev mentioned this pull request Jun 6, 2025

[aotd] Support mutations in reordering_to_mimic_autograd_engine #155353

Closed

pytorch-bot bot added ciflow/inductor release notes: fx release notes category labels Jun 6, 2025

facebook-github-bot added the fx label Jun 6, 2025

IvanKobzarev requested a review from zou3519 June 6, 2025 18:58

IvanKobzarev added the topic: not user facing topic category label Jun 6, 2025

bdhirsh reviewed Jun 6, 2025

View reviewed changes

torch/fx/experimental/proxy_tensor.py Show resolved Hide resolved

IvanKobzarev added a commit that referenced this pull request Jun 6, 2025

[aotd] Support mutations of the same input in fw and bw

6786980

ghstack-source-id: 95a4053 Pull Request resolved: #155354

IvanKobzarev added a commit that referenced this pull request Jun 9, 2025

[aotd] Support mutations of the same input in fw and bw

576c0cc

ghstack-source-id: 9fb4349 Pull Request resolved: #155354

albanD removed their request for review June 9, 2025 15:18

IvanKobzarev added a commit that referenced this pull request Jun 10, 2025

[aotd] Support mutations of the same input in fw and bw

ff6d1c4

ghstack-source-id: 73ebf15 Pull Request resolved: #155354

This was referenced Jun 11, 2025

In-place operations are reordered across the forward-backward in autograd function #154820

Closed

[2/2] proxy_tensor do not clobber for mutating ops #155716

Closed

pytorchmergebot reopened this Jun 24, 2025

IvanKobzarev added a commit that referenced this pull request Jun 24, 2025

[aotd] Support mutations of the same input in fw and bw

8a39bd2

ghstack-source-id: 4db1990 Pull Request resolved: #155354

pytorchmergebot added the merging label Jun 25, 2025

pytorchmergebot removed the merging label Jun 25, 2025

IvanKobzarev added a commit that referenced this pull request Jun 25, 2025

[aotd] Support mutations of the same input in fw and bw

ae7c587

ghstack-source-id: 485f601 Pull Request resolved: #155354

IvanKobzarev added a commit that referenced this pull request Jun 26, 2025

[aotd] Support mutations of the same input in fw and bw

03f4411

ghstack-source-id: 9d633c6 Pull Request resolved: #155354

pytorchmergebot added the merging label Jun 26, 2025

pytorchmergebot closed this in 2f94f69 Jun 26, 2025

pytorchmergebot removed the merging label Jun 26, 2025

zou3519 mentioned this pull request Jun 26, 2025

Add AOTDispatcher config to set backward autocast behavior #156356

Closed

bdhirsh mentioned this pull request Dec 15, 2025

fix input mutation handling for subclasses that perform intermediate compute during copy_ (DTensor) #170467

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[aotd] Support mutations of the same input in fw and bw#155354

[aotd] Support mutations of the same input in fw and bw#155354
IvanKobzarev wants to merge 16 commits intogh/IvanKobzarev/114/basefrom
gh/IvanKobzarev/114/head

IvanKobzarev commented Jun 6, 2025 •

edited by pytorch-bot bot

Loading

Uh oh!

pytorch-bot bot commented Jun 6, 2025 •

edited

Loading

Uh oh!

Uh oh!

IvanKobzarev commented Jun 25, 2025

Uh oh!

pytorchmergebot commented Jun 25, 2025

Uh oh!

pytorchmergebot commented Jun 25, 2025

Uh oh!

IvanKobzarev commented Jun 26, 2025

Uh oh!

pytorchmergebot commented Jun 26, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

IvanKobzarev commented Jun 6, 2025 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jun 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/155354

⏳ No Failures, 1 Pending

Uh oh!

Uh oh!

IvanKobzarev commented Jun 25, 2025

Uh oh!

pytorchmergebot commented Jun 25, 2025

Merge started

Uh oh!

pytorchmergebot commented Jun 25, 2025

Merge failed

Uh oh!

IvanKobzarev commented Jun 26, 2025

Uh oh!

pytorchmergebot commented Jun 26, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

IvanKobzarev commented Jun 6, 2025 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Jun 6, 2025 •

edited

Loading