[precompile] Integrate AOTI as a backend.#167338
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/167338
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (4 Unrelated Failures)As of commit 62bea25 with merge base d8384e2 ( BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
929d9b9 to
13b69e4
Compare
| with torch.device("cuda"): | ||
| from torch._dynamo.hooks import Hooks | ||
|
|
||
| mod = SimpleLinearModule() |
There was a problem hiding this comment.
Instead of testing it this way, let's test it with actual aot_compile_module, passing in the AOTI backend with ModelInput.
jamesjwu
left a comment
There was a problem hiding this comment.
Let's get tests passing and add some tests around aot_compile_module. Otherwise, though, this looks good, hopefully we can land it quickly.
|
|
||
| fake_mode = detect_fake_mode(inputs_) | ||
| ctx = ( | ||
| mock.patch.object(fake_mode, "allow_non_fake_inputs", True) |
There was a problem hiding this comment.
Why do we need to do this? Also isn't it equivalent to fake_mode.allow_non_fake_inputs?
There was a problem hiding this comment.
I think normally aoti assumes this flag has been set from upper layer of the call stack, but in our case we are a new toplevel function so we need to set this properly.
| reset_cudagraph_trees() | ||
|
|
||
|
|
||
| class _TorchCompileAOTInductorWrapper(_TorchCompileInductorWrapper): |
There was a problem hiding this comment.
Let's say I'm a random user with torch.compile, how do I use this? Do I have to pass in this specific backend to fullgraph_compile? Should we make it lower friction somehow?
I.e. backend="inductor" + some config = AOTI
backend="inductor" without config = Python Wrapper
13b69e4 to
cc82255
Compare
|
Updated with:
|
|
@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 |
Merge failedReason: 1 jobs have failed, first few of them are: trunk / macos-py3-arm64 / test (mps, 1, 1, macos-m2-15) Details for Dev Infra teamRaised by workflow job |
torch/_inductor/output_code.py
Outdated
| if self.device_type.startswith("cuda"): | ||
| current_callable = ( | ||
| torch._C._aoti.AOTIModelContainerRunnerCuda( # type: ignore[call-arg] | ||
| current_callable, 1, self.device_type |
There was a problem hiding this comment.
One perf trick here is to set run_single_threaded to True, otherwise it won't compose with cudagraphs, see #148601 for more backgrounds.
35425f9 to
ffff272
Compare
|
@pytorchbot merge |
|
@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 |
|
@pytorchbot revert -m "seems to be breaking internal tests and builds, see D86919103" -c ghfirst |
|
@pytorchbot successfully started a revert job. Check the current status here. |
This reverts commit 273babe. Reverted #167338 on behalf of https://github.com/jeanschmidt due to seems to be breaking internal tests and builds, see D86919103 ([comment](#167338 (comment)))
|
@zhxchen17 your PR has been successfully reverted. |
44985e1 to
e9cc1a8
Compare
e9cc1a8 to
62bea25
Compare
|
@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 |
Fixes #ISSUE_NUMBER Pull Request resolved: pytorch#167338 Approved by: https://github.com/jamesjwu
This reverts commit 273babe. Reverted pytorch#167338 on behalf of https://github.com/jeanschmidt due to seems to be breaking internal tests and builds, see D86919103 ([comment](pytorch#167338 (comment)))
Fixes #ISSUE_NUMBER Pull Request resolved: pytorch#167338 Approved by: https://github.com/jamesjwu
Fixes #ISSUE_NUMBER
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @Lucaskabela