[export] Update swap's forward function #137102
Conversation
Downstream APS code was failing to run the previously swapped module because of some fx.GraphModule forward function weirdness. So to fix this, I just attached a custom forward function which matches the unflattened module's forward function. Differential Revision: [D63683422](https://our.internmc.facebook.com/intern/diff/D63683422/) [ghstack-poisoned]
angelayi
requested review from
avikchaudhuri,
tugsbayasgalan,
ydwu4 and
zhxchen17
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/137102
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 1338d9d with merge base f2d174c ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
This pull request was exported from Phabricator. Differential Revision: D63683422 |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
torch/export/unflatten.py
Outdated
| flat_args, in_spec = pytree.tree_flatten( | ||
| (args, reordered_kwargs) | ||
| ) |
There was a problem hiding this comment.
| flat_args, in_spec = pytree.tree_flatten( | |
| (args, reordered_kwargs) | |
| ) | |
| flat_args, in_spec = pytree.tree_flatten((args, reordered_kwargs)) |
torch/export/_unlift.py
Outdated
| # TODO(suo) this should not be optional, but is since we still ahve | ||
| # capture_pre_autograd_graph grr | ||
| graph_signature: Optional[ExportGraphSignature] = None, | ||
| _check_input_constraints: bool = True, |
There was a problem hiding this comment.
Hmm why would we not want to run this check when swapping?
There was a problem hiding this comment.
I think it's an additional QPS optimization? This is currently something the unflattener also supports, and is turned off in APS here.
Downstream APS code was failing to run the previously swapped module because of some fx.GraphModule forward function weirdness. So to fix this, I just attached a custom forward function which matches the unflattened module's forward function. Differential Revision: [D63683422](https://our.internmc.facebook.com/intern/diff/D63683422/) [ghstack-poisoned]
|
This pull request was exported from Phabricator. Differential Revision: D63683422 |
torch/export/unflatten.py
Outdated
| flat_args, in_spec = pytree.tree_flatten( | ||
| (args, reordered_kwargs) | ||
| ) |
There was a problem hiding this comment.
| flat_args, in_spec = pytree.tree_flatten( | |
| (args, reordered_kwargs) | |
| ) | |
| flat_args, in_spec = pytree.tree_flatten((args, reordered_kwargs)) |
Downstream APS code was failing to run the previously swapped module because of some fx.GraphModule forward function weirdness (P1594789677). So to fix this, I just attached a custom forward function which matches the unflattened module's forward function. Differential Revision: [D63683422](https://our.internmc.facebook.com/intern/diff/D63683422/) [ghstack-poisoned]
|
This pull request was exported from Phabricator. Differential Revision: D63683422 |
avikchaudhuri
left a comment
avikchaudhuri
left a comment
There was a problem hiding this comment.
Dislike the _check_input_constraints flag polluting the API. What does it even do? If it's something we want to support officially, OK let's do it more systematically. But as is it feels very weird.
Downstream APS code was failing to run the previously swapped module because of some fx.GraphModule forward function weirdness (P1594789677). So to fix this, I just attached a custom forward function which matches the unflattened module's forward function. Differential Revision: [D63683422](https://our.internmc.facebook.com/intern/diff/D63683422/) [ghstack-poisoned]
|
This pull request was exported from Phabricator. Differential Revision: D63683422 |
Downstream APS code was failing to run the previously swapped module because of some fx.GraphModule forward function weirdness (P1594789677). So to fix this, I just attached a custom forward function which matches the unflattened module's forward function. Differential Revision: [D63683422](https://our.internmc.facebook.com/intern/diff/D63683422/) [ghstack-poisoned]
|
This pull request was exported from Phabricator. Differential Revision: D63683422 |
|
@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) |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Refactor so that unflattener doesn't become too messy Differential Revision: [D63719648](https://our.internmc.facebook.com/intern/diff/D63719648/) Pull Request resolved: #137134 Approved by: https://github.com/avikchaudhuri ghstack dependencies: #136191, #137102
Stack from ghstack (oldest at bottom):
Downstream APS code was failing to run the previously swapped module because of some fx.GraphModule forward function weirdness (P1594789677). So to fix this, I just attached a custom forward function which matches the unflattened module's forward function.
Differential Revision: D63683422