Combining torch.compile with cueq #1175
Merged
Conversation
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
Collaborator
Author
|
@ilyes319 , still more work to do for compilation but I wanted to check that the choices for the methods here look ok to you. |
Contributor
|
mmm good question @mariogeiger, is this the intended settings? |
mariogeiger
reviewed
Sep 12, 2025
Contributor
mariogeiger
left a comment
mariogeiger
left a comment
There was a problem hiding this comment.
Did I answer the questions @ilyes319 ?
| shared_weights=shared_weights, | ||
| internal_weights=internal_weights, | ||
| use_fallback=True, | ||
| method="naive", |
Contributor
There was a problem hiding this comment.
Good.
Just a question, is this for the skipTP? is one of the input always a one-hot vector? If so, why not indexing the weights instead of contract with a one-hot?
Contributor
There was a problem hiding this comment.
yeah indeed, does cueq provide the option to pass the weight in the forward?
There was a problem hiding this comment.
Hi @ilyes319, since 0.6.0 we support indexing in the linear: you just need to provide the number of weight_classes at init and pass the weight_indices in the fwd.
If your indices are sorted you can also use the experimental "indexed_linear" method, but even in the "naive" case this would first index into the weights, then compute a linear for each element, which should still be faster.
Contributor
|
@mariogeiger Yes thank you!!! |
Collaborator
Author
|
The latest commit on this PR works in my limited testing but really just changes the scope of the workaround of using symbolic tracing. @mariogeiger would be good to have your take on this - the cueq torch ops perform a number of runtime checks on input tensors. This is obviously a great thing for developer UX but breaks the symbolic tracing that is currently used in MACE. This is a broader design choice since having runtime errors that introduce side-effects that are incompatible with the static control flow requirement for symbolic tracing. A couple of choices are possible here:
Keen to hear thought or suggestions! |
Contributor
|
Since you are using the argument "method=..." you should maybe update the version in your setup.cfg -> |
|
Hi, we were discussing about this: in the past we had some of these checks with a wrapper like Would something like this, with the appropriate mode, work to disable the checks for compiled models? Also, can you tell me exactly what kind of compilation command you're trying to enable here? Because we're testing some like |
Collaborator
Author
At the moment mace is first passed through |
Contributor
|
Hey @hatemhelal, curious what the status of the PR? should it be merged, is there anything else to work on? I can think about the following point to use that most effectively in MD:
|
Collaborator
Author
|
@ilyes319 think this should be ok to merge, the follow on discussion is a tangential direction for making this more robust. Also recall that @ThomasWarford was looking to add benchmarks for this in #1184. Agree with all the follow on points you suggest, just not sure when I might get to do them so can't commit at the moment! |
Contributor
|
ok I merged it for now, @ThomasWarford tell me if you are up to try implementing these. |
Contributor
|
I tried implementing it today and I get errors running The first was fixed by Here's the first other error: I believe #1184 should work if this is resolved, but I do still get another error with keys (which I posted there), rather than this error. This is python version 3.12, CUDA 12.6, PyTorch 2.9.0+cu126 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 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.
This PR provides a simple fix to enable the use of
torch.compilein conjunction withcueq.The main change is to provide a tighter scope for the symbolic tracing that was previously applied to the
NonLinearReadoutBlockand instead just simplify the activation constructor.Also updated tests to parameterise over enabling cueq