Add BundledAOTAutogradSerializableCallable #162170

jamesjwu · 2025-09-04T15:01:34Z

Stack from ghstack (oldest at bottom):

This PR hooks up the python wrapper inductor backend to aot_compile. This is not the best way for us to grab the output of AOTAutograd; that involves a refactor to make AOTAutograd itself return a serializable callable. I'll do that refactor soon, but I want a basic interface to test with for now.

In the medium term, we'll want aot_compile to call AOTAutograd directly, instead of using the TorchInductorWrapper's callback through compile_fx.

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

[ghstack-poisoned]

pytorch-bot · 2025-09-04T15:01:39Z

🔗 Helpful Links

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

📄 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

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

⏳ No Failures, 43 Pending

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

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

[ghstack-poisoned]

jamesjwu · 2025-09-07T15:34:15Z

@pytorchbot merge

pytorchmergebot · 2025-09-07T15:37:09Z

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-09-07T18:12:37Z

Merge failed

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

Auto-merging torch/_dynamo/aot_compile.py
CONFLICT (content): Merge conflict in torch/_dynamo/aot_compile.py
error: could not apply 4413729db12... [easy] [precompile] Convert CompileArtifacts to callable
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

[ghstack-poisoned]

jamesjwu · 2025-09-07T22:16:34Z

@pytorchbot merge

pytorchmergebot · 2025-09-07T22:18:18Z

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-09-07T22:50:03Z

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / win-vs2022-cpu-py3 / test (default, 1, 3, windows.4xlarge.nonephemeral)

Details for Dev Infra team

Raised by workflow job

jamesjwu · 2025-09-07T23:29:35Z

@pytorchbot merge -i

pytorchmergebot · 2025-09-07T23:31:40Z

Merge started

Your change will be merged while ignoring the following 1 checks: trunk / win-vs2022-cpu-py3 / test (default, 1, 3, windows.4xlarge.nonephemeral)

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

huydhn · 2025-09-09T03:36:49Z

@pytorchbot rebase -b main

pytorchmergebot · 2025-09-09T03:38:24Z

@pytorchbot started a rebase job onto refs/remotes/origin/main. Check the current status here

pytorchmergebot · 2025-09-09T03:38:30Z

Rebase failed due to

Aborting rebase because rebasing the branch resulted in the same sha as the target branch.
This usually happens because the PR has already been merged.  Please rebase locally and push.

Raised by https://github.com/pytorch/pytorch/actions/runs/17570981566

huydhn · 2025-09-09T04:57:09Z

It turns out that I need to revert #162183 for breaking macos job in trunk, so let me reland this as it would not be conflicted anymore

huydhn · 2025-09-09T05:29:22Z

@pytorchbot merge -f 'Reland'

pytorch-bot · 2025-09-09T05:29:26Z

You need to provide a reason for using force merge, in the format @pytorchbot merge -f 'Explanation'.
The explanation needs to be clear on why this is needed. Here are some good examples:

Bypass checks due to unrelated upstream failures from ...
This is a minor fix to ..., which shouldn't break anything
This is pre-tested in a previous CI run
Bypass flaky ... check

huydhn · 2025-09-09T05:40:22Z

@pytorchbot merge -f 'Have not seen this check in action for a while, I am relanding the same change #162170 (comment)'

pytorchmergebot · 2025-09-09T05:42:00Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

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 hooks up the python wrapper inductor backend to aot_compile. This is *not* the best way for us to grab the output of AOTAutograd; that involves a refactor to make AOTAutograd itself return a serializable callable. I'll do that refactor soon, but I want a basic interface to test with for now. In the medium term, we'll want aot_compile to call AOTAutograd directly, instead of using the TorchInductorWrapper's callback through compile_fx. Pull Request resolved: pytorch#162170 Approved by: https://github.com/zhxchen17 ghstack dependencies: pytorch#162169

This reverts commit 5babb4d. Reverted pytorch#162170 on behalf of https://github.com/huydhn due to This PR has a merge conflict with D81793200 on aot_compile.py where PRs and diffs are landed in reverted order ([comment](pytorch#162170 (comment)))

This PR hooks up the python wrapper inductor backend to aot_compile. This is *not* the best way for us to grab the output of AOTAutograd; that involves a refactor to make AOTAutograd itself return a serializable callable. I'll do that refactor soon, but I want a basic interface to test with for now. In the medium term, we'll want aot_compile to call AOTAutograd directly, instead of using the TorchInductorWrapper's callback through compile_fx. Pull Request resolved: pytorch#162170 Approved by: https://github.com/zhxchen17 ghstack dependencies: pytorch#162169

This reverts commit 5babb4d. Reverted pytorch#162170 on behalf of https://github.com/huydhn due to This PR has a merge conflict with D81793200 on aot_compile.py where PRs and diffs are landed in reverted order ([comment](pytorch#162170 (comment)))

This PR hooks up the python wrapper inductor backend to aot_compile. This is *not* the best way for us to grab the output of AOTAutograd; that involves a refactor to make AOTAutograd itself return a serializable callable. I'll do that refactor soon, but I want a basic interface to test with for now. In the medium term, we'll want aot_compile to call AOTAutograd directly, instead of using the TorchInductorWrapper's callback through compile_fx. Pull Request resolved: pytorch#162170 Approved by: https://github.com/zhxchen17 ghstack dependencies: pytorch#162169

This reverts commit 5babb4d. Reverted pytorch#162170 on behalf of https://github.com/huydhn due to This PR has a merge conflict with D81793200 on aot_compile.py where PRs and diffs are landed in reverted order ([comment](pytorch#162170 (comment)))

This PR hooks up the python wrapper inductor backend to aot_compile. This is *not* the best way for us to grab the output of AOTAutograd; that involves a refactor to make AOTAutograd itself return a serializable callable. I'll do that refactor soon, but I want a basic interface to test with for now. In the medium term, we'll want aot_compile to call AOTAutograd directly, instead of using the TorchInductorWrapper's callback through compile_fx. Pull Request resolved: pytorch#162170 Approved by: https://github.com/zhxchen17 ghstack dependencies: pytorch#162169

This reverts commit 5babb4d. Reverted pytorch#162170 on behalf of https://github.com/huydhn due to This PR has a merge conflict with D81793200 on aot_compile.py where PRs and diffs are landed in reverted order ([comment](pytorch#162170 (comment)))

This PR hooks up the python wrapper inductor backend to aot_compile. This is *not* the best way for us to grab the output of AOTAutograd; that involves a refactor to make AOTAutograd itself return a serializable callable. I'll do that refactor soon, but I want a basic interface to test with for now. In the medium term, we'll want aot_compile to call AOTAutograd directly, instead of using the TorchInductorWrapper's callback through compile_fx. Pull Request resolved: pytorch#162170 Approved by: https://github.com/zhxchen17 ghstack dependencies: pytorch#162169

Update

6c8e3bd

[ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: dynamo labels Sep 4, 2025

This was referenced Sep 4, 2025

[easy] [precompile] Convert CompileArtifacts to callable #162169

Closed

[Precompile] [RFC] Implement aot_compile_module #162171

Closed

jamesjwu requested review from ezyang and zhxchen17 September 4, 2025 15:04

jamesjwu added the topic: not user facing topic category label Sep 4, 2025

jamesjwu marked this pull request as ready for review September 4, 2025 15:06

zhxchen17 approved these changes Sep 4, 2025

View reviewed changes

jamesjwu added 3 commits September 4, 2025 10:22

Update

6bb7791

[ghstack-poisoned]

Code review

7a3c081

[ghstack-poisoned]

Update

9bd005e

[ghstack-poisoned]

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Sep 7, 2025

pytorchmergebot added the merging label Sep 7, 2025

pytorchmergebot removed the merging label Sep 7, 2025

Rebaes

c24b85b

[ghstack-poisoned]

pytorchmergebot added the merging label Sep 7, 2025

pytorchmergebot removed the merging label Sep 7, 2025

pytorchmergebot added the merging label Sep 7, 2025

pytorchmergebot added Reverted ci-no-td Do not run TD on this PR labels Sep 9, 2025

pytorchmergebot reopened this Sep 9, 2025

huydhn mentioned this pull request Sep 9, 2025

testing infra and some fixes #162183

Closed

pytorch deleted a comment from pytorch-bot bot Sep 9, 2025

pytorchmergebot added the merging label Sep 9, 2025

pytorchmergebot closed this in d85392a Sep 9, 2025

pytorchmergebot removed the merging label Sep 9, 2025

github-actions bot deleted the gh/jamesjwu/189/head branch October 10, 2025 02:09

Add BundledAOTAutogradSerializableCallable #162170

Add BundledAOTAutogradSerializableCallable #162170

Uh oh!

Conversation

jamesjwu commented Sep 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 4, 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/162170

⏳ No Failures, 43 Pending

Uh oh!

jamesjwu commented Sep 7, 2025

Uh oh!

pytorchmergebot commented Sep 7, 2025

Merge started

Uh oh!

pytorchmergebot commented Sep 7, 2025

Merge failed

Uh oh!

jamesjwu commented Sep 7, 2025

Uh oh!

pytorchmergebot commented Sep 7, 2025

Merge started

Uh oh!

pytorchmergebot commented Sep 7, 2025

Merge failed

Uh oh!

jamesjwu commented Sep 7, 2025

Uh oh!

pytorchmergebot commented Sep 7, 2025

Merge started

Uh oh!

huydhn commented Sep 9, 2025

Uh oh!

pytorchmergebot commented Sep 9, 2025

Uh oh!

pytorchmergebot commented Sep 9, 2025

Uh oh!

huydhn commented Sep 9, 2025

Uh oh!

huydhn commented Sep 9, 2025

Uh oh!

pytorch-bot bot commented Sep 9, 2025

Uh oh!

huydhn commented Sep 9, 2025

Uh oh!

pytorchmergebot commented Sep 9, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

jamesjwu commented Sep 4, 2025 •

edited

Loading

pytorch-bot bot commented Sep 4, 2025 •

edited

Loading