Move autoparallel to use leaner export API #181
Merged
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
meta-cla
bot
added
the
CLA Signed
tugsbayasgalan
commented
Oct 1, 2025
| 3) Be more careful about tensor constants names. | ||
| """ | ||
| with torch._dynamo.config.patch(install_free_tensors=True): | ||
| gm = _dynamo_graph_capture_for_export(model)(*inputs) |
Contributor
Author
There was a problem hiding this comment.
This function is used for new strict export and it adds implicitly sets up right calling convention using pytree. We might need to change this a bit to rely on bytecode instead.
tugsbayasgalan
commented
Oct 1, 2025
| """ | ||
| with torch._dynamo.config.patch(install_free_tensors=True): | ||
| gm = _dynamo_graph_capture_for_export(model)(*inputs) | ||
| _restore_state_dict(model, gm) |
Contributor
Author
There was a problem hiding this comment.
autoparallel wants the original fqn to be stored.
ezyang
approved these changes
Oct 1, 2025
tugsbayasgalan
force-pushed
the
tugsuu_export
branch
from
8bd72ef to
e66dfb1
Compare
SherlockNoMad
approved these changes
Oct 2, 2025
Member
|
rebase for the ci fixes |
tugsbayasgalan
force-pushed
the
tugsuu_export
branch
from
e66dfb1 to
efd8617
Compare
tugsbayasgalan
added a commit
that referenced
this pull request
Oct 10, 2025
This reverts commit f1887eb.
tugsbayasgalan
added a commit
that referenced
this pull request
Oct 10, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
it needs to land after: pytorch/pytorch#164401
All the fullgraph_capture based APIs are moving towards using a common datastructure called GraphCaptureOutput. The idea is that this data structure would be used as source of truth to pull all necessary information to build a runnable compiled artifact. By doing so, we aim to enable various downstream users to build their products using modular components. Today autoparallel is using export directly which carries lot of extra stuff that autoparallel probably doesn't care about. So this PR shows that we can indeed just use a simple lean API instead of export.
Since there might be some churn in the future on the exact API, I abstract it away in a util function.
Some things we need to add are: