[export] refactor and add same_signature flag to dynamo.export #106569
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/106569
Note: Links to docs will display an error until the docs builds have been completed. ✅ 3 Unrelated FailuresAs of commit f238848: BROKEN TRUNK - The following jobs failed but were present on the merge base 4c46ea5:👉 Rebase onto the `viable/strict` branch to avoid these failures
UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
ydwu4
added
release notes: export
ciflow/trunk
ydwu4
force-pushed
the
refactor_export_sig
branch
from
0004952 to
f238848
Compare
|
Can you help us understand what is code motion and what is new? Or split into code motion only, and then adding the new knob. |
Sure! No new code except the knob. It just moves existing functions into rewrite_signature so that we don't need to have to check the knob here and there. |
|
Come to Friday export meeting, we'll talk about how to properly expose this knob with others there, when generally talking about export phases. |
|
@pytorchbot merge |
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 |
…ch#106569) This PR adds a **same_signature** flag to dynamo.export. **Motivation:** In pytorch#105679, we experimented on **using dynamo to inspect the UDFs** for cond in eager mode (without torch.compile). This helps us to normalize the inputs (e.g. lifting closure to inputs) and makes higher order operator more robust (e.g. forbid python side effects) and less error-prone in general. We decided to use dynamo.export (instead of torch.compile) to do the inspection (pointed out by @voznesenskym @zou3519): - We'd like a **whole-graph capture** for the UDF. - We'd like the dynamo inspection to be **stateless**. Using torch.compile would require resetting dynamo context before and after the inspection because the compile flags may be different from users' torch.compile. This will clear all dynamo cache. - We can still implement some **caching** based on the guards. However, this requires export to be able to handle the case where it cannot always rewrite signature: e.g. closure lifted as input. This PR makes the rewrite optional. **Implementation:** We just put all the code that are related to signature rewriting into a function called rewrite_signature and use a same_signature flag to optionally to the transformation. **Test Plan:** existing tests. Pull Request resolved: pytorch#106569 Approved by: https://github.com/ezyang
3 participants
This PR adds a same_signature flag to dynamo.export.
Motivation:
In #105679, we experimented on using dynamo to inspect the UDFs for cond in eager mode (without torch.compile). This helps us to normalize the inputs (e.g. lifting closure to inputs) and makes higher order operator more robust (e.g. forbid python side effects) and less error-prone in general.
We decided to use dynamo.export (instead of torch.compile) to do the inspection (pointed out by @voznesenskym @zou3519):
However, this requires export to be able to handle the case where it cannot always rewrite signature: e.g. closure lifted as input.
This PR makes the rewrite optional.
Implementation:
We just put all the code that are related to signature rewriting into a function called rewrite_signature and use a same_signature flag to optionally to the transformation.
Test Plan:
existing tests.
cc @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @aakhundov