How can I natively parallelize my batch calculations and leverage the GPU with an accelerated simulator backend

PennyLane Help
1

Performance Bottleneck: When running this loop on a CUDA-enabled session, active GPU utilization scales down completely to 0.00% while a single thread on the CPU spikes to 100%. If I remove the explicit `.cpu()` cast from the internal elements to force direct GPU execution, the framework runtime halts immediately with the following error layout:

import torch
import torch.nn as nn
import torch.optim as optim
import pennylane as qml

# 1. Global Simulation Environment (8 Wires)
n_qubits = 8
dev = qml.device("default.qubit", wires=n_qubits)

@qml.qnode(dev, interface="torch")
def batch_processing_circuit(inputs, weights):
    # Standard Angle Embedding configuration
    qml.AngleEmbedding(inputs, wires=range(n_qubits), rotation='X')
    
    # Parameterized single-qubit rotations
    for i in range(n_qubits):
        qml.RX(weights[i], wires=i)
        
    # Fixed entanglement chain
    for i in range(0, n_qubits, 2):
        qml.CNOT(wires=[i, (i + 1) % n_qubits])
        
    return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]

# 2. Hybrid Network Wrapper Architecture
class HybridModel(nn.Module):
    def __init__(self):
        super().__init__()
        # Classical feature extractor
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(12 * 8 * 8, 64),
            nn.ReLU(),
            nn.Linear(64, n_qubits)
        )
        # Variational circuit weights
        self.q_weights = nn.Parameter(torch.randn(12))
        self.output_head = nn.Linear(n_qubits, 1000)
        
    def forward(self, x):
        # x arrives initialized on the GPU target
        features = torch.tanh(self.encoder(x)) * 3.14
        
        # Core Bottleneck Loop: Processing elements line-by-line via CPU fallback
        quantum_outputs = []
        for sample in features:
            # Forcing data to .cpu() prevents cross-device tensor crashes, 
            # but completely stalls GPU utilization down to 0%
            q_out = torch.stack(batch_processing_circuit(sample.cpu(), self.q_weights.cpu()))
            quantum_outputs.append(q_out)
            
        final_features = torch.stack(quantum_outputs).to(x.device).float()
        return self.output_head(final_features)

# 3. Execution Harness (Simulates a single batch forward evaluation pass)
if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Active Device Target: {device}")
    
    model = HybridModel().to(device)
    
    # Standard random noise mock tensor representing a 32-sample batch
    mock_batch = torch.randn(32, 12, 8, 8).to(device)
    
    output = model(mock_batch)
    print("Forward pass finished.")

If you want help with diagnosing an error, please put the full error message below:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

**
import** pennylane as qml

qml.about()

Name: pennylane
Version: 0.45.1
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: 
Author: 
License: 
Location: /usr/local/lib/python3.12/dist-packages
Platform info:           Linux-6.12.90+-x86_64-with-glibc2.35
Python version:          3.12.13
Numpy version:           2.0.2
Scipy version:           1.16.3
JAX version:             0.7.2
Catalyst version:        None
Installed devices:
- default.clifford (pennylane-0.45.1)
- default.gaussian (pennylane-0.45.1)
- default.mixed (pennylane-0.45.1)
- default.qubit (pennylane-0.45.1)
- default.qutrit (pennylane-0.45.1)
- default.qutrit.mixed (pennylane-0.45.1)
- default.tensor (pennylane-0.45.1)
- null.qubit (pennylane-0.45.1)
- reference.qubit (pennylane-0.45.1)
- lightning.qubit (pennylane_lightning-0.45.0)
2

Hi @Jay_Prakash , welcome to the Forum!

Are you planning on running this on only 8 qubits, or potentially more in the future? I’m asking because if you only need 8 qubits using a GPU will likely make things slower, not faster.

If this is just an example and you actually want to run on more qubits please let me know how many (more or less) so that I can suggest the best alternative.

3

Thank you so much for the quick feedback, I understand about the memory overhead for a singular 8-qubit system, but in my use case ,I am training a hybrid deep learning model (a Quantum CNN) on a massive dataset of 14 million chess states using PyTorch.

Because of the size of the dataset, I am trying to run a large batch size (512 samples per step(increased it over each step to reduce the time taken to train(since im running on colab time constraint). I am currently compressing the classical features down to 8 inputs, embedding them via AngleEmbedding, and evaluating expectation values across the 8 qubits.

Since I am simulating 512 independent states simultaneously at every single training step, I am trying to figure out the best way to leverage hardware parallelization.

I originally picked 8 qubits because:

  1. It naturally maps to the 8x8 spatial grid topology of a chess board.

  2. I assumed keeping the qubit count small would keep the simulator lightning-fast, letting us rely on a fast trainable classical encoder to compress the board features before embedding.

  3. currently running a large batch size (512 samples per step), meaning we are simulating 512 independent 8-qubit states simultaneously to keep our PyTorch training stable.

My Question: Given that maximizing throughput per second for these 512-batch steps is my absolute primary bottleneck, is keeping the circuit at 8 qubits and using a classical pre-encoder the most efficient way to design this?
Alternatively, if widening the circuit to more qubits (e.g., 12 to 16 wires) would unlock better native GPU parallelization or allow it to embed the chess features more natively without the classical bottleneck, we are completely open to scaling it up!!!

Dataset I’m using: https://www.kaggle.com/datasets/zdfowler/encoded-chess-games

4

Although i don’t have the exact code when I posted this query but this what I’m working with now:(So you could get an idea on what I’m working with)
Pseudo+Code:

import os
import time
import torch
import torch.nn as tn
import torch.optim as optim
import pennylane as qml
import numpy as np

==========================================

1. HARDWARE & QUANTUM BACKEND PIPELINE

==========================================

n_qubits = 8
dev = qml.device(“lightning.gpu”, wires=n_qubits)

def Q_convolution(weights):

[PROPRIETARY ARCHITECTURE]: Custom variational entangling

and rotation gates mapping across the 8-qubit space.

pass

def Q_pooling(weights):

[PROPRIETARY ARCHITECTURE]: Custom controlled-unitary operations

performing quantum dimensionality reduction.

pass

@qml.qnode(dev, interface=“torch”)
def pure_qcnn_circuit(inputs, weights):
qml.AngleEmbedding(inputs, wires=range(n_qubits), rotation=‘X’)
Q_convolution(weights[0:8])
Q_pooling(weights[8:12])
return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]

==========================================

2. HYBRID QUANTUM DEEP LEARNING MODEL

==========================================

class Standalone_QCNN(tn.Module):
def init(self):
super(Standalone_QCNN, self).init()

12 Quantum weights (8 for conv, 4 for pooling)

self.q_weights = tn.Parameter(torch.randn(12))

Massive classical classification projection head (Maps to full action space)

self.policy_head = tn.Linear(n_qubits, 65536)
self.value_head = tn.Linear(n_qubits, 1)

def forward(self, precomputed_inputs):

CORE BOTTLENECK QUESTION:

precomputed_inputs shape is [BATCH_SIZE=512, n_qubits=8]

quantum_outputs =

Current Implementation: Iterating row-by-row through the batch

for sample in precomputed_inputs:
q_out = torch.stack(pure_qcnn_circuit(sample, self.q_weights))
quantum_outputs.append(q_out)

quantum_features = torch.stack(quantum_outputs).to(precomputed_inputs.device)
quantum_features = quantum_features.float()

policy_logits = self.policy_head(quantum_features)
value_prediction = self.value_head(quantum_features)

return policy_logits, value_prediction

==========================================

3. PRE-COMPUTED DATA EXTRACTION (CHUNKS)

==========================================

Note: The raw dataset contains ~14.1 Million sequence trajectory games.

We multi-threadedly read the files, apply a fixed classical mapping projection,

and save out precomputed .npy feature arrays to disk to optimize CPU load.

5

Thanks for sharing this additional information @Jay_Prakash !

I don’t have an answer for all of your questions but what I do know is that for 8 qubits you’re better off using “default.qubit” or maybe “lightning.qubit”.

I would recommend that you start with “default.qubit”.

I encourage you to test it out and work on optimizing your code to achieve better performance. Note however that quantum computers and simulators are generally bad at working with large amounts of data, and they’re bad at certain tasks where classical computers excel. So if you feel like the quantum circuit is just too slow, it may be just the nature of things, it doesn’t necessarily mean that you’re doing something wrong.

You can see an example in this demo. Feel free to explore the PennyLane demos library for more examples that can help you for inspiration or for improving your approach.

I hope this helps!