feat(torch.compile): partial PET application#1029
Draft
HaoZeke wants to merge 1 commit intometatensor:mainfrom
Draft
feat(torch.compile): partial PET application#1029HaoZeke wants to merge 1 commit intometatensor:mainfrom
HaoZeke wants to merge 1 commit intometatensor:mainfrom
Conversation
Co-authored-by: sirmarcel <sirmarcel@users.noreply.github.com>
Luthaf
reviewed
Feb 2, 2026
| self._compile_enabled = self.hypers.get("compile", False) | ||
| self._compiled = False | ||
|
|
||
| def _maybe_compile(self) -> None: |
Member
There was a problem hiding this comment.
So this does not do a full torch compile for the full model, only in-place replacing of some modules, correct?
Do you know if full model compilation can be done?
Member
Author
There was a problem hiding this comment.
So this does not do a full torch compile for the full model, only in-place replacing of some modules, correct?
Yup.
Do you know if full model compilation can be done?
Not without a rewrite (@sirmarcel has one)
Luthaf
reviewed
Feb 2, 2026
| if self._compile_enabled and not self._compiled: | ||
| # Compile the GNN layers - this is the computational bottleneck | ||
| for i, layer in enumerate(self.gnn_layers): | ||
| self.gnn_layers[i] = torch.compile(layer, mode="reduce-overhead") |
Member
There was a problem hiding this comment.
Does this work with fullgraph=True? This is what we would need for torch export down the line
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Bare minimal changes to "use"
torch.compile. @sirmarcel has a much nicer UPET port which can be both compiled and exported.Of course the main blocker is that
metatomicneeds TorchScript at the moment.Not something to be merged, just discussed.
Contributor (creator of pull-request) checklist
Maintainer/Reviewer checklist
📚 Documentation preview 📚: https://metatrain--1029.org.readthedocs.build/en/1029/