Make CUDAQ objects pickleable
Opened this issue · 1 comments
zohimchandani commented
Required prerequisites
- Search the issue tracker to check if your feature has already been mentioned or rejected in other issues.
Describe the feature
I am trying to manually hand-code the hamiltonian batching workflow where I distribute via MPI:
import cudaq
from mpi4py import MPI
import cupy as cp
from cudaq import spin
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
num_ranks = comm.Get_size()
cp.cuda.runtime.setDevice(rank)
print(f'rank {rank} running with GPU {cp.cuda.Device().pci_bus_id}')
cudaq.set_target("nvidia")
qubit_count = 4
@cudaq.kernel
def kernel(qubit_count: int):
qvector = cudaq.qvector(qubit_count)
h(qvector[0])
for i in range(1, qubit_count):
cx(qvector[0], qvector[i])
# we have a hamiltonian
h = spin.x(0) + spin.y(1) + spin.z(2) + spin.x(3)
# we split it into 4 batches for distribution across 4 gpus via mpi
batched_h = [[spin.x(0)], [spin.y(1)], [spin.z(2)], [spin.x(3)]]
# we distribute each term across the gpus available
comm.scatter(batched_h, root = 0)
root@gorby:/home/cudaq/cudaq_work# mpirun -np 4 --allow-run-as-root python3 mpi.py
rank 2 running with GPU 0000:81:00.0
rank 3 running with GPU 0000:C2:00.0
rank 1 running with GPU 0000:47:00.0
rank 0 running with GPU 0000:01:00.0
TypeError: cannot pickle 'cudaq.mlir._mlir_libs._quakeDialects.cudaq_runtime.SpinOperator' object
Can we make CUDAQ types pickleable?
Thanks team.
amccaskey commented
You could probably get away with serializing / deserializing the spin operator
serialized = op.serialize() # list of floats
# distribute / scatter the data
deserialized = cudaq.SpinOperator(serialized, numQubits)