[论文解读] Massively Parallel Tensor Network State Algorithms on Hybrid CPU-GPU Based Architectures

Andor Menczer, Örs Legeza|arXiv (Cornell University)|May 9, 2023

Distributed and Parallel Computing Systems被引用 6

一句话总结

本文提出了一种大规模并行、混合CPU-GPU张量网络态算法，专为高性能计算优化，引入了Maze-Runner任务调度、TTCache内存管理以及SBMM4S批量矩阵乘法等新技术。该算法在8张A100 GPU上实现了接近艾浮级（exaflop）的性能，相较于纯CPU执行最高可提速20倍，能够高效模拟希尔伯特空间维度高达2.88×10^36的强关联体系。

ABSTRACT

The interplay of quantum and classical simulation and the delicate divide between them is in the focus of massively parallelized tensor network state (TNS) algorithms designed for high performance computing (HPC). In this contribution, we present novel algorithmic solutions together with implementation details to extend current limits of TNS algorithms on HPC infrastructure building on state-of-the-art hardware and software technologies. Benchmark results obtained via large-scale density matrix renormalization group (DMRG) simulations are presented for selected strongly correlated molecular systems addressing problems on Hilbert space dimensions up to $2.88 imes10^{36}$.

研究动机与目标

• 通过利用张量网络态（TNS）方法，克服经典计算下量子多体模拟的指数级扩展问题。
• 通过大规模并行化，在高性能计算（HPC）架构上扩展密度矩阵密度矩阵重正化群（DMRG）模拟的极限。
• 为混合CPU-GPU系统开发高效、可扩展且可移植的算法解决方案，以处理量子化学与凝聚态物理中的大规模张量运算。
• 通过优化硬件利用与低开销任务调度，实现对具有长程相互作用的强关联体系的千万亿次（petascale）模拟。

提出的方法

• 提出Maze-Runner，一种轻量级、无角色的任务并行模型，通过任务分组实现同质化线程使用，简化调度机制，降低在波动工作负载下的开销。
• 提出TTCache，一种分层虚拟内存寻址方案，将数据分解为属性并映射至执行块，最小化冗余I/O、内存碎片化与分配成本。
• 开发SBMM4S，一种零成本累加归约的批量矩阵乘法内核，通过分步内存访问与偏移批量处理，实现多个矩阵-向量运算的高效交错计算。
• 采用多级并行化策略，结合CPU与GPU的低级SIMD执行与高级HPC作业调度，实现跨多个硬件与软件抽象层的独立执行。
• 通过将大型张量运算替换为可扩展的独立向量与矩阵运算，对一般模型哈密顿量（包括长程相互作用）的DMRG算法进行优化。
• 将MPI协议迁移至混合CPU-多GPU框架，支持跨多个节点的分布式计算，实现千万亿次规模模拟。

实验结果

研究问题

• RQ1 混合CPU-GPU架构能否在强关联量子体系的张量网络态模拟中实现接近艾浮级的性能？
• RQ2 在典型量子多体模拟中常见的不规则、波动性工作负载下，如何最小化任务调度开销？
• RQ3 在大规模张量网络计算中，内存访问模式在多大程度上可被优化，以减少I/O、碎片化与分配成本？
• RQ4 具有零成本累加归约的批量矩阵乘法能否显著加速DMRG模拟中的对角化步骤？
• RQ5 随着GPU加速器数量的增加，所提算法的扩展行为如何？是否可达到现代硬件的理论性能极限？

主要发现

• 即使仅使用单张GPU，混合CPU-多GPU实现相较纯CPU执行仍实现2–4倍提速，且在多个加速器上性能呈线性扩展。
• 在8张NVIDIA A100-PCIE-40GB GPU下，系统达到77.6 TFLOPS的性能，接近该硬件配置的理论峰值性能77.6 TFLOPS。
• 对于较大的键维数D，总模拟时间在双重对数尺度下呈线性下降，相较于纯CPU执行，使用8张GPU时提速超过20倍。
• 该算法成功对F2分子在CAS(18,18)空间中的有效哈密顿量完成对角化，处理了高达2.88×10^36的希尔伯特空间维度。
• 作为量子化学中的基准体系，FeMoco簇被该优化的GPU加速DMRG方法高效模拟，证明了其在真实强关联体系中的适用性。
• 所提出的Maze-Runner、TTCache与SBMM4S技术实现了高效、低开销且高度可扩展的执行，为具有长程相互作用的量子多体体系的千万亿次模拟铺平了道路。