[invoke_subgraph] Do not cache fake tensors for AOTDispatcher first pass #150450

anijain2305 · 2025-04-01T19:05:30Z

Stack from ghstack (oldest at bottom):

[ghstack-poisoned]

pytorch-bot · 2025-04-01T19:05:33Z

🔗 Helpful Links

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

📄 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.

✅ You can merge normally! (3 Unrelated Failures)

As of commit 513c1a2 with merge base 15dbad2 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

inductor / unit-test / linux-jammy-cpu-py3.12-gcc11-inductor-halide / build (gh) (trunk failure)

UNSTABLE - The following jobs are marked as unstable, possibly due to flakiness on trunk:

pull / cuda12.4-py3.10-gcc9-sm75 / test (pr_time_benchmarks, 1, 1, linux.g4dn.metal.nvidia.gpu) (gh) (#149370)
REGRESSION: benchmark ('aotdispatcher_training_nosubclass_cpu', 'compile_time_instruction_count') failed, actual result 3636524535 is 1.61% higher than expected 3579000000 ±+1.50% if this is an expected regression, please update the expected results.
pull / linux-jammy-xpu-2025.0-py3.9 / build (gh) (#150430)
/usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/sstream:152:52: error: expected value in expression

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

angelayi · 2025-04-01T19:58:56Z

torch/_subclasses/fake_tensor.py

+                        # AOTDispatcher first pass does not run make_fx on
+                        # dynamo graphs. As a result, it can have non OpOverload
+                        # ops.
+                        if not isinstance(op, torch._ops.OpOverload):


This looks fine to me but I'm not sure if there are any other problems from doing so, so will leave it to Richard for the stamp :3

zou3519

Seems better than the previous state. But I think we should be tracing through the inside of the HOP subgraph so that the operator.add call gets desugared into an aten.add call, instead of manually iterating through the nodes and calling _validate_cache_key on each node. Does that make sense?

anijain2305 · 2025-04-01T21:27:35Z

Seems better than the previous state. But I think we should be tracing through the inside of the HOP subgraph so that the operator.add call gets desugared into an aten.add call, instead of manually iterating through the nodes and calling _validate_cache_key on each node. Does that make sense?

I plan to work on this. But till then, it would be good to get this in to make trunk healthy, as Lazos, me and Angela are working with invoke_subgraph recently.

zou3519

sgtm

anijain2305 · 2025-04-01T22:49:57Z

@pytorchbot merge

pytorchmergebot · 2025-04-01T22:51:42Z

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

…alse (#150486) I am not sure if this is the right way. Pull Request resolved: #150486 Approved by: https://github.com/zou3519 ghstack dependencies: #150082, #150450

…t module (#150556) Pull Request resolved: #150556 Approved by: https://github.com/bdhirsh, https://github.com/zou3519 ghstack dependencies: #150082, #150450, #150486

…150561) I am unable to come up with a testcase. It passes many end-to-end tests that fail with ReshapeError at https://ossci-raw-job-status.s3.amazonaws.com/log/39717218372 ![image](https://github.com/user-attachments/assets/8509b485-3897-4538-968b-bbe05af63a59) Pull Request resolved: #150561 Approved by: https://github.com/zou3519, https://github.com/bdhirsh ghstack dependencies: #150082, #150450, #150486, #150556

…ass (pytorch#150450) Pull Request resolved: pytorch#150450 Approved by: https://github.com/zou3519 ghstack dependencies: pytorch#150082

…alse (pytorch#150486) I am not sure if this is the right way. Pull Request resolved: pytorch#150486 Approved by: https://github.com/zou3519 ghstack dependencies: pytorch#150082, pytorch#150450

…t module (pytorch#150556) Pull Request resolved: pytorch#150556 Approved by: https://github.com/bdhirsh, https://github.com/zou3519 ghstack dependencies: pytorch#150082, pytorch#150450, pytorch#150486

…ytorch#150561) I am unable to come up with a testcase. It passes many end-to-end tests that fail with ReshapeError at https://ossci-raw-job-status.s3.amazonaws.com/log/39717218372 ![image](https://github.com/user-attachments/assets/8509b485-3897-4538-968b-bbe05af63a59) Pull Request resolved: pytorch#150561 Approved by: https://github.com/zou3519, https://github.com/bdhirsh ghstack dependencies: pytorch#150082, pytorch#150450, pytorch#150486, pytorch#150556

ghstack-source-id: fd95d6d Pull Request resolved: pytorch/pytorch#150450

[invoke_subgraph] Do not cache fake tensors for AOTDispatcher first pass

513c1a2

[ghstack-poisoned]

anijain2305 mentioned this pull request Apr 1, 2025

[invoke_subgraph] Support None in the fwd output #150082

Closed

pytorch-bot bot added the ciflow/inductor label Apr 1, 2025

anijain2305 requested review from angelayi and zou3519 April 1, 2025 19:05

anijain2305 added the topic: not user facing topic category label Apr 1, 2025

angelayi reviewed Apr 1, 2025

View reviewed changes

zou3519 reviewed Apr 1, 2025

View reviewed changes

zou3519 approved these changes Apr 1, 2025

View reviewed changes

anijain2305 added the ciflow/trunk Trigger trunk jobs on your pull request label Apr 1, 2025

pytorchmergebot added the merging label Apr 1, 2025

anijain2305 mentioned this pull request Apr 1, 2025

[invoke_subgraph] Filter out grad_out where fw_out requires_grad is False #150486

Closed

pytorchmergebot added the Merged label Apr 2, 2025

pytorchmergebot closed this in 61ebe99 Apr 2, 2025

pytorchmergebot removed the merging label Apr 2, 2025

This was referenced Apr 2, 2025

[invoke_subgraph][min-cut partitioner] Fix bug to use the correct root module #150556

Closed

[invoke_subgraph] Force grad_outs to be contiguous at tracing time #150561

Closed

This was referenced Apr 4, 2025

[invoke_subgraph] Lazy backward #150666

Closed

[invoke_subgraph][fake_tensor] Run the subgraph with fake tensor mode to validate cache #150704

Closed

Divigroup-RAP pushed a commit to Divigroup-RAP/PYTORCH that referenced this pull request Apr 22, 2025

[invoke_subgraph] Do not cache fake tensors for AOTDispatcher first pass

0dffffc

ghstack-source-id: fd95d6d Pull Request resolved: pytorch/pytorch#150450

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[invoke_subgraph] Do not cache fake tensors for AOTDispatcher first pass #150450

[invoke_subgraph] Do not cache fake tensors for AOTDispatcher first pass #150450

Uh oh!

anijain2305 commented Apr 1, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Apr 1, 2025 •

edited

Loading

Uh oh!

angelayi Apr 1, 2025

Uh oh!

zou3519 left a comment

Uh oh!

anijain2305 commented Apr 1, 2025

Uh oh!

zou3519 left a comment

Uh oh!

anijain2305 commented Apr 1, 2025

Uh oh!

pytorchmergebot commented Apr 1, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

[invoke_subgraph] Do not cache fake tensors for AOTDispatcher first pass #150450

[invoke_subgraph] Do not cache fake tensors for AOTDispatcher first pass #150450

Uh oh!

Conversation

anijain2305 commented Apr 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Apr 1, 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/150450

✅ You can merge normally! (3 Unrelated Failures)

Uh oh!

angelayi Apr 1, 2025

Choose a reason for hiding this comment

Uh oh!

zou3519 left a comment

Choose a reason for hiding this comment

Uh oh!

anijain2305 commented Apr 1, 2025

Uh oh!

zou3519 left a comment

Choose a reason for hiding this comment

Uh oh!

anijain2305 commented Apr 1, 2025

Uh oh!

pytorchmergebot commented Apr 1, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

anijain2305 commented Apr 1, 2025 •

edited

Loading

pytorch-bot bot commented Apr 1, 2025 •

edited

Loading