cirq + custatevec on multiple GPUs

[hthayko]

Hi,

I am using cuQuantum Appliance 22.07-Cirq to experiment with cirq + custatevec simulator.
I am able to run up to 32 qubit simulations with single NVidia A100 40GB gpu as expected.

However I am having trouble getting it to run on multiple GPUs. I am using QSimOptions.gpu_mode = 2 to achieve this as explained in cuquantum docs but I only see one gpu being used thourgh nvidia-smi command, and I run out of memory for 33 qubits.

here is a minimal reproducer:

import cirq
import qsimcirq

def load_test(num_gpus = 1, depth = 4, num_qubits = 30):
	circuit = cirq.testing.random_circuit(
		qubits = num_qubits,
		n_moments = depth,
		op_density = 1.0,
		random_state = 1)
	num_gates = len(list(circuit.all_operations()))
	options = {"gpu_mode": num_gpus, "n_subsvs": num_gpus}
	qsim_simulator = qsimcirq.QSimSimulator(options)
	result = qsim_simulator.simulate(circuit)
	print (f"DONE with qubits: {num_qubits} \t gates: {num_gates} \t depth: {depth} \t")

load_test(num_gpus = 2, depth = 2, num_qubits = 30) # uses  ~8GB memory on 1 gpu
load_test(num_gpus = 2, depth = 2, num_qubits = 32) # uses ~32GB memory on 1 gpu
load_test(num_gpus = 2, depth = 2, num_qubits = 33) # CUDA error: out of memory vector_mgpu.h 116

I have 2 NVidia A100 gpu's on my machine and here is the output of nvidia-smi for 32 qubit case:

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 510.47.03    Driver Version: 510.47.03    CUDA Version: 11.6     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA A100-SXM...  Off  | 00000000:00:04.0 Off |                    0 |
| N/A   35C    P0    82W / 400W |  33561MiB / 40960MiB |     23%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA A100-SXM...  Off  | 00000000:00:05.0 Off |                    0 |
| N/A   43C    P0    75W / 400W |      2MiB / 40960MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     26614      C   python                          33559MiB |
+-----------------------------------------------------------------------------+

Am I missing something?

[leofang]

Hi @hthayko thank you for using the cuQuantum container! This problem is due to the fact that in Qsim the public option names do not match the keys used by QsimSimulator internally, see https://github.com/quantumlib/qsim/blob/11e8afadc7f03c1691c0e506b9620311b26be2d6/qsimcirq/qsim_simulator.py#L76-L101. We had to follow the Google team's convention. Therefore, if instead of passing a dict you create/pass a QSimOptions object

options = qsimcirq.QSimOptions(gpu_mode=num_gpus)

it'd be correct. Extra tips:

• you don't have to set n_subsvs explicitly
• a new 22.11 container is out, which includes some perf optimizations

Hope this helps 🙂

[hthayko]

Worked like a charm. Thanks 🙏