HPC talk

State-vector simulation on HPC platforms

Hybrid quantum programs, run as memory-bound kernels on CPUs, GPUs, and MPI jobs.

Based on Hybrid quantum programming with PennyLane Lightning on HPC platforms (arXiv:2403.02512).

The working set

The state vector is the working set

  • An n-qubit circuit is a vector of 2n complex amplitudes. One more qubit doubles it.
  • At 24 qubits, with thousands of parameters, the paper says the memory is already hundreds of gigabytes.
  • Differentiating every gate on top of that vector is what makes a hybrid loop expensive. The simulators keep the device a black box and hand the framework a custom gradient.
n n+1 n+2 n+3 ×2 each qubit shape only

Stated size, not bar height: 24 qubits and thousands of parameters can take hundreds of gigabytes.

Workload shapes

Three ways the work shows up

01 · Task parallel

Many small circuits

Independent fragments, each happy with a slice of a GPU. Schedule them. Do not invent a bigger vector than the fragment needs.

02 · One device

Fill a socket or a GPU

One circuit owns the CPU socket or the local GPU. Kernel choice and overheads decide the runtime.

03 · Distributed

The vector does not fit

Multi-GPU or multi-node. The state is partitioned, and gates that touch remote amplitudes pay for the network.

The paper sets out to cover these three. A fourth, mixing them in one job, is left for later.

Kernel

A gate is a streaming update

  • A one-qubit gate touches disjoint pairs of amplitudes. Indices come from bit masks, so the 2n×2n matrix is never built.
  • Disjoint pairs are a parallel loop. OpenMP on the gate is a compile-time option.
  • An n-controlled gate skips every amplitude whose controls are off, so the paper expects about 2−n the time of the uncontrolled kernel.
  • Hand-written kernels cover 1 to 4 qubits. Past that, a general kernel takes over.
state vector thread A thread B pairs do not overlap

Backends

One interface, three backends

lightning.qubit

C++17/20 on the CPU. OpenMP over observables, and optionally over the gate.

At runtime the x86_64 build picks a portable kernel, AVX2, or AVX-512. No recompile.

lightning.gpu

NVIDIA cuQuantum on CUDA compute capability 7.0 and newer.

The same API runs distributed, through MPI and cuStateVec, including adjoint gradients.

lightning.kokkos

The same kernels in Kokkos: OpenMP, CUDA from Pascal through Ampere, and ROCm on Instinct MI100 and MI200.

Useful when the CPU has many cores and few observables, or the GPU is not NVIDIA.

Every backend is templated on 32-bit or 64-bit floats, from Python down to the kernel. The benchmarks below are double precision.

CPU microbenchmark

SIMD helps until the cache says no

  • 30-qubit RX gate, C++ binary, double precision, Sapphire Rapids on an EC2 R7i.16xlarge (32 physical cores).
  • AVX-512 beats AVX2 and the portable kernel at 1, 4, and 16 threads.
  • At higher thread counts the portable kernel wins. perf c2c pointed at cache eviction from extra loads and stores.
  • Streaming stores (_mm512_stream_pd in place of _mm512_store_pd) give AVX-512 the lead back when the thread count is high.
1, 4, 16 threads AVX-512 ahead of AVX2 and the portable kernel.
Higher thread counts Portable kernel leads. Intrinsics saturate.
Streaming stores AVX-512 recovers at high thread counts.

Same machine, Hadamard, RX, and CNOT against other simulators: best single-thread time and the lowest variance. At the top thread count the spread is on the order of 10−3 seconds.

Why the 1-thread number matters

Gradients are the real workload

Adjoint

Linear in the gates

Adjoint differentiation applies each parametric gate a constant number of times. Parameter-shift and finite differences are quadratic in the parameter count.

OpenMP

One thread per observable

The gradient farms observables out as OpenMP threads. Inside each thread the kernel is often single-threaded, so the 1-thread SIMD result is the one the pipeline multiplies.

Do not compare naively

Different complexity

A simulator without this gradient path is solving a harder loop. The paper treats that as an unfair comparison, not a missing speedup.

Single node · VQE

Where the GPU overtakes the CPU

  • Energy plus gradient, 10 times, release v0.34.0. CPU side: 32-core AMD EPYC Milan, up to 32 threads, AVX2 dispatch.
  • GPU side: a 40 GB NVIDIA A100, and Kokkos built for CUDA.
  • Under 20 qubits, overheads win and the CPU backend is the fastest.
  • H2, CC-PVDZ: 20 qubits, 2,951 Hamiltonian terms, GPUs about 2×.
  • C2H2: 24 qubits, 6,401 terms. The two GPU backends match, and both are around 7× Lightning-Qubit on Sapphire Rapids.
2× ~7× 20 qubits 2,951 terms 24 qubits 6,401 terms 1× CPU

Task farm · Perlmutter

Forty-five thousand independent circuits

79 qubits, before the cut QAOA MaxCut, one layer. Three clusters of 25 nodes, two edges between clusters.
45,711 fragments, 27–29 qubits Circuit cutting, not counting the offset identity. Each fragment is its own GPU task.
  • lightning.gpu and Ray on Perlmutter, from 1 to 64 nodes, which is 4 to 256 GPUs.
  • Packing a batch of fragments onto each GPU beat allocating memory on demand for every circuit.
  • The batched schedule reaches about 80 circuits per second at 256 GPUs.

Distributed VQE · LUMI

One molecule, many GPUs

  • C2H4 with no term grouping: a 28-qubit ansatz, over 3,000 gates, and close to 9,000 Pauli words.
  • Each word is an independent expectation. lightning.kokkos on HIP, with an mpi4py process pool, farms the observables across LUMI.
  • The runs take at most 10 optimizer steps. Energies matched on every distribution that finished.
  • One node, 8 GPUs, hit the 48-hour wall clock. That point is not a full measurement: it is estimated from the mean and standard deviation of 8 completed steps. The larger node counts use the measured times.

This is task parallelism over the Hamiltonian, not one state vector sliced across ranks. The next slide is the other kind of MPI.

MPI state vector · Perlmutter

When one vector spans the cluster

8% weak-scaling efficiency, sampling 104 shots of a strongly entangling layered circuit. Largest run: 256 GPUs.
30% weak-scaling efficiency, gradients Adjoint Jacobian, one Pauli-Z per qubit. Largest run: 512 GPUs.

80 GB A100s, Cray MPICH, CUDA 11.7, December 2023. Elsewhere the paper’s examples reach 30 qubits on one device, a 36-qubit MPI illustration, and 41 qubits across nodes.

The distributed GPU backend returns that adjoint Jacobian. The paper presents it as the distributed state-vector implementation that supports this gradient.

Scheduling

Place the job by its memory

What you have What the paper’s numbers say
Under 20 qubits CPU. Overheads make the GPU slower. Thread the observables, and let the runtime pick AVX2 or AVX-512.
20–24 qubits, thousands of terms One GPU. About 2× at 20 qubits and 2,951 terms, around 7× at 24 qubits and 6,401 terms.
Tens of thousands of fragments Batch them. 45,711 fragments reached about 80 circuits per second on 256 GPUs. On-demand allocation did worse.
One vector past a device MPI. Budget for weak-scaling efficiency of 8% (sampling, 256 GPUs) or 30% (gradients, 512 GPUs), out to 41 qubits in these examples.

Takeaways

What to take back to the machine

  • The hot loop is a streaming update of disjoint amplitude pairs. SIMD is the first lever. Streaming stores are the second, once cache traffic wins.
  • Hybrid programs spend their time in gradients. Adjoint mode keeps that linear, and the useful parallelism is one observable per thread.
  • Past one GPU, a task farm and a distributed state vector are different machines. The farm scaled to about 80 circuits per second. The distributed vector held 8% and 30% weak-scaling efficiency.

Based on Hybrid quantum programming with PennyLane Lightning on HPC platforms (arXiv:2403.02512).