Deadlock with lightning.gpu backend on 14-qubit circuit (1000 observables, example case)

PennyLane Help
1

Hi! I’m experiencing a deadlock when using the lightning.gpu backend in a JAX training setup.
The program hangs indefinitely (GPU at 100% utilization, no output) — but runs fine with lightning.qubit.

This minimal script defines a 14-qubit circuit with 1000 observables — the large observable count is only for demonstration, not real training.
Still, this setup consistently causes the GPU backend to freeze.

Can anyone help me? Thanks

import pennylane as qml
import jax, jax.numpy as jnp, numpy as np, os

# --- GPU env setup ---
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

# --- circuit config ---
N_QUBITS = 14
N_OBS = 1000

# --- device ---
dev = qml.device("lightning.gpu", wires=N_QUBITS)

# --- dummy observable set (for test only) ---
ALL_OBSERVABLES = [qml.PauliZ(i % N_QUBITS) for i in range(N_OBS)]

@qml.qnode(dev, interface="jax")
def circuit(params):
    for i in range(N_QUBITS):
        qml.RX(params[i], wires=i)
        qml.RY(params[i]*0.5, wires=i)
    for i in range(N_QUBITS-1):
        qml.CNOT(wires=[i, i+1])
    return [qml.expval(obs) for obs in ALL_OBSERVABLES]

params = jax.random.uniform(jax.random.PRNGKey(0), (N_QUBITS,), 0, 3.14)
print("Running 14-qubit circuit with 1000 observables (demo only)...")
out = circuit(params)
print("Output shape:", jnp.shape(out))
2

Note: The architecture consists of a quantum circuit that outputs 1000 expectation values (observables). These features are then passed through two dense layers trained with Adam optimizer and cross-entropy loss.

3

Hii @Nhan_Hoang_Quang , welcome to the Forum!

I just ran your code on Google Colab and I got an error because jax.random.uniform has some unexpected arguments 0, 3.14. As soon as I removed this, your code ran very quickly with lightning.gpu and everything else the same. I only changed the params to params = jax.random.uniform(jax.random.PRNGKey(0), (N_QUBITS,)).

If you want to try this on Google Colab you can change the runtime type to GPU and pip install the following in the first line:

!pip install custatevec_cu12
!pip install pennylane-lightning-gpu

If this works for you on Google Colab but not on your own machine then there may be a compatibility or installation issue.

I hope this helps!

4

Thank you very much for your help❤️ I will check again my enviroment and codes.

1 Like