KEMBAR78
[export] fix joint graph metadata by pianpwk · Pull Request #136011 · pytorch/pytorch · GitHub
Skip to content

[export] fix joint graph metadata #136011

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 1 commit into from
Closed

[export] fix joint graph metadata #136011

wants to merge 1 commit into from

Conversation

@pianpwk
Copy link
Contributor

@pianpwk pianpwk commented Sep 13, 2024

Differential Revision: D62652832

@pytorch-bot
Copy link

pytorch-bot bot commented Sep 13, 2024 •
edited
Loading

🔗 Helpful Links

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

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

✅ You can merge normally! (2 Unrelated Failures)

As of commit 31f509c with merge base c0a930b (image):

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

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

@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D62652832

facebook-github-bot pushed a commit to pytorch/executorch that referenced this pull request Sep 17, 2024
@pianpwk @facebook-github-bot
fix joint graph metadata (#136011)
a09d817 
Summary: Pull Request resolved: pytorch/pytorch#136011

Differential Revision: D62652832
@pianpwk pianpwk mentioned this pull request Sep 17, 2024
Open
facebook-github-bot pushed a commit that referenced this pull request Oct 2, 2024
@pianpwk @facebook-github-bot
[export] fix joint graph metadata (#136011)
5b57952 
Summary:
Fixes up missing metadata on export -> joint graph export, with patching in AOTAutograd so tracer knows to record nn_module_stack

X-link: pytorch/executorch#5447


Test Plan: test_experimental

Differential Revision: D62652832
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D62652832

tugsbayasgalan
tugsbayasgalan reviewed Oct 2, 2024
View reviewed changes
torch/_functorch/aot_autograd.py Outdated
return *fw_outs, *output_gradients

fx_g = make_fx(flattened_joint)(*full_args)
flattened_joint._orig_mod = fx_g
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we use functools.wraps to get the orig mod? Can you also reference this note here https://github.com/pytorch/pytorch/pull/115454/files#diff-594ef3e84a8e3555bfa4a6167383949da5f321e93f602246eb9c0a797d5ecc00R617
cc: @bdhirsh

@tugsbayasgalan tugsbayasgalan requested a review from bdhirsh October 2, 2024 17:37
bdhirsh
bdhirsh reviewed Oct 2, 2024
View reviewed changes
test/export/test_experimental.py
inputs = (x, labels)

ep = export(net, inputs)
ep = _export_forward_backward(ep)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what does the test failure look like without your fix? (if _orig_mod isn't properly plumbed in AOTAutograd).

It feels pretty easy to add a new wrapper at some point and forget to propagate the orig_mod, so we probably want it to error in a loud/obvious way.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@bdhirsh it crashes at the ExportedProgram verifier, as nn_module_stack and other metadata is missing for the graph.

I think it affects this line, where ModuleStackTracer populates the metadata:

pytorch/torch/fx/experimental/proxy_tensor.py

Lines 1872 to 1876 in 4d3c0fc

if hasattr(f, "_orig_mod") and self.record_module_stack:
scope_root = f._orig_mod
self.fx_tracer = _ModuleStackTracer(scope_root)
else:
self.fx_tracer = PythonKeyTracer()

@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D62652832

facebook-github-bot pushed a commit that referenced this pull request Oct 3, 2024
@pianpwk @facebook-github-bot
[export] fix joint graph metadata (#136011)
156b78e 
Summary:
Fixes up missing metadata on export -> joint graph export, with patching in AOTAutograd so tracer knows to record nn_module_stack

X-link: pytorch/executorch#5447


Test Plan: test_experimental

Differential Revision: D62652832
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D62652832

@pianpwk
Copy link
Contributor Author

pianpwk commented Oct 3, 2024

updated to use functool.wraps on flattened_joint, let's see if this breaks anything

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 3, 2024
facebook-github-bot pushed a commit that referenced this pull request Oct 4, 2024
@pianpwk @facebook-github-bot
[export] fix joint graph metadata (#136011)
56ab88e 
Summary:
Fixes up missing metadata on export -> joint graph export, with patching in AOTAutograd so tracer knows to record nn_module_stack

X-link: pytorch/executorch#5447


Test Plan: test_experimental

Differential Revision: D62652832
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D62652832

@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D62652832

@pianpwk @facebook-github-bot
[export] fix joint graph metadata (#136011)
31f509c 
Summary:
Fixes up missing metadata on export -> joint graph export, with patching in AOTAutograd so tracer knows to record nn_module_stack

X-link: pytorch/executorch#5447


Test Plan: test_experimental

Reviewed By: tugsbayasgalan

Differential Revision: D62652832
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D62652832

@facebook-github-bot
Copy link
Contributor

@pytorchbot merge -f 'Landed internally'

(Initiating merge automatically since Phabricator Diff has merged, using force because this PR might not pass merge_rules.json but landed internally)

@pytorchmergebot
Copy link
Collaborator

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

@github-actions github-actions bot deleted the export-D62652832 branch November 7, 2024 02:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@bdhirsh bdhirsh bdhirsh left review comments

@tugsbayasgalan tugsbayasgalan tugsbayasgalan approved these changes

@avikchaudhuri avikchaudhuri Awaiting requested review from avikchaudhuri avikchaudhuri is a code owner

@zhxchen17 zhxchen17 Awaiting requested review from zhxchen17 zhxchen17 is a code owner

@ydwu4 ydwu4 Awaiting requested review from ydwu4 ydwu4 is a code owner

@angelayi angelayi Awaiting requested review from angelayi angelayi is a code owner

Assignees

No one assigned

Labels

ciflow/inductor ciflow/trunk Trigger trunk jobs on your pull request fb-exported Merged release notes: AO frontend

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@pianpwk @facebook-github-bot @pytorchmergebot @bdhirsh @tugsbayasgalan