Add a how-to for catalyst-compiling "Symmetry-invariant quantum machine learning force fields" #1222
Conversation
…ne learning force fields" The how-to contains the full code listing, and some primitive tutorial words.
|
👋 Hey, looks like you've updated some demos! 🐘 Don't forget to update the Please hide this comment once the field(s) are updated. Thanks! |
|
We should reach out to the original author before deciding whether this should be an additional section to the original demo or a separate new demo. Until then I won't add too many documentation/metadata.json boilerplate, just because it might not be necessary. @josh146 |
There was a problem hiding this comment.
Thanks @paul0403! It would be worth running black on the file to fix up some formatting, but otherwise looks good to me! 🎉
|
@paul0403 It looks like a meta data file is needed so we can look at the built version :) |
paul0403
force-pushed
the
eqnn_forcefield_qjit
branch
from
8d5569a to
c72a374
Compare
There was a problem hiding this comment.
Thanks @paul0403! This looks fine to me 😄. I made a few suggestions.
|
@isaacdevlugt given that this change is quite small, should we make any updates to the metadata? |
|
@paul0403 oh yes this is a good point! You'll want to change the "Updated on:" field in the metadata. @alvaro-at-xanadu do you have any other recommendations for metadata changes, or is the updated date sufficient? Cliff notes for this PR: we replaced |
Done |
@isaacdevlugt This should be all that's needed! |
|
@isaacdevlugt I don't think there's any leftover work here, so I think we should deploy it and mark it as done. Not sure about CI though. |
| loss, grads = jax.value_and_grad(cost, argnums=0)(net_params, loss_data) | ||
|
|
||
| loss = cost(net_params, loss_data) | ||
| grads = catalyst.grad(cost, method="fd", h=1e-13, argnums=0)(net_params, loss_data) |
There was a problem hiding this comment.
When changing to catalyst, it was discovered that QubitUnitary cannot be differentiated:
catalyst.utils.exceptions.DifferentiableCompileError: QubitUnitary is non-differentiable on 'lightning.qubit' device
To make the demo work, I had to manually change gradient method to finite difference. This causes significant performance degradation.
Possible paths forward:
- Find another demo to convert
- Still convert this demo, but either (a) make lightning work with differentiating through qubit unitary, or (b) decompose qubit unitary into "standard" gateset, although unsure whether this decomposition itself can be jitted or not, or (c) decrease the batch size and the number of traning steps
|
@paul0403 @isaacdevlugt I recommend not merging this demo if it is using finite-difference or a reduction in the timesteps that deviates from the current demo. @paul0403 you should try option (b): QubitUnitary.compute_decomposition(U, wires=...)assuming the unitary matrix is small enough to support decomposition. |
The problem is that decomposition essentially makes the circuit (hence the IR) different for different unitaries. I strongly suspect this decomposition itself is not jittable, as catalyst itself handles decomposition at the frontend, so essentially the decomposition would make qjit recompile every single run. Edit: actually @mudit2812 and I found out that this method will actually fail with infinite recursion and/or a custom |
|
Thank you for opening this pull request. You can find the built site at this link. Deployment Info:
Note: It may take several minutes for updates to this pull request to be reflected on the deployed site. |
So I tried this: #qml.QubitUnitary(U, wires=wires, id="U")
decomp = qml.QubitUnitary.compute_decomposition(U, wires=wires)#, id="U")
for op in decomp:
qml.apply(decomp)This raises the following warning and error: I am unsure what this error means, as it seems like the queue contains operator only. |
The differentiability warning is certainly a concern, since that is a requirement for this demo. The actual error is probably because of this typo though: |
OH! Good catch 💦 However after fixing this we get an even more mysterious jax error now: Also, this approach likely just does not work at all, as seen by this factored out example: import pennylane as qml
import catalyst
from catalyst import qjit
import jax
from jax import numpy as jnp
dev = qml.device("lightning.qubit", wires=1)
@qjit(keep_intermediate=False)
@qml.qnode(dev)
def circuit(x):
U = jnp.array([[1,0],[0,x]])
decomp = qml.QubitUnitary.compute_decomposition(U, wires=0)
for op in decomp:
qml.apply(op)
return qml.probs()
print(circuit(1.0)) # fine, [1. 0.]
def f(x):
probs = circuit(x)
return probs[0] + probs[1]
print(catalyst.grad(f, argnums=0)(1.0)) |
Weird, it's just a data layout issue though which can probably fixed relatively easily.
This one seems worse, must be a gap in our stablehlo support 🤔 |
|
I opened an issue so we can formally tackle the underlying problem in the future: |
The how-to contains the full code listing, and some primitive tutorial words.
Before submitting
Please complete the following checklist when submitting a PR:
Ensure that your tutorial executes correctly, and conforms to the
guidelines specified in the README.
Remember to do a grammar check of the content you include.
All tutorials conform to
PEP8 standards.
To auto format files, simply
pip install black, and thenrun
black -l 100 path/to/file.py.When all the above are checked, delete everything above the dashed
line and fill in the pull request template.
Title: How to Quantum Just-In-Time Compile "Symmetry-invariant quantum machine learning force fields" with Catalyst
Summary:
As part of Catalyst's work on identifying 10 demos to compile with catalyst, we convert "Symmetry-invariant quantum machine learning force fields", a machine learning workflow that calls a training step function repeatedly and thus can have significant performance boosts with catalyst compilation.
Note: the original demo is from an external user, so we should NOT merge this until we have reached out to them. Update: we have obtained their permission.
Relevant references:
Possible Drawbacks:
Because lightning does not support differentiation through
QubitUnitary, we have to resort to finite difference method for gradients, which brings a big performance degradation.Related GitHub Issues:
[sc-72938]
If you are writing a demonstration, please answer these questions to facilitate the marketing process.
GOALS — Why are we working on this now?
Promote Catalyst by demonstrating the performance improvements it offers by QJIT compiling a relatively complex end-to-end quantum workflow.
AUDIENCE — Who is this for?
Chemistry and quantum machine learning researchers.
KEYWORDS — What words should be included in the marketing post?
Which of the following types of documentation is most similar to your file?
(more details here)