[论文解读] Hybrid quantum programming with PennyLane Lightning on HPC platforms

Ali Asadi, Amintor Dusko+6 more|arXiv (Cornell University)|Mar 4, 2024

Quantum Computing Algorithms and Architecture被引用 4

一句话总结

该论文介绍了PennyLane Lightning，一个基于C++的高性能状态矢量模拟器套件，专为CPU、GPU和高性能计算（HPC）环境优化。它通过SIMD加速、多线程处理、批量执行以及跨多个节点的分布式MPI模拟，实现了对混合量子-经典工作负载（如VQE、QAOA和基于梯度的优化）的高效支持，可在GPU和多节点HPC平台中实现显著的性能提升，最高支持41量子比特电路的模拟。

ABSTRACT

We introduce PennyLane's Lightning suite, a collection of high-performance state-vector simulators targeting CPU, GPU, and HPC-native architectures and workloads. Quantum applications such as QAOA, VQE, and synthetic workloads are implemented to demonstrate the supported classical computing architectures and showcase the scale of problems that can be simulated using our tooling. We benchmark the performance of Lightning with backends supporting CPUs, as well as NVidia and AMD GPUs, and compare the results to other commonly used high-performance simulator packages, demonstrating where Lightning's implementations give performance leads. We show improved CPU performance by employing explicit SIMD intrinsics and multi-threading, batched task-based execution across multiple GPUs, and distributed forward and gradient-based quantum circuit executions across multiple nodes. Our data shows we can comfortably simulate a variety of circuits, giving examples with up to 30 qubits on a single device or node, and up to 41 qubits using multiple nodes.

研究动机与目标

• 开发一个高性能、可移植的量子模拟器套件，支持包括CPU、NVIDIA和AMD GPU以及HPC集群在内的多种经典硬件后端。
• 通过自动微分支持，实现对变分量子算法（如VQE、QAOA）和基于梯度的优化等混合量子-经典工作负载的高效执行。
• 通过MPI实现跨多个节点的可扩展、分布式量子电路模拟，突破单台设备的限制，支持大规模电路模拟。
• 在所有后端（包括分布式GPU环境）中提供高性能、原生集成的量子梯度计算（采用伴随法）。
• 为研究人员提供准确的运行时估算指南，用于HPC环境中的资源分配，从而支持大规模量子模拟的工作负载高效规划。

提出的方法

• 基于C++20设计并实现模块化模拟器堆栈，利用SIMD内嵌函数（AVX-512、AVX2）优化门核函数，提升CPU上的单线程和多线程性能。
• 为lightning.qubit（CPU）、lightning.gpu（NVIDIA CUDA）和lightning.kokkos（AMD ROCm）实现设备后端，并通过PennyLane的设备无关框架提供统一接口。
• 利用CUDA流和多GPU任务调度，实现跨多个GPU的独立电路批量、基于任务的执行。
• 通过MPI实现分布式状态矢量模拟，支持跨集群的多节点执行，涵盖正向和基于梯度的电路评估。
• 在所有后端中原生集成伴随法Jacobian计算，支持在JAX、PyTorch和TensorFlow工作流中高效实现反向模式自动微分。
• 利用现代C++20特性与抽象层，确保在Perlmutter、LUMI和AWS EC2等多样化HPC平台上的可移植性与高性能。

实验结果

研究问题

• RQ1 在单节点、多GPU和分布式多节点HPC环境中，PennyLane Lightning的性能与其它高性能模拟器相比如何？
• RQ2 SIMD内嵌函数和多线程处理在多大程度上能提升变分算法中基于CPU的量子电路模拟性能？
• RQ3 使用MPI在多个节点和GPU上进行分布式量子电路模拟的可扩展性和效率如何？
• RQ4 在统一框架中，能否高效且原生地支持CPU、GPU和分布式HPC后端的伴随法梯度计算？
• RQ5 研究人员如何利用该工具集在HPC环境中估算大规模量子模拟的资源需求和运行时间？

主要发现

• 在512块GPU上，Lightning在基于梯度的工作负载（如电路梯度）中表现出比其他模拟器更强的弱缩放效率，分别达到8%和30%的效率，提升幅度达30%。
• 在30量子比特电路中，使用AVX-512的Lightning CPU后端在门应用时间方面相比基线实现有显著性能提升，尤其在特定门类型上表现更优。
• GPU后端（lightning.gpu）在电路超过20量子比特后相比CPU展现出性能优势，且随着电路深度和量子比特数的增加，优势愈发明显。
• lightning.gpu是唯一原生支持多节点环境下伴随法梯度计算的分布式状态矢量模拟器，可在大规模模拟中实现高效的反向模式微分。
• 该框架通过多节点分布式执行支持高达41量子比特的电路模拟，证明了其在大规模量子算法原型设计中的可行性。
• 提供了适用于各类工作负载的运行时估算指南，帮助研究人员在Perlmutter和LUMI等HPC系统上规划资源分配，相关基准测试已在真实HPC硬件上验证。