Add option for TorchDispatchMode to ignore torch.compile internals #161648
Conversation
If TorchDispatchMode.ignore_compile_internals() is True, then we turn off the TorchDispatchMode during the compilation process, instead turning it back on during runtime of the compiled artifact. Test Plan: - new test [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/161648
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 46106a1 with merge base cd87f30 ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
torch/utils/_python_dispatch.py
Outdated
| f(x) | ||
| The above example will not log anything if | ||
| ``LoggingMode.ignore_compile_internals`` is True. |
There was a problem hiding this comment.
in this docstring ignore_compile_internals looks like an attribute, not a class method. Do you want to to keep it as a class method? (my first thought would be that attribute is simpler and easier for a mode to override if they want to, although i'm not sure if ability-to-override is your goal
There was a problem hiding this comment.
is_infra_mode() is a method, so I made this a method. I'm open to making it a class attribute though.
There was a problem hiding this comment.
nvm i see you override the class method in your test, either seems fine
…nternals" If TorchDispatchMode.ignore_compile_internals() is True, then we turn off the TorchDispatchMode during the compilation process, instead turning it back on during runtime of the compiled artifact. Test Plan: - new test cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames Lucaskabela [ghstack-poisoned]
test/dynamo/test_misc.py
Outdated
| foo(x) | ||
|
|
||
| self.assertEqual(len(_checksums), 2) | ||
| # If you are getting "3" here, then that means the .abs().sum() |
There was a problem hiding this comment.
hmm i'm a bit confused by the test - even without your change, wouldnt' the .abs().sum() never get compiled? we only call that in the case when our TorchDispatchMode intercepts the custom op foo, but it doesn't look like the test ever calls mylib.foo in a compiled region.
There was a problem hiding this comment.
I'm just going to delete the comment because I don't know what exactly is going on here. The .abs().sum() isn't being compiled. When Dynamo falls back to eager, it turns out that dynamo decides to compile each of the ops individually (the mul, sin, cos). This means that it had an opportunity to turn itself back on, but I don't know why or how.
There was a problem hiding this comment.
This was a bug in my implementation (great catch!)
torch/_dynamo/convert_frame.py
Outdated
| continue | ||
| if mode.is_infra_mode(): | ||
| continue | ||
| return True |
There was a problem hiding this comment.
hmm do you mind explaining what the behavior is when you have multiple modes active, and one has ignore_compile_internals set but the other does not?
There was a problem hiding this comment.
Dynamo will fallback to eager and all modes will see exactly the same ops that they saw in eager mode
…nternals" If TorchDispatchMode.ignore_compile_internals() is True, then we turn off the TorchDispatchMode during the compilation process, instead turning it back on during runtime of the compiled artifact. Test Plan: - new test cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames Lucaskabela [ghstack-poisoned]
…nternals" If TorchDispatchMode.ignore_compile_internals() is True, then we turn off the TorchDispatchMode during the compilation process, instead turning it back on during runtime of the compiled artifact. Test Plan: - new test cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames Lucaskabela [ghstack-poisoned]
|
@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 |
…ytorch#161648) If TorchDispatchMode.ignore_compile_internals() is True, then we turn off the TorchDispatchMode during the compilation process, instead turning it back on during runtime of the compiled artifact. Test Plan: - new test Pull Request resolved: pytorch#161648 Approved by: https://github.com/bdhirsh
3 participants
Stack from ghstack (oldest at bottom):
If TorchDispatchMode.ignore_compile_internals() is True, then we turn
off the TorchDispatchMode during the compilation process, instead
turning it back on during runtime of the compiled artifact.
Test Plan:
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @Lucaskabela