Add support for tracing vmap in pre-dispatch export #154650
Closed
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
tugsbayasgalan
requested review from
angelayi,
avikchaudhuri,
bobrenjc93,
laithsakka,
ydwu4 and
zhxchen17
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/154650
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 8459179 with merge base 5ee464d ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
This pull request was exported from Phabricator. Differential Revision: D75623875 |
bdhirsh
reviewed
May 29, 2025
bdhirsh
reviewed
May 29, 2025
bdhirsh
reviewed
May 29, 2025
bdhirsh
reviewed
May 29, 2025
bdhirsh
reviewed
May 29, 2025
bdhirsh
reviewed
May 29, 2025
bdhirsh
reviewed
May 29, 2025
zou3519
reviewed
Jun 5, 2025
There was a problem hiding this comment.
Our conclusion from the meeting on Tuesday was:
- Yes, we're going to put all the API calls into the graph
- We should only interpose on these API calls when non-strict export is on. These shouldn't go through regular torch_function, because they are private APIs.
- It might be easier to do this by creating a python function wrapper around e.g. add_batch_dim, and then putting the torch_function handler and export checks into said function. This may require you to add these new functions to Dynamo skiplists to not break the Dynamo side of things
tugsbayasgalan
force-pushed
the
export-D75623875
branch
from
7472d71 to
bd6401b
Compare
pytorch-bot bot
pushed a commit
that referenced
this pull request
Jul 23, 2025
Summary: ONNX team ran into this error and we also ran into during our export benchmarking. This diff makes it possible to trace through vmap implementation in pre-dispatch IR. Note that we don't support serializing functorch ops in pre-dispatch IR and in the future, we should desugar them to post-grad ops. Test Plan: CI Differential Revision: D75623875
|
This pull request was exported from Phabricator. Differential Revision: D75623875 |
1 similar comment
|
This pull request was exported from Phabricator. Differential Revision: D75623875 |
tugsbayasgalan
added a commit
to tugsbayasgalan/pytorch
that referenced
this pull request
Jul 23, 2025
Summary: Pull Request resolved: pytorch#154650 ONNX team ran into this error and we also ran into during our export benchmarking. This diff makes it possible to trace through vmap implementation in pre-dispatch IR. Note that we don't support serializing functorch ops in pre-dispatch IR and in the future, we should desugar them to post-grad ops. Test Plan: CI Differential Revision: D75623875
tugsbayasgalan
force-pushed
the
export-D75623875
branch
from
bd6401b to
c0a8978
Compare
tugsbayasgalan
force-pushed
the
export-D75623875
branch
from
71f718a to
11f079b
Compare
tugsbayasgalan
added a commit
to tugsbayasgalan/pytorch
that referenced
this pull request
Aug 20, 2025
Summary: ONNX team ran into this error and we also ran into during our export benchmarking. This diff makes it possible to trace through vmap implementation in pre-dispatch IR. Note that we don't support serializing functorch ops in pre-dispatch IR and in the future, we should desugar them to post-grad ops. Test Plan: CI Differential Revision: D75623875
|
This pull request was exported from Phabricator. Differential Revision: D75623875 |
tugsbayasgalan
added a commit
to tugsbayasgalan/pytorch
that referenced
this pull request
Aug 20, 2025
Summary: ONNX team ran into this error and we also ran into during our export benchmarking. This diff makes it possible to trace through vmap implementation in pre-dispatch IR. Note that we don't support serializing functorch ops in pre-dispatch IR and in the future, we should desugar them to post-grad ops. Test Plan: CI Differential Revision: D75623875
tugsbayasgalan
force-pushed
the
export-D75623875
branch
from
11f079b to
b327868
Compare
|
This pull request was exported from Phabricator. Differential Revision: D75623875 |
tugsbayasgalan
added a commit
to tugsbayasgalan/pytorch
that referenced
this pull request
Aug 20, 2025
Summary: ONNX team ran into this error and we also ran into during our export benchmarking. This diff makes it possible to trace through vmap implementation in pre-dispatch IR. Note that we don't support serializing functorch ops in pre-dispatch IR and in the future, we should desugar them to post-grad ops. Test Plan: CI Differential Revision: D75623875
tugsbayasgalan
force-pushed
the
export-D75623875
branch
from
b327868 to
6592b95
Compare
|
This pull request was exported from Phabricator. Differential Revision: D75623875 |
tugsbayasgalan
added a commit
to tugsbayasgalan/pytorch
that referenced
this pull request
Aug 20, 2025
Summary: ONNX team ran into this error and we also ran into during our export benchmarking. This diff makes it possible to trace through vmap implementation in pre-dispatch IR. Note that we don't support serializing functorch ops in pre-dispatch IR and in the future, we should desugar them to post-grad ops. Test Plan: CI Differential Revision: D75623875
tugsbayasgalan
force-pushed
the
export-D75623875
branch
from
6592b95 to
d4b6021
Compare
|
This pull request was exported from Phabricator. Differential Revision: D75623875 |
tugsbayasgalan
added a commit
to tugsbayasgalan/pytorch
that referenced
this pull request
Aug 20, 2025
Summary: ONNX team ran into this error and we also ran into during our export benchmarking. This diff makes it possible to trace through vmap implementation in pre-dispatch IR. Note that we don't support serializing functorch ops in pre-dispatch IR and in the future, we should desugar them to post-grad ops. Test Plan: CI Differential Revision: D75623875
tugsbayasgalan
force-pushed
the
export-D75623875
branch
from
d4b6021 to
65fb23a
Compare
|
This pull request was exported from Phabricator. Differential Revision: D75623875 |
zou3519
reviewed
Aug 20, 2025
zou3519
approved these changes
Aug 20, 2025
Summary: ONNX team ran into this error and we also ran into during our export benchmarking. This diff makes it possible to trace through vmap implementation in pre-dispatch IR. Note that we don't support serializing functorch ops in pre-dispatch IR and in the future, we should desugar them to post-grad ops. Test Plan: CI Reviewed By: zou3519 Differential Revision: D75623875
tugsbayasgalan
force-pushed
the
export-D75623875
branch
from
65fb23a to
8459179
Compare
|
This pull request was exported from Phabricator. Differential Revision: D75623875 |
|
@pytorchbot merge (Initiating merge automatically since Phabricator Diff has merged) |
Merge startedYour 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 |
markc-614
pushed a commit
to markc-614/pytorch
that referenced
this pull request
Sep 17, 2025
Summary: ONNX team and recent transformer upgrade ran into this error and we also ran into during our export benchmarking. This diff makes it possible to trace through vmap implementation in pre-dispatch IR. Note that we don't support serializing functorch ops in pre-dispatch IR and in the future, we should desugar them to post-grad ops. The implementation strategy is: 1. We add python wrappers around vmap APIs so that we attach custom torch function handler that is only on during non-strict export. The reason is we don't want to add this to default torch_function handler because it will break BC. 2. Some dynamo changes to make sure it picks up new python wrapper APIs. The reason is when we do strict export, we need to re-materialize these APIs in pre-dispatch IR from torch IR. We can avoid this by special casing in dynamo for export to proxy different API calls but i feel that is too much chaos because you need to be able to proxy 2 different variants of same vmap API. Test Plan: CI Differential Revision: D75623875 Pull Request resolved: pytorch#154650 Approved by: https://github.com/ezyang, https://github.com/zou3519
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
8 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary: ONNX team and recent transformer upgrade ran into this error and we also ran into during our export benchmarking. This diff makes it possible to trace through vmap implementation in pre-dispatch IR. Note that we don't support serializing functorch ops in pre-dispatch IR and in the future, we should desugar them to post-grad ops.
The implementation strategy is:
Test Plan: CI
Differential Revision: D75623875
cc @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv @voznesenskym @penguinwu @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @jiayisunx @chenyang78 @kadeng @chauhang @amjames @Lucaskabela