Hello everyone. I have an issue with memory not releasing when using Pennylane both for CPUs and GPUs. When I try running the following code, the memory usage doubles instead of the first dev allocation releasing the memory since the variable is being overwritten. Can you please help me in proper usage of the library?
import pennylane as qml
dev = qml.device("lightning.qubit", wires=28)
#small circuit
@qml.qnode(dev)
def circuit(theta):
qml.Hadamard(0)
qml.RZ(theta, wires=0)
qml.adjoint(qml.RX(theta, wires=0))
return qml.expval(qml.PauliZ(0))
circuit(0.5)
del circuit
dev = qml.device("lightning.qubit", wires=28)
Garbage collection in Python is usually automated, and so it handles memory release when it thinks it can be safely handled. Lightning devices are deallocated in the C++ destructors of their respective statevector classes, which are bound to Python. To remove all Python references to these objects calling del helps, but is often still no guarantee that the garbage collector will be called immediately. Note also that the dev handle is where the statevector memory is held.
You can try the following, and let us know if it works for your case:
import gc
import pennylane as qml
dev = qml.device("lightning.qubit", wires=28)
#small circuit
@qml.qnode(dev)
def circuit(theta):
qml.Hadamard(0)
qml.RZ(theta, wires=0)
qml.adjoint(qml.RX(theta, wires=0))
return qml.expval(qml.PauliZ(0))
circuit(0.5)
del circuit
dev = qml.device("lightning.qubit", wires=28)
gc.collect()
Memory management is generally a product of your operating system and how Python decides to handle things, so even though you deleted the QNode and replaced the device handle, the old buffer many hang around until the garbage collector runs. On my system, the allocate/deallocate calls happen immediately.
No, unfortunately that does not help. I had tried it already. Even if I do del dev and gc.collect(), the memory doesn’t get released. It works only if I don’t run the circuit (circuit(0.5)), but otherwise I don’t know how I can make it free the memory. It seems that after running a circuit dev gets additional references.
Thanks again for the feedback. It seems that for your workload our older device API is caching the statevector in an internal LRU cache entity, which is preventing it from being released. We have redesigned the device API in PennyLane for which LightningQubit has been ported to for the upcoming 0.36 release. This issue should no longer be present in LQ once this goes live.
However, for 0.35 and devices like LightningGPU, we can force a reset of this cache by calling dev.map_wires.cache_clear(), which forces the process-global cache to remove all data from that given method.
As an example, you can see the effect of this with the following minimum example:
import pennylane as qml
import gc
def fun(idx: int):
dev = qml.device("lightning.qubit", wires=28)
dev.map_wires.cache_clear() # Comment this out, and the issue returns
tape = qml.tape.QuantumScript([qml.Hadamard(0), qml.RX(0.5, wires=0)], [qml.expval(qml.Z(0))])
out = dev.execute(tape)
return out
if __name__ == "__main__":
for i in range(10):
res = fun(i)
While we complete the update of our devices to the new API, the above should suffice to allow you to redefine the device/qnode as needed. For most workloads, we usually avoid redefining the device or QNode, which is why this hasn’t appeared to us in the past, instead favoring defining the device once, and reusing it as needed. Let us know if this helps.
Yes, it helped a lot. Thank you for providing this solution, it solved my issue. I was getting this image of references to dev before using dev.map_wires.cache_clear().
Hi @Tigran_Mamikonyan
Glad to hear it. We are also discussing an update to PennyLane that will allow this issue to be resolved across the ecosystem. We are hoping to have this ready sometime in the coming days, after which the suggested cache clear method should no longer be required. We will keep you posted when this goes live.
